From 8dc749b9dd4977774fb76cf77683165ec47e4ed8 Mon Sep 17 00:00:00 2001 From: Vibhu Pandey Date: Wed, 7 May 2025 13:48:13 +0530 Subject: [PATCH] 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 --- ee/sqlstore/postgressqlstore/provider.go | 4 +++ pkg/sqlmigration/029_drop_groups.go | 33 +++++++++++++++++++----- pkg/sqlstore/sqlitesqlstore/dialect.go | 25 ++++++++++++------ pkg/sqlstore/sqlstore.go | 7 +++++ pkg/sqlstore/sqlstoretest/dialect.go | 8 ++++++ 5 files changed, 62 insertions(+), 15 deletions(-) diff --git a/ee/sqlstore/postgressqlstore/provider.go b/ee/sqlstore/postgressqlstore/provider.go index 6c098b448f..dd2797a1fa 100644 --- a/ee/sqlstore/postgressqlstore/provider.go +++ b/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 +} diff --git a/pkg/sqlmigration/029_drop_groups.go b/pkg/sqlmigration/029_drop_groups.go index 211b9b68ec..31f225cb9f 100644 --- a/pkg/sqlmigration/029_drop_groups.go +++ b/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 } diff --git a/pkg/sqlstore/sqlitesqlstore/dialect.go b/pkg/sqlstore/sqlitesqlstore/dialect.go index a27f4b7763..fe5ceae147 100644 --- a/pkg/sqlstore/sqlitesqlstore/dialect.go +++ b/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 } diff --git a/pkg/sqlstore/sqlstore.go b/pkg/sqlstore/sqlstore.go index 072a06c3f4..79a392d280 100644 --- a/pkg/sqlstore/sqlstore.go +++ b/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 } diff --git a/pkg/sqlstore/sqlstoretest/dialect.go b/pkg/sqlstore/sqlstoretest/dialect.go index 897f946ed8..087ed1444e 100644 --- a/pkg/sqlstore/sqlstoretest/dialect.go +++ b/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 +}