db_cmd.go

package peerdb

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"

	internalSite "gitlab.com/peerdb/peerdb/internal/site"

	"github.com/hashicorp/go-cleanhttp"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/rs/zerolog"
	"gitlab.com/tozd/go/errors"
	"gitlab.com/tozd/go/x"
	"gitlab.com/tozd/identifier"

	"gitlab.com/peerdb/peerdb/document"
	"gitlab.com/peerdb/peerdb/indexer"
	"gitlab.com/peerdb/peerdb/internal/export"
	internalSearch "gitlab.com/peerdb/peerdb/internal/search"
	internalStore "gitlab.com/peerdb/peerdb/internal/store"
)

// startAndWaitSite starts the base for a site, runs optional beforeStart before it does, then waits for
// indexing to catch up, and refreshes the ElasticSearch indices.
//
// beforeStart runs before the base is started, which is where whatever prepares the store for the indexing
// belongs: once the base is started, its bridge is catching up on the commit log and its river client is
// running reindex jobs. It shares the counters progress is reported through with the indexing which follows,
// so that the site is reported on as one run.
func startAndWaitSite(
	ctx context.Context, logger zerolog.Logger, site internalSite.Site,
	beforeStart func(ctx context.Context, logger zerolog.Logger, site internalSite.Site, count, size *x.Counter) errors.E,
) (func(), errors.E) {
	// We set fallback context values which are used to set application name on PostgreSQL connections.
	ctx = internalStore.WithFallbackDBContext(ctx, site.Schema, "db")

	count := x.NewCounter(0)
	size := x.NewCounter(0)
	progress := indexer.Progress(logger, "indexing", nil)
	ticker := x.NewTicker(ctx, count, size, indexer.ProgressPrintRate)
	defer ticker.Stop()
	go func() {
		for p := range ticker.C {
			progress(ctx, p)
		}
	}()

	if beforeStart != nil {
		errE := beforeStart(ctx, logger, site, count, size)
		if errE != nil {
			return nil, errE
		}
	}

	documents, errE := converterDocuments(ctx, site.Base)
	if errE != nil {
		return nil, errE
	}

	onShutdown, errE := site.Start(ctx, documents)
	if errE != nil {
		return onShutdown, errE
	}

	errE = site.Base.WaitUntilCaughtUp(ctx, count, size)
	if errE != nil {
		return onShutdown, errE
	}

	for _, index := range site.LevelIndexes() {
		_, err := site.ESClient.Indices.Refresh().Index(index).Do(ctx)
		if err != nil {
			errE := internalSearch.WithESError(ctx, err)
			errors.Details(errE)["index"] = index
			return onShutdown, errE
		}
	}

	logger.Info().
		Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).
		Int64("count", count.Count()).
		Int64("total", size.Count()).
		Msg("indexing done")

	return onShutdown, nil
}

// enqueueForReindex enqueues every document of the site into the reindex queue. The queue is drained by the
// job the base submits when it is started, so this runs before that.
func enqueueForReindex(ctx context.Context, logger zerolog.Logger, site internalSite.Site, count, size *x.Counter) errors.E {
	enqueued, errE := site.Base.EnqueueAllForReindex(ctx, count, size)
	if errE != nil {
		return errE
	}

	logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Int("enqueued", enqueued).Msg("enqueued documents for reindex")

	return nil
}

// Run executes the db wait command which initializes the base,
// waits until all pending indexing is complete, and then exits.
func (c *DBWaitCommand) Run(globals *Globals) errors.E {
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	ctx, cancel := context.WithCancel(ctx)

	onShutdownInit, errE := Init(ctx, globals)
	if onShutdownInit != nil {
		defer onShutdownInit()
	}
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()
	if errE != nil {
		return errE
	}

	onShutdown := []func(){}
	onShutdownF := func() {
		for _, f := range onShutdown {
			if f == nil {
				continue
			}
			f()
		}
	}
	defer onShutdownF()
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("waiting for indexing")

		onS, errE := startAndWaitSite(ctx, globals.Logger, site, nil)
		onShutdown = append(onShutdown, onS)
		if errE != nil {
			return errE
		}
	}

	globals.Logger.Info().Msg("db wait done")

	return nil
}

// Run executes the db reindex command which re-indexes every document and then exits. It re-renders each
// document's current state through the reindex queue, from the store and the derived state (the document
// metadata and the materialized state), which are trusted: when they are not, run db repair first. With
// --recreate-index the ElasticSearch indices are deleted first and recreated from the current mapping
// during initialization, so a mapping change is applied without losing source data.
func (c *DBReindexCommand) Run(globals *Globals) errors.E {
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	// When recreating the index, delete it before Init so that the base's EnsureIndex (run during
	// startup) recreates it from the current mapping. The documents are then re-rendered into the fresh
	// index below. Deletion resolves the level name through its alias to the concrete index, if it is an
	// alias (EnsureIndex creates such alias layout).
	if c.RecreateIndex {
		esClient, errE := internalSearch.GetClient(cleanhttp.DefaultPooledClient(), globals.Logger, globals.Elastic.URL)
		if errE != nil {
			return errE
		}
		for _, site := range globals.Sites {
			for _, index := range site.LevelIndexes() {
				errE = internalSearch.DeleteIndex(ctx, esClient, index)
				if errE != nil {
					return errE
				}
				globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("index", index).Msg("index deleted for recreation")
			}
		}
	}

	ctx, cancel := context.WithCancel(ctx)

	onShutdownInit, errE := Init(ctx, globals)
	if onShutdownInit != nil {
		defer onShutdownInit()
	}
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()
	if errE != nil {
		return errE
	}

	onShutdown := []func(){}
	onShutdownF := func() {
		for _, f := range onShutdown {
			if f == nil {
				continue
			}
			f()
		}
	}
	defer onShutdownF()
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("reindexing")

		// The reindex re-renders every document's current state through the reindex queue: it enqueues every
		// document and the drain reads each at its latest version (so deleted documents are skipped, never
		// transiently re-created), rendering it from the store and the materialized state without writing
		// either.
		onS, errE := startAndWaitSite(ctx, globals.Logger, site, enqueueForReindex)
		onShutdown = append(onShutdown, onS)
		if errE != nil {
			return errE
		}

		// Re-indexing rewrites each document once, leaving its superseded (deleted) Lucene version behind.
		// Those deletes linger per shard until merged and bloat the index (they also skew per-shard term
		// statistics). Now that the site is caught up, expunge them. We use only_expunge_deletes rather than
		// max_num_segments because the index keeps receiving live writes after a reindex, and full-merging an
		// index that is still written to is discouraged.
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Msg("expunging deletes")

		for _, index := range site.LevelIndexes() {
			_, err := site.ESClient.Indices.Forcemerge().Index(index).OnlyExpungeDeletes(true).Do(ctx)
			if err != nil {
				errE := internalSearch.WithESError(ctx, err)
				errors.Details(errE)["index"] = index
				return errE
			}
			globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("index", index).Msg("deletes expunged")
		}
	}

	globals.Logger.Info().Msg("db reindex done")

	return nil
}

// repairSite rederives the derived state of one site from its stored documents: first the extracted
// content in the document metadata, repaired in place at every stored revision, and then the
// materialized state (the references, inverse-relations, and embedding tables), rebuilt from the latest
// versions of the documents. Both passes report progress through one set of counters, so the site is
// reported on as one run.
func repairSite(ctx context.Context, logger zerolog.Logger, site internalSite.Site, dryRun bool) errors.E {
	// We set fallback context values which are used to set application name on PostgreSQL connections.
	ctx = internalStore.WithFallbackDBContext(ctx, site.Schema, "db")

	count := x.NewCounter(0)
	size := x.NewCounter(0)
	progress := indexer.Progress(logger, "repairing", nil)
	ticker := x.NewTicker(ctx, count, size, indexer.ProgressPrintRate)
	defer ticker.Stop()
	go func() {
		for p := range ticker.C {
			progress(ctx, p)
		}
	}()

	examined, repaired, errE := site.Base.RepairDocumentsMetadata(ctx, dryRun, count, size)
	if errE != nil {
		return errE
	}
	event := logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Int64("examined", examined).Int64("repaired", repaired)
	if dryRun {
		event.Msg("document metadata which would be repaired")
	} else {
		event.Msg("document metadata repaired")
	}

	// The schema documents are listed through the just-repaired metadata (see
	// store.DocumentMetadata.InstanceOf), which is why the metadata pass runs first. With dryRun the
	// metadata stays as it is, so the drift reported below is computed under the current (possibly
	// unrepaired) metadata and can differ from what a repair would then compute. We say so when the
	// metadata does need repair, because the converters are then built from an incomplete schema and
	// the drift numbers cannot be taken at face value.
	if dryRun && repaired > 0 {
		logger.Warn().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Int64("repaired", repaired).
			Msg("document metadata needs repair, so the materialized state drift is computed from an incomplete schema and can be wrong")
	}

	documents, errE := converterDocuments(ctx, site.Base)
	if errE != nil {
		return errE
	}

	drift, errE := site.Base.RepairMaterializedState(ctx, documents, dryRun, count, size)
	if errE != nil {
		return errE
	}
	event = logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).
		Int64("referencesExtra", drift.ReferencesExtra).Int64("referencesMissing", drift.ReferencesMissing).
		Int64("inverseRelationsExtra", drift.InverseRelationsExtra).Int64("inverseRelationsMissing", drift.InverseRelationsMissing).
		Int64("embeddingExtra", drift.EmbeddingExtra).Int64("embeddingMissing", drift.EmbeddingMissing)
	switch {
	case dryRun:
		event.Msg("materialized state drift")
	case drift.Any():
		event.Msg("materialized state rebuilt")
	default:
		event.Msg("materialized state unchanged")
	}

	return nil
}

// Run executes the db repair command which repairs the database of every site, taking its stored
// documents as the single source of truth, and then exits. With --dry-run it only reports what would
// be repaired. The repaired state is not re-rendered into ElasticSearch: run db reindex for that. It
// runs against a stopped deployment: the base is never started here, and nothing else may be writing
// the store or rendering from the materialized state while it runs.
func (c *DBRepairCommand) Run(globals *Globals) errors.E {
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	onShutdownInit, errE := Init(ctx, globals)
	if onShutdownInit != nil {
		defer onShutdownInit()
	}
	if errE != nil {
		return errE
	}

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("repairing")

		errE := repairSite(ctx, globals.Logger, site, c.DryRun)
		if errE != nil {
			return errE
		}
	}

	globals.Logger.Info().Msg("db repair done")

	return nil
}

// vacuumSchema runs VACUUM on every table in the given PostgreSQL schema to reclaim space held by
// dead tuples and to refresh planner statistics. VACUUM cannot run inside a transaction, so we
// acquire a single connection and run each statement directly in autocommit mode. We enumerate the
// schema's tables instead of running a bare VACUUM because the latter vacuums every schema in the
// database, which would redundantly revisit other sites sharing the same database.
func vacuumSchema(ctx context.Context, dbpool *pgxpool.Pool, schema string) errors.E {
	conn, err := dbpool.Acquire(ctx)
	if err != nil {
		return internalStore.WithPgxError(err)
	}
	defer conn.Release()

	// A VACUUM of a large table takes longer than the statement timeout the pool sets, and the timeout
	// applies to it like to any other statement. It is set on the session and not with SET LOCAL, which
	// has no effect outside a transaction block, so it is reset before the connection goes back to the
	// pool: releasing a connection resets only the application name and the search path, so the lifted
	// timeout would otherwise be inherited by whoever acquires it next. The reset runs on an uncancelled
	// context so that it still happens when ctx is already done.
	_, err = conn.Exec(ctx, `SET statement_timeout = 0`)
	if err != nil {
		return internalStore.WithPgxError(err)
	}
	defer func() {
		_, err := conn.Exec(context.WithoutCancel(ctx), `RESET statement_timeout`)
		if err != nil {
			zerolog.Ctx(ctx).Error().Err(internalStore.WithPgxError(err)).Msg(`unable to reset "statement_timeout" after vacuuming`)
		}
	}()

	rows, err := conn.Query(ctx, `SELECT tablename FROM pg_tables WHERE schemaname = $1 ORDER BY tablename`, schema)
	if err != nil {
		return internalStore.WithPgxError(err)
	}
	tables, err := pgx.CollectRows(rows, pgx.RowTo[string])
	if err != nil {
		return internalStore.WithPgxError(err)
	}

	for _, table := range tables {
		_, err := conn.Exec(ctx, fmt.Sprintf(`VACUUM (ANALYZE) "%s"."%s"`, schema, table))
		if err != nil {
			return internalStore.WithPgxError(err)
		}
	}

	return nil
}

// Run executes the db vacuum command which reclaims dead tuples in PostgreSQL
// and expunges deleted documents from ElasticSearch for all configured sites.
func (c *DBVacuumCommand) Run(globals *Globals) errors.E {
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	// We use context.WithoutCancel here because we want to cancel the pool ourselves and not when context
	// is cancelled (so that cleanup code which needs PostgreSQL access can continue to use connections).
	dbpool, dbpoolCleanup, errE := internalStore.InitPostgres(
		context.WithoutCancel(ctx),
		string(globals.Postgres.URL),
		globals.Logger,
		getRequestWithFallback(),
	)
	if errE != nil {
		return errE
	}
	defer dbpoolCleanup()

	esClient, errE := internalSearch.GetClient(cleanhttp.DefaultPooledClient(), globals.Logger, globals.Elastic.URL)
	if errE != nil {
		return errE
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("vacuuming")

		// We set fallback context values which are used to set application name on PostgreSQL connections.
		siteCtx := internalStore.WithFallbackDBContext(ctx, site.Schema, "vacuum")

		errE = vacuumSchema(siteCtx, dbpool, site.Schema)
		if errE != nil {
			return errE
		}

		globals.Logger.Info().Str("schema", site.Schema).Msg("schema vacuumed")

		// We expunge deleted Lucene documents per shard so the superseded versions that versioned writes
		// and reindexing accumulate do not bloat the index or skew per-shard term statistics. We use
		// only_expunge_deletes rather than max_num_segments because the index keeps receiving live writes,
		// and full-merging an index that is still written to is discouraged.
		for _, index := range site.LevelIndexes() {
			_, err := esClient.Indices.Forcemerge().Index(index).OnlyExpungeDeletes(true).Do(siteCtx)
			if err != nil {
				errE := internalSearch.WithESError(ctx, err)
				errors.Details(errE)["index"] = index
				return errE
			}
			globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("index", index).Msg("deletes expunged")
		}
	}

	globals.Logger.Info().Msg("db vacuum done")

	return nil
}

// Run executes the db export command which exports documents to CSV or JSON.
func (c *DBExportCommand) Run(globals *Globals) (returnErr errors.E) { //nolint:nonamedreturns
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	ctx, cancel := context.WithCancel(ctx)

	onShutdownInit, errE := Init(ctx, globals)
	if onShutdownInit != nil {
		defer onShutdownInit()
	}
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()
	if errE != nil {
		return errE
	}

	onShutdown := []func(){}
	onShutdownF := func() {
		for _, f := range onShutdown {
			if f == nil {
				continue
			}
			f()
		}
	}
	defer onShutdownF()
	// It is safe to call cancel multiple times. We want it to be
	// called before any onShutdown waits.
	defer cancel()

	// Determine output writer.
	var w io.Writer
	if c.Output == "-" {
		w = os.Stdout
	} else {
		f, err := os.Create(c.Output)
		if err != nil {
			return errors.WithStack(err)
		}
		defer func() {
			errE := f.Close()
			if errE != nil && returnErr == nil {
				returnErr = errors.WithStack(errE)
			}
		}()
		w = f
	}

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("exporting")

		// We set fallback context values which are used to set application name on PostgreSQL connections.
		siteCtx := internalStore.WithFallbackDBContext(ctx, site.Schema, "export")

		onS, errE := startAndWaitSite(siteCtx, globals.Logger, site, nil)
		onShutdown = append(onShutdown, onS)
		if errE != nil {
			return errE
		}

		getDoc := func(ctx context.Context, id identifier.Identifier) (*document.D, errors.E) {
			doc, _, _, _, errE := site.Base.GetDocumentLatestDoc(ctx, id)
			return doc, errE
		}

		errE = export.Export(siteCtx, w, site.ESClient, site.TopIndex(), getDoc, export.Config{
			Format:     c.Format,
			InstanceOf: c.InstanceOf,
			Properties: c.Property,
		})
		if errE != nil {
			return errE
		}
	}

	globals.Logger.Info().Msg("db export done")

	return nil
}

// Run executes the db diagram command which writes a Mermaid ER diagram describing classes and fields.
func (c *DBDiagramCommand) Run(globals *Globals) (returnErr errors.E) { //nolint:nonamedreturns
	var w io.Writer
	if c.Output == "-" {
		w = os.Stdout
	} else {
		f, err := os.Create(c.Output)
		if err != nil {
			return errors.WithStack(err)
		}
		defer func() {
			errE := f.Close()
			if errE != nil && returnErr == nil {
				returnErr = errors.WithStack(errE)
			}
		}()
		w = f
	}

	return export.Diagram(globals.Logger, w, c.SkipCore)
}

// Run executes the db wipe command which drops PostgreSQL schemas
// and deletes ElasticSearch indices for all configured sites.
func (c *DBWipeCommand) Run(globals *Globals) errors.E {
	// We stop gracefully on ctrl-c and TERM signal.
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	ctx = globals.Logger.WithContext(ctx)

	errE := InitSites(globals)
	if errE != nil {
		return errE
	}

	// We use context.WithoutCancel here because we want to cancel the pool ourselves and not when context
	// is cancelled (so that cleanup code which needs PostgreSQL access can continue to use connections).
	dbpool, dbpoolCleanup, errE := internalStore.InitPostgres(
		context.WithoutCancel(ctx),
		string(globals.Postgres.URL),
		globals.Logger,
		getRequestWithFallback(),
	)
	if errE != nil {
		return errE
	}
	defer dbpoolCleanup()

	esClient, errE := internalSearch.GetClient(cleanhttp.DefaultPooledClient(), globals.Logger, globals.Elastic.URL)
	if errE != nil {
		return errE
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	for _, site := range globals.Sites {
		globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("schema", site.Schema).Msg("wiping")

		// We set fallback context values which are used to set application name on PostgreSQL connections.
		siteCtx := internalStore.WithFallbackDBContext(ctx, site.Schema, "wipe")

		errE = internalStore.RetryTransaction(siteCtx, dbpool, pgx.ReadWrite, func(ctx context.Context, tx pgx.Tx) errors.E {
			_, err := tx.Exec(ctx, fmt.Sprintf(`DROP SCHEMA IF EXISTS "%s" CASCADE`, site.Schema))
			return internalStore.WithPgxError(err)
		})
		if errE != nil {
			return errE
		}

		globals.Logger.Info().Str("schema", site.Schema).Msg("schema dropped")

		for _, index := range site.LevelIndexes() {
			errE = internalSearch.DeleteIndex(siteCtx, esClient, index)
			if errE != nil {
				return errE
			}
			globals.Logger.Info().Str("indexPrefix", site.IndexPrefix).Str("index", index).Msg("index deleted")
		}
	}

	// File contents live on disk in the storage directory, content-addressed and shared across all
	// sites, so we clear it once as part of the wipe.
	errE = clearDirContents(globals.Storage.Dir)
	if errE != nil {
		return errE
	}
	globals.Logger.Info().Str("dir", globals.Storage.Dir).Msg("storage directory cleared")

	globals.Logger.Info().Msg("db wipe done")

	return nil
}

// clearDirContents removes everything inside dir, leaving dir itself in place. A missing directory is
// treated as already empty.
func clearDirContents(dir string) errors.E {
	entries, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return nil
	}
	if err != nil {
		errE := errors.WithStack(err)
		errors.Details(errE)["path"] = dir
		return errE
	}
	for _, entry := range entries {
		path := filepath.Join(dir, entry.Name())
		err := os.RemoveAll(path)
		if err != nil {
			errE := errors.WithStack(err)
			errors.Details(errE)["path"] = path
			return errE
		}
	}
	return nil
}