feat(cmd): add migrate command

cmd/enterprise/main.go

@@ -1,10 +1,18 @@
 package main
 
 import (
+	"context"
 	"log/slog"
 
 	"github.com/SigNoz/signoz/cmd"
+	"github.com/SigNoz/signoz/ee/sqlschema/postgressqlschema"
+	"github.com/SigNoz/signoz/ee/sqlstore/postgressqlstore"
+	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/instrumentation"
+	"github.com/SigNoz/signoz/pkg/signoz"
+	"github.com/SigNoz/signoz/pkg/sqlschema"
+	"github.com/SigNoz/signoz/pkg/sqlstore"
+	"github.com/SigNoz/signoz/pkg/sqlstore/sqlstorehook"
 )
 
 func main() {
@@ -14,5 +22,21 @@ func main() {
 	// register a list of commands to the root command
 	registerServer(cmd.RootCmd, logger)
 
+	// TODO(grandwizard28): DRY this code
+	sqlstoreFactories := signoz.NewSQLStoreProviderFactories()
+	if err := sqlstoreFactories.Add(postgressqlstore.NewFactory(sqlstorehook.NewLoggingFactory())); err != nil {
+		logger.ErrorContext(context.TODO(), "failed to add postgressqlstore factory", "error", err)
+		panic(err)
+	}
+
+	cmd.RegisterSQL(cmd.RootCmd, logger, func(sqlstore sqlstore.SQLStore) factory.NamedMap[factory.ProviderFactory[sqlschema.SQLSchema, sqlschema.Config]] {
+		existingFactories := signoz.NewSQLSchemaProviderFactories(sqlstore)
+		if err := existingFactories.Add(postgressqlschema.NewFactory(sqlstore)); err != nil {
+			panic(err)
+		}
+
+		return existingFactories
+	}, sqlstoreFactories)
+
 	cmd.Execute(logger)
 }

cmd/sql.go

@@ -0,0 +1,155 @@
+package cmd
+
+import (
+	"log/slog"
+
+	"github.com/SigNoz/signoz/pkg/factory"
+	"github.com/SigNoz/signoz/pkg/instrumentation"
+	"github.com/SigNoz/signoz/pkg/signoz"
+	"github.com/SigNoz/signoz/pkg/sqlmigration"
+	"github.com/SigNoz/signoz/pkg/sqlmigrator"
+	"github.com/SigNoz/signoz/pkg/sqlschema"
+	"github.com/SigNoz/signoz/pkg/sqlstore"
+	"github.com/SigNoz/signoz/pkg/version"
+	"github.com/spf13/cobra"
+)
+
+// TODO(grandwizard28): DRY this code
+func RegisterSQL(parentCmd *cobra.Command, logger *slog.Logger, sqlSchemaProviderFactories func(sqlstore.SQLStore) factory.NamedMap[factory.ProviderFactory[sqlschema.SQLSchema, sqlschema.Config]], sqlstoreProviderFactories factory.NamedMap[factory.ProviderFactory[sqlstore.SQLStore, sqlstore.Config]]) {
+	sqlCmd := &cobra.Command{
+		Use:   "sql",
+		Short: "Run commands to interact with the SQL",
+	}
+
+	migrateCmd := &cobra.Command{
+		Use:   "migrate",
+		Short: "Run migrations for the SQL database",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			config, err := NewSigNozConfig(ctx, signoz.DeprecatedFlags{})
+			if err != nil {
+				return err
+			}
+
+			// Initialize instrumentation
+			instrumentation, err := instrumentation.New(ctx, config.Instrumentation, version.Info, "signoz")
+			if err != nil {
+				return err
+			}
+
+			// Get the provider settings from instrumentation
+			providerSettings := instrumentation.ToProviderSettings()
+
+			// Initialize sqlstore from the available sqlstore provider factories
+			sqlstore, err := factory.NewProviderFromNamedMap(
+				ctx,
+				providerSettings,
+				config.SQLStore,
+				sqlstoreProviderFactories,
+				config.SQLStore.Provider,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Initialize sqlschema from the available sqlschema provider factories
+			sqlschema, err := factory.NewProviderFromNamedMap(
+				ctx,
+				providerSettings,
+				config.SQLSchema,
+				sqlSchemaProviderFactories(sqlstore),
+				config.SQLStore.Provider,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Run migrations on the sqlstore
+			sqlmigrations, err := sqlmigration.New(
+				ctx,
+				providerSettings,
+				config.SQLMigration,
+				signoz.NewSQLMigrationProviderFactories(sqlstore, sqlschema),
+			)
+			if err != nil {
+				return err
+			}
+
+			err = sqlmigrator.New(ctx, providerSettings, sqlstore, sqlmigrations, config.SQLMigrator).Migrate(ctx)
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+	}
+
+	rollbackCmd := &cobra.Command{
+		Use:   "rollback",
+		Short: "Rollback the last migration",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			config, err := NewSigNozConfig(ctx, signoz.DeprecatedFlags{})
+			if err != nil {
+				return err
+			}
+
+			// Initialize instrumentation
+			instrumentation, err := instrumentation.New(ctx, config.Instrumentation, version.Info, "signoz")
+			if err != nil {
+				return err
+			}
+
+			// Get the provider settings from instrumentation
+			providerSettings := instrumentation.ToProviderSettings()
+
+			// Initialize sqlstore from the available sqlstore provider factories
+			sqlstore, err := factory.NewProviderFromNamedMap(
+				ctx,
+				providerSettings,
+				config.SQLStore,
+				sqlstoreProviderFactories,
+				config.SQLStore.Provider,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Initialize sqlschema from the available sqlschema provider factories
+			sqlschema, err := factory.NewProviderFromNamedMap(
+				ctx,
+				providerSettings,
+				config.SQLSchema,
+				sqlSchemaProviderFactories(sqlstore),
+				config.SQLStore.Provider,
+			)
+			if err != nil {
+				return err
+			}
+
+			// Run migrations on the sqlstore
+			sqlmigrations, err := sqlmigration.New(
+				ctx,
+				providerSettings,
+				config.SQLMigration,
+				signoz.NewSQLMigrationProviderFactories(sqlstore, sqlschema),
+			)
+			if err != nil {
+				return err
+			}
+
+			err = sqlmigrator.New(ctx, providerSettings, sqlstore, sqlmigrations, config.SQLMigrator).Rollback(ctx)
+			if err != nil {
+				return err
+			}
+
+			return nil
+		},
+	}
+
+	sqlCmd.AddCommand(migrateCmd)
+	sqlCmd.AddCommand(rollbackCmd)
+	parentCmd.AddCommand(sqlCmd)
+}

pkg/sqlmigration/s100sqlmigration/100_v100.go

@@ -467,7 +467,7 @@ func (migration *v100) Up(ctx context.Context, db *bun.DB) error {
 				{Name: sqlschema.ColumnName("org_id"), DataType: sqlschema.DataTypeText, Nullable: false, Default: ""},
 				{Name: sqlschema.ColumnName("version"), DataType: sqlschema.DataTypeInteger, Nullable: true, Default: ""},
 				{Name: sqlschema.ColumnName("element_type"), DataType: sqlschema.DataTypeText, Nullable: false, Default: ""},
-				{Name: sqlschema.ColumnName("deploy_status"), DataType: sqlschema.DataTypeText, Nullable: false, Default: "dirty"},
+				{Name: sqlschema.ColumnName("deploy_status"), DataType: sqlschema.DataTypeText, Nullable: false, Default: "'dirty'"},
 				{Name: sqlschema.ColumnName("deploy_sequence"), DataType: sqlschema.DataTypeInteger, Nullable: true, Default: ""},
 				{Name: sqlschema.ColumnName("deploy_result"), DataType: sqlschema.DataTypeText, Nullable: true, Default: ""},
 				{Name: sqlschema.ColumnName("hash"), DataType: sqlschema.DataTypeText, Nullable: true, Default: ""},

pkg/sqlmigrator/migrator.go

@@ -53,7 +53,7 @@ func (migrator *migrator) Migrate(ctx context.Context) error {
 
 	group, err := migrator.migrator.Migrate(ctx)
 	if err != nil {
-		migrator.settings.Logger().ErrorContext(ctx, "failed to run sqlstore migrations", "error", err, "dialect", migrator.dialect, "group", group.String(), "applied", group.Migrations.Applied())
+		migrator.settings.Logger().ErrorContext(ctx, "failed to run sqlstore migrations", "error", err, "dialect", migrator.dialect, "group", group.String(), "applied", group.Migrations.Applied().String())
 		return err
 	}
 
@@ -82,7 +82,7 @@ func (migrator *migrator) Rollback(ctx context.Context) error {
 		return nil
 	}
 
-	migrator.settings.Logger().InfoContext(ctx, "rolled back", "group", group.String(), "dialect", migrator.dialect)
+	migrator.settings.Logger().InfoContext(ctx, "rolled back", "group", group.String(), "dialect", migrator.dialect, "applied", group.Migrations.Applied().String())
 	return nil
 }
 

pkg/sqlschema/operator.go

@@ -67,14 +67,22 @@ func (operator *Operator) AlterTable(oldTable *Table, oldTableUniqueConstraints
 		sql = append(sql, operator.AlterColumn(oldTable, oldTableUniqueConstraints, column)...)
 	}
 
-	// Primary key constraints
+	// If the old table does not have a primary key constraint and the new table does, we need to create it.
-	if oldTable.PrimaryKeyConstraint != newTable.PrimaryKeyConstraint {
+	if oldTable.PrimaryKeyConstraint == nil {
+		if newTable.PrimaryKeyConstraint != nil {
+			sql = append(sql, operator.CreateConstraint(oldTable, oldTableUniqueConstraints, newTable.PrimaryKeyConstraint)...)
+		}
+	}
+
+	if oldTable.PrimaryKeyConstraint != nil {
 		if newTable.PrimaryKeyConstraint == nil {
 			sql = append(sql, operator.DropConstraint(oldTable, oldTableUniqueConstraints, oldTable.PrimaryKeyConstraint)...)
 		} else {
+			if !oldTable.PrimaryKeyConstraint.Equals(newTable.PrimaryKeyConstraint) {
 				sql = append(sql, operator.CreateConstraint(oldTable, oldTableUniqueConstraints, newTable.PrimaryKeyConstraint)...)
 			}
 		}
+	}
 
 	// Drop foreign key constraints that are in the old table but not in the new table.
 	for _, fkConstraint := range oldTable.ForeignKeyConstraints {
@@ -248,6 +256,8 @@ func (operator *Operator) CreateConstraint(table *Table, uniqueConstraints []*Un
 	if constraint.Type() == ConstraintTypePrimaryKey {
 		if operator.support.SCreateAndDropConstraint {
 			if table.PrimaryKeyConstraint != nil {
+				// TODO(grandwizard28): this is a hack to drop the primary key constraint.
+				// We need to find a better way to do this.
 				sqls = append(sqls, table.PrimaryKeyConstraint.ToDropSQL(operator.fmter, table.Name))
 			}
 		}

pkg/sqlstore/bun.go

@@ -19,6 +19,7 @@ type BunDB struct {
 
 func NewBunDB(settings factory.ScopedProviderSettings, sqldb *sql.DB, dialect schema.Dialect, hooks []SQLStoreHook, opts ...bun.DBOption) *BunDB {
 	db := bun.NewDB(sqldb, dialect, opts...)
+	bun.SetLogger(&bunLogger{settings.Logger()})
 
 	for _, hook := range hooks {
 		db.AddQueryHook(hook)

pkg/sqlstore/logger.go

@@ -0,0 +1,12 @@
+package sqlstore
+
+import "log/slog"
+
+type bunLogger struct {
+	logger *slog.Logger
+}
+
+func (l *bunLogger) Printf(format string, v ...interface{}) {
+	// the no lint directive is needed because the bun logger does not accept context
+	l.logger.Info(format, v...) //nolint:sloglint
+}