From fcfe5bd37c8c19c5b1e55cdb6304eab63bed66f7 Mon Sep 17 00:00:00 2001 From: Rogee Date: Thu, 13 Aug 2026 17:13:17 +0800 Subject: [PATCH] H-60: harden Captain migration rollback and concurrency (#10) Co-authored-by: Rogee --- .../captain_binding_migration_test.go | 86 +++++++++++++++++++ .../service/captain_binding_atomicity_test.go | 58 +++++++++++-- ...0079_make_captain_bindings_unique.down.sql | 11 ++- ...000079_make_captain_bindings_unique.up.sql | 7 ++ 4 files changed, 153 insertions(+), 9 deletions(-) create mode 100644 backend/internal/database/captain_binding_migration_test.go diff --git a/backend/internal/database/captain_binding_migration_test.go b/backend/internal/database/captain_binding_migration_test.go new file mode 100644 index 00000000..02ee8ab7 --- /dev/null +++ b/backend/internal/database/captain_binding_migration_test.go @@ -0,0 +1,86 @@ +package database + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func TestCaptainBindingMigrationDeduplicatesAndRejectsDown(t *testing.T) { + if os.Getenv("GOCHAT_TEST_DB") == "sqlite" { + t.Skip("requires PostgreSQL migration semantics") + } + dsn := os.Getenv("GOCHAT_TEST_DB_URL") + if dsn == "" { + dsn = "host=localhost port=5432 user=postgres password=postgres dbname=gochat_test sslmode=disable" + } + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + + schema := "captain_migration_" + time.Now().Format("20060102150405000000000") + require.NoError(t, db.Exec("CREATE SCHEMA "+schema).Error) + t.Cleanup(func() { _ = db.Exec("DROP SCHEMA " + schema + " CASCADE").Error }) + require.NoError(t, db.Exec("SET search_path TO "+schema).Error) + require.NoError(t, db.Exec(` + CREATE TABLE agent_bots (id BIGINT PRIMARY KEY, account_id BIGINT, bot_type TEXT, config JSONB); + CREATE TABLE agent_bot_inboxes (id BIGINT PRIMARY KEY, agent_bot_id BIGINT, inbox_id BIGINT); + CREATE TABLE captain_inboxes (id BIGINT PRIMARY KEY, inbox_id BIGINT, deleted_at TIMESTAMPTZ); + CREATE TABLE conversations (id BIGINT PRIMARY KEY, assignee_agent_bot_id BIGINT); + CREATE TABLE agent_bot_presence_events (id BIGINT PRIMARY KEY, agent_bot_id BIGINT); + CREATE TABLE bot_rules (id BIGINT PRIMARY KEY, agent_bot_id BIGINT); + CREATE TABLE bot_trigger_configs (id BIGINT PRIMARY KEY, agent_bot_id BIGINT); + INSERT INTO agent_bots VALUES + (10, 1, 'captain', '{"assistant_id":7}'), + (11, 1, 'captain', '{"assistant_id":7}'); + INSERT INTO agent_bot_inboxes VALUES (20, 10, 42), (21, 11, 42); + INSERT INTO captain_inboxes VALUES (30, 42, NULL), (31, 42, NULL); + INSERT INTO conversations VALUES (40, 11); + INSERT INTO agent_bot_presence_events VALUES (50, 11); + INSERT INTO bot_rules VALUES (60, 11); + INSERT INTO bot_trigger_configs VALUES (70, 11); + `).Error) + + up, err := os.ReadFile(filepath.Join("..", "..", "migrations", "000079_make_captain_bindings_unique.up.sql")) + require.NoError(t, err) + require.NoError(t, db.Transaction(func(tx *gorm.DB) error { return tx.Exec(string(up)).Error })) + + for _, table := range []string{"agent_bots", "agent_bot_inboxes", "captain_inboxes"} { + var count int64 + require.NoError(t, db.Table(table).Count(&count).Error) + assert.Equal(t, int64(1), count, table) + } + for table, column := range map[string]string{ + "conversations": "assignee_agent_bot_id", + "agent_bot_presence_events": "agent_bot_id", + "bot_rules": "agent_bot_id", + "bot_trigger_configs": "agent_bot_id", + } { + var botID int64 + require.NoError(t, db.Table(table).Select(column).Scan(&botID).Error) + assert.Equal(t, int64(10), botID, table) + } + for table, index := range map[string]string{ + "agent_bots": "idx_agent_bots_captain_assistant", + "agent_bot_inboxes": "idx_agent_bot_inboxes_bot_inbox", + "captain_inboxes": "idx_captain_inboxes_active_inbox", + } { + assert.True(t, db.Migrator().HasIndex(table, index), index) + } + + down, err := os.ReadFile(filepath.Join("..", "..", "migrations", "000079_make_captain_bindings_unique.down.sql")) + require.NoError(t, err) + err = db.Transaction(func(tx *gorm.DB) error { return tx.Exec(string(down)).Error }) + require.ErrorContains(t, err, "migration 000079 is irreversible") + assert.True(t, db.Migrator().HasColumn("agent_bots", "captain_assistant_id")) +} diff --git a/backend/internal/service/captain_binding_atomicity_test.go b/backend/internal/service/captain_binding_atomicity_test.go index 4fc69ddf..363c4078 100644 --- a/backend/internal/service/captain_binding_atomicity_test.go +++ b/backend/internal/service/captain_binding_atomicity_test.go @@ -3,13 +3,18 @@ package service import ( "context" "errors" + "fmt" + "net/url" + "os" + "strings" "sync" "testing" + "time" "github.com/gochat/gochat/internal/model" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "gorm.io/driver/sqlite" + "gorm.io/driver/postgres" "gorm.io/gorm" "gorm.io/gorm/logger" ) @@ -42,31 +47,74 @@ func TestWidgetConversationRollsBackCaptainBindingOnCreateFailure(t *testing.T) } func TestEnsureCaptainAgentBotBindingConcurrentCallsStayUnique(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + if os.Getenv("GOCHAT_TEST_DB") == "sqlite" { + t.Skip("requires PostgreSQL conflict handling") + } + dsn := os.Getenv("GOCHAT_TEST_DB_URL") + if dsn == "" { + dsn = "host=localhost port=5432 user=postgres password=postgres dbname=gochat_test sslmode=disable" + } + admin, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + schema := fmt.Sprintf("captain_binding_%d", time.Now().UnixNano()) + require.NoError(t, admin.Exec("CREATE SCHEMA "+schema).Error) + t.Cleanup(func() { + _ = admin.Exec("DROP SCHEMA " + schema + " CASCADE").Error + if sqlDB, dbErr := admin.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") { + dsnURL, parseErr := url.Parse(dsn) + require.NoError(t, parseErr) + query := dsnURL.Query() + query.Set("search_path", schema) + dsnURL.RawQuery = query.Encode() + dsn = dsnURL.String() + } else { + dsn += " search_path=" + schema + } + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + DisableForeignKeyConstraintWhenMigrating: true, + }) require.NoError(t, err) require.NoError(t, db.AutoMigrate(&model.CaptainAssistant{}, &model.AgentBot{}, &model.AgentBotInbox{})) sqlDB, err := db.DB() require.NoError(t, err) - sqlDB.SetMaxOpenConns(1) + const calls = 8 + sqlDB.SetMaxOpenConns(calls) t.Cleanup(func() { _ = sqlDB.Close() }) assistant := &model.CaptainAssistant{AccountID: 1, Name: "Concurrent", Status: model.AssistantStatusActive} require.NoError(t, db.Create(assistant).Error) - const calls = 8 errs := make(chan error, calls) ids := make(chan uint, calls) + ready := make(chan struct{}, calls) + start := make(chan struct{}) var wg sync.WaitGroup for range calls { wg.Add(1) go func() { defer wg.Done() - bot, err := ensureCaptainAgentBotBinding(context.Background(), db, assistant, 42) + var bot *model.AgentBot + err := db.Transaction(func(tx *gorm.DB) error { + ready <- struct{}{} + <-start + var bindErr error + bot, bindErr = ensureCaptainAgentBotBinding(context.Background(), tx, assistant, 42) + return bindErr + }) errs <- err if bot != nil { ids <- bot.ID } }() } + for range calls { + <-ready + } + close(start) wg.Wait() close(errs) close(ids) diff --git a/backend/migrations/000079_make_captain_bindings_unique.down.sql b/backend/migrations/000079_make_captain_bindings_unique.down.sql index f74b2e12..172430ea 100644 --- a/backend/migrations/000079_make_captain_bindings_unique.down.sql +++ b/backend/migrations/000079_make_captain_bindings_unique.down.sql @@ -1,4 +1,7 @@ -DROP INDEX IF EXISTS idx_captain_inboxes_active_inbox; -DROP INDEX IF EXISTS idx_agent_bot_inboxes_bot_inbox; -DROP INDEX IF EXISTS idx_agent_bots_captain_assistant; -ALTER TABLE agent_bots DROP COLUMN IF EXISTS captain_assistant_id; +-- The up migration merges duplicate bots/bindings and rewrites references. +-- Those deleted rows and original references cannot be reconstructed safely. +DO $$ +BEGIN + RAISE EXCEPTION 'migration 000079 is irreversible: restore from backup instead of migrating down'; +END +$$; diff --git a/backend/migrations/000079_make_captain_bindings_unique.up.sql b/backend/migrations/000079_make_captain_bindings_unique.up.sql index ed0e12dc..9f15ddd6 100644 --- a/backend/migrations/000079_make_captain_bindings_unique.up.sql +++ b/backend/migrations/000079_make_captain_bindings_unique.up.sql @@ -1,5 +1,12 @@ ALTER TABLE agent_bots ADD COLUMN IF NOT EXISTS captain_assistant_id BIGINT; +-- This data repair and its unique indexes must see a stable snapshot. These +-- locks block concurrent writes while allowing reads, so run this migration +-- in a maintenance window sized for the deduplication. +LOCK TABLE agent_bots, agent_bot_inboxes, captain_inboxes, conversations, + agent_bot_presence_events, bot_rules, bot_trigger_configs + IN SHARE ROW EXCLUSIVE MODE; + UPDATE agent_bots SET captain_assistant_id = (config ->> 'assistant_id')::BIGINT WHERE bot_type = 'captain'