diff --git a/server/config/audit_store.go b/server/config/audit_store.go index 8adb9d97..00f61931 100644 --- a/server/config/audit_store.go +++ b/server/config/audit_store.go @@ -125,6 +125,9 @@ func (w *FileAuditWriter) Query(ctx context.Context, query ConfigAuditQuery) ([] scanner.Buffer(make([]byte, 4096), 256*1024) lineNumber := 0 for scanner.Scan() { + if err := auditContextError(ctx); err != nil { + return nil, err + } lineNumber++ var event ConfigAuditEvent if err := json.Unmarshal(scanner.Bytes(), &event); err != nil { @@ -214,6 +217,9 @@ func (w *FileAuditWriter) Cleanup(ctx context.Context, now time.Time) (int, erro scanner.Buffer(make([]byte, 4096), 256*1024) lineNumber := 0 for scanner.Scan() { + if err := auditContextError(ctx); err != nil { + return deleted, err + } lineNumber++ var event ConfigAuditEvent if err := json.Unmarshal(scanner.Bytes(), &event); err != nil { @@ -241,6 +247,9 @@ func (w *FileAuditWriter) Cleanup(ctx context.Context, now time.Time) (int, erro if err := file.Close(); err != nil { return deleted, ErrAuditStorageUnavailable } + if err := auditContextError(ctx); err != nil { + return deleted, err + } if deleted == 0 { return 0, nil } diff --git a/server/config/security.go b/server/config/security.go index 4bc41b77..74794705 100644 --- a/server/config/security.go +++ b/server/config/security.go @@ -94,6 +94,9 @@ func (w *MemoryAuditWriter) Query(ctx context.Context, query ConfigAuditQuery) ( defer w.mu.Unlock() result := make([]ConfigAuditEvent, 0, minAuditLimit(limit, len(w.events))) for _, event := range w.events { + if err := auditContextError(ctx); err != nil { + return nil, err + } if !auditEventMatches(event, query) { continue } diff --git a/server/services/dbmigration/service.go b/server/services/dbmigration/service.go index 9db63777..e6bebdd9 100644 --- a/server/services/dbmigration/service.go +++ b/server/services/dbmigration/service.go @@ -170,6 +170,9 @@ func (s *Service) execute(ctx context.Context, id, token string, actor config.Au if s == nil || strings.TrimSpace(id) == "" || strings.TrimSpace(token) == "" { return Record{}, ErrInvalidRequest } + if ctx == nil { + ctx = context.Background() + } now := s.currentTime() s.mu.Lock() item, ok := s.records[id] diff --git a/server/services/dbmigration/sql_executor.go b/server/services/dbmigration/sql_executor.go index 5e3f547c..83b7f26c 100644 --- a/server/services/dbmigration/sql_executor.go +++ b/server/services/dbmigration/sql_executor.go @@ -184,15 +184,21 @@ func (c *SQLCopier) Copy(ctx context.Context, source, target config.DatabaseConf if ctx == nil { ctx = context.Background() } + if err := contextError(ctx); err != nil { + return err + } adminDB, err := c.open(ctx, source, "") - if err != nil { + if err != nil || adminDB == nil { return ErrSQLConnectionUnavailable } defer adminDB.Close() + if err := contextError(ctx); err != nil { + return err + } switch source.Type { case config.DatabaseTypeMySQL: sourceDB, err := c.open(ctx, source, source.Name) - if err != nil { + if err != nil || sourceDB == nil { return ErrSQLConnectionUnavailable } defer sourceDB.Close() @@ -243,6 +249,9 @@ func (c *SQLCopier) copyMySQL(ctx context.Context, sourceDB, adminDB *sql.DB, so tableNames := make([]string, 0) sourceRowCounts := make(map[string]int64) for rows.Next() { + if err := contextError(ctx); err != nil { + return err + } var tableName, tableType string if err := rows.Scan(&tableName, &tableType); err != nil { return fmt.Errorf("read mysql table: %w", err) @@ -266,6 +275,9 @@ func (c *SQLCopier) copyMySQL(ctx context.Context, sourceDB, adminDB *sql.DB, so return fmt.Errorf("read mysql tables: %w", err) } for _, tableName := range tableNames { + if err := contextError(ctx); err != nil { + return err + } sourceTable := quoteMySQLIdentifier(sourceName) + "." + quoteMySQLIdentifier(tableName) targetTable := quoteMySQLIdentifier(targetName) + "." + quoteMySQLIdentifier(tableName) if _, err := adminDB.ExecContext(ctx, "INSERT INTO "+targetTable+" SELECT * FROM "+sourceTable); err != nil { @@ -343,6 +355,9 @@ func (c *SQLCopier) copyMySQLForeignKeys(ctx context.Context, database *sql.DB, foreignKeys := make([]mysqlForeignKey, 0) indexes := make(map[string]int) for rows.Next() { + if err := contextError(ctx); err != nil { + return err + } var item mysqlForeignKey var column, referencedColumn string if err := rows.Scan(&item.tableName, &item.constraintName, &column, &item.referencedSchema, &item.referencedTable, &referencedColumn, &item.updateRule, &item.deleteRule); err != nil { @@ -365,6 +380,9 @@ func (c *SQLCopier) copyMySQLForeignKeys(ctx context.Context, database *sql.DB, return fmt.Errorf("%w: read mysql foreign keys", ErrSQLConnectionUnavailable) } for _, foreignKey := range foreignKeys { + if err := contextError(ctx); err != nil { + return err + } referencedSchema := foreignKey.referencedSchema if referencedSchema == sourceName { referencedSchema = targetName diff --git a/server/services/dbmigration/sql_executor_test.go b/server/services/dbmigration/sql_executor_test.go index 337298d5..25e28601 100644 --- a/server/services/dbmigration/sql_executor_test.go +++ b/server/services/dbmigration/sql_executor_test.go @@ -2,6 +2,7 @@ package dbmigration import ( "context" + "database/sql" "errors" "net/url" "strings" @@ -39,6 +40,31 @@ func validStructuredPlan(t *testing.T) Plan { } } +func TestSQLCopierRejectsNilDatabaseWithoutPanicking(t *testing.T) { + copier := NewSQLCopier(func(context.Context, config.DatabaseConfig, string) (*sql.DB, error) { + return nil, nil + }) + if err := copier.Copy(context.Background(), testDatabase(t), testDatabase(t)); !errors.Is(err, ErrSQLConnectionUnavailable) { + t.Fatalf("nil admin database error = %v, want connection unavailable", err) + } +} + +func TestSQLCopierStopsAfterAdminConnectionWhenContextIsCanceled(t *testing.T) { + openCalls := 0 + copier := NewSQLCopier(func(context.Context, config.DatabaseConfig, string) (*sql.DB, error) { + openCalls++ + return nil, errors.New("open should not continue") + }) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := copier.Copy(ctx, testDatabase(t), testDatabase(t)); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled copy error = %v, want context canceled", err) + } + if openCalls != 0 { + t.Fatalf("open calls = %d, want zero for canceled context", openCalls) + } +} + func TestSQLExecutorCopiesAndSwitchesWithOrderedStatuses(t *testing.T) { copier := &recordingStructuredCopier{} var actor config.AuditActor