otel-collector/cmd/o11yotelcollector/migrate/sync_check.go

package migrate

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"time"

	"github.com/cenkalti/backoff/v4"
	"github.com/hanzo-ds/go"
	"github.com/hanzoai/otel-collector/cmd/o11yotelcollector/config"
	schemamigrator "github.com/hanzoai/otel-collector/cmd/o11yschemamigrator/schema_migrator"
	"github.com/hanzoai/otel-collector/constants"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

type syncCheck struct {
	conn             datastore.Conn
	timeout          time.Duration
	migrationManager *schemamigrator.MigrationManager
	logger           *zap.Logger
}

func registerSyncCheck(parentCmd *cobra.Command, logger *zap.Logger) {
	syncCheckCommand := &cobra.Command{
		Use:          "check",
		Short:        "Checks the status of migrations for the store by checking the status of migrations in the migration table.",
		SilenceUsage: true,
		RunE: func(cmd *cobra.Command, args []string) error {
			check, err := newSyncCheck(config.Datastore.DSN, config.Datastore.Cluster, config.Datastore.Replication, config.MigrateSyncCheck.Timeout, logger)
			if err != nil {
				return err
			}

			err = check.Run(cmd.Context())
			if err != nil {
				return err
			}

			return nil
		},
	}

	config.MigrateSyncCheck.RegisterFlags(syncCheckCommand)

	parentCmd.AddCommand(syncCheckCommand)
}

func newSyncCheck(dsn string, cluster string, replication bool, timeout time.Duration, logger *zap.Logger) (*syncCheck, error) {
	opts, err := datastore.ParseDSN(dsn)
	if err != nil {
		return nil, err
	}

	conn, err := datastore.Open(opts)
	if err != nil {
		return nil, err
	}

	migrationManager, err := schemamigrator.NewMigrationManager(
		schemamigrator.WithClusterName(cluster),
		schemamigrator.WithReplicationEnabled(replication),
		schemamigrator.WithConn(conn),
		schemamigrator.WithConnOptions(*opts),
		schemamigrator.WithLogger(logger),
	)
	if err != nil {
		return nil, err
	}

	return &syncCheck{
		conn:             conn,
		timeout:          timeout,
		migrationManager: migrationManager,
		logger:           logger,
	}, nil
}

func (cmd *syncCheck) Run(ctx context.Context) error {
	backoff := backoff.NewExponentialBackOff()
	backoff.MaxElapsedTime = cmd.timeout

	for {
		err := cmd.Check(ctx)
		if err == nil {
			break
		}

		cmd.logger.Info("Error occurred while checking for sync migrations to complete, retrying", zap.Error(err))
		nextBackOff := backoff.NextBackOff()
		if nextBackOff == backoff.Stop {
			return errors.New("timed out waiting for sync migrations to complete within the configured timeout")
		}
		time.Sleep(nextBackOff)
	}

	return nil
}

func (cmd *syncCheck) Check(ctx context.Context) error {
	tracesLastSyncMigrationID, err := cmd.getLastSyncMigration(schemamigrator.TracesMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yTracesDB, tracesLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", tracesLastSyncMigrationID, schemamigrator.O11yTracesDB)
		}
	}

	logsMigrations := schemamigrator.LogsMigrations
	if constants.EnableLogsMigrationsV2 {
		logsMigrations = schemamigrator.LogsMigrationsV2
	}

	logsLastSyncMigrationID, err := cmd.getLastSyncMigration(logsMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yLogsDB, logsLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", logsLastSyncMigrationID, schemamigrator.O11yLogsDB)
		}
	}

	metricsLastSyncMigrationID, err := cmd.getLastSyncMigration(schemamigrator.MetricsMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yMetricsDB, metricsLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", metricsLastSyncMigrationID, schemamigrator.O11yMetricsDB)
		}
	}

	metadataLastSyncMigrationID, err := cmd.getLastSyncMigration(schemamigrator.MetadataMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yMetadataDB, metadataLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", metadataLastSyncMigrationID, schemamigrator.O11yMetadataDB)
		}
	}

	analyticsLastSyncMigrationID, err := cmd.getLastSyncMigration(schemamigrator.AnalyticsMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yAnalyticsDB, analyticsLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", analyticsLastSyncMigrationID, schemamigrator.O11yAnalyticsDB)
		}
	}

	meterLastSyncMigrationID, err := cmd.getLastSyncMigration(schemamigrator.MeterMigrations)
	if err == nil {
		ok, err := cmd.migrationManager.CheckMigrationStatus(ctx, schemamigrator.O11yMeterDB, meterLastSyncMigrationID, schemamigrator.FinishedStatus)
		if err != nil {
			return err
		}

		if !ok {
			return fmt.Errorf("migration with ID %d for database '%s' has not been completed", meterLastSyncMigrationID, schemamigrator.O11yMeterDB)
		}
	}

	return nil
}

func (cmd *syncCheck) getLastSyncMigration(migrations []schemamigrator.SchemaMigrationRecord) (uint64, error) {
	for _, migration := range slices.Backward(migrations) {
		if cmd.migrationManager.IsSync(migration) {
			return migration.MigrationID, nil
		}
	}

	return 0, fmt.Errorf("no sync migration found")
}