fix(migration): fix cascading drops in sqlite (#7844)

* fix(foreign-key): fix cascading drops in sqlite

* fix(foreign-key): fix comments

* fix(foreign-key): fix function names

* fix(foreign-key): fix order of migration

---------

Co-authored-by: Vikrant Gupta <vikrant@signoz.io>

ee/sqlstore/postgressqlstore/provider.go

@@ -106,3 +106,7 @@ func (provider *provider) WrapAlreadyExistsErrf(err error, code errors.Code, for
 
 	return err
 }
+
+func (dialect *dialect) ToggleForeignKeyConstraint(ctx context.Context, bun *bun.DB, enable bool) error {
+	return nil
+}

pkg/sqlmigration/029_drop_groups.go

@@ -33,13 +33,6 @@ func (migration *dropGroups) Register(migrations *migrate.Migrations) error {
 }
 
 func (migration *dropGroups) Up(ctx context.Context, db *bun.DB) error {
-	tx, err := db.BeginTx(ctx, nil)
-	if err != nil {
-		return err
-	}
-
-	defer tx.Rollback()
-
 	type Group struct {
 		bun.BaseModel `bun:"table:groups"`
 
@@ -49,6 +42,27 @@ func (migration *dropGroups) Up(ctx context.Context, db *bun.DB) error {
 		Name  string `bun:"name,type:text,notnull,unique" json:"name"`
 	}
 
+	exists, err := migration.sqlstore.Dialect().TableExists(ctx, db, new(Group))
+	if err != nil {
+		return err
+	}
+
+	if !exists {
+		return nil
+	}
+
+	// Disable foreign keys temporarily
+	if err := migration.sqlstore.Dialect().ToggleForeignKeyConstraint(ctx, db, false); err != nil {
+		return err
+	}
+
+	tx, err := db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	defer tx.Rollback()
+
 	type existingUser struct {
 		bun.BaseModel `bun:"table:users"`
 
@@ -133,6 +147,11 @@ func (migration *dropGroups) Up(ctx context.Context, db *bun.DB) error {
 		return err
 	}
 
+	// Enable foreign keys
+	if err := migration.sqlstore.Dialect().ToggleForeignKeyConstraint(ctx, db, true); err != nil {
+		return err
+	}
+
 	return nil
 }
 

pkg/sqlstore/sqlitesqlstore/dialect.go

@@ -428,6 +428,15 @@ func (dialect *dialect) AddPrimaryKey(ctx context.Context, bun bun.IDB, oldModel
 }
 
 func (dialect *dialect) DropColumnWithForeignKeyConstraint(ctx context.Context, bunIDB bun.IDB, model interface{}, column string) error {
+	var isForeignKeyEnabled bool
+	if err := bunIDB.QueryRowContext(ctx, "PRAGMA foreign_keys").Scan(&isForeignKeyEnabled); err != nil {
+		return err
+	}
+
+	if isForeignKeyEnabled {
+		return errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput, "foreign keys are enabled, please disable them before running this migration")
+	}
+
 	existingTable := bunIDB.Dialect().Tables().Get(reflect.TypeOf(model))
 	columnExists, err := dialect.ColumnExists(ctx, bunIDB, existingTable.Name, column)
 	if err != nil {
@@ -455,11 +464,6 @@ func (dialect *dialect) DropColumnWithForeignKeyConstraint(ctx context.Context,
 		}
 	}
 
-	// Disable foreign keys temporarily
-	if _, err := bunIDB.ExecContext(ctx, "PRAGMA foreign_keys = OFF"); err != nil {
-		return err
-	}
-
 	if _, err = createTableQuery.Exec(ctx); err != nil {
 		return err
 	}
@@ -479,10 +483,15 @@ func (dialect *dialect) DropColumnWithForeignKeyConstraint(ctx context.Context,
 		return err
 	}
 
-	// Re-enable foreign keys
-	if _, err := bunIDB.ExecContext(ctx, "PRAGMA foreign_keys = ON"); err != nil {
+	return nil
+}
+
+func (dialect *dialect) ToggleForeignKeyConstraint(ctx context.Context, bun *bun.DB, enable bool) error {
+	if enable {
+		_, err := bun.ExecContext(ctx, "PRAGMA foreign_keys = ON")
 		return err
 	}
 
-	return nil
+	_, err := bun.ExecContext(ctx, "PRAGMA foreign_keys = OFF")
+	return err
 }

pkg/sqlstore/sqlstore.go

@@ -82,4 +82,11 @@ type SQLDialect interface {
 
 	// Drops the column and the associated foreign key constraint for the given table and column.
 	DropColumnWithForeignKeyConstraint(context.Context, bun.IDB, interface{}, string) error
+
+	// Checks if a table exists.
+	TableExists(ctx context.Context, bun bun.IDB, table interface{}) (bool, error)
+
+	// Toggles foreign key constraint for the given database. This makes sense only for sqlite. This cannot take a transaction as an argument and needs to take the db
+	// as an argument.
+	ToggleForeignKeyConstraint(ctx context.Context, bun *bun.DB, enable bool) error
 }

pkg/sqlstore/sqlstoretest/dialect.go

@@ -60,3 +60,11 @@ func (dialect *dialect) IndexExists(ctx context.Context, bun bun.IDB, table stri
 func (dialect *dialect) DropColumnWithForeignKeyConstraint(ctx context.Context, bun bun.IDB, model interface{}, column string) error {
 	return nil
 }
+
+func (dialect *dialect) TableExists(ctx context.Context, bun bun.IDB, table interface{}) (bool, error) {
+	return true, nil
+}
+
+func (dialect *dialect) ToggleForeignKeyConstraint(ctx context.Context, bun *bun.DB, enable bool) error {
+	return nil
+}