H-60: harden Captain migration rollback and concurrency (#10)

Co-authored-by: Rogee <rogee@ipao.vip>
This commit is contained in:
Rogee
2026-08-13 17:13:17 +08:00
committed by GitHub
co-authored by rogee
parent 0540ae3791
commit fcfe5bd37c
4 changed files with 153 additions and 9 deletions
@@ -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"))
}
@@ -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)
@@ -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
$$;
@@ -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'