Skip to content

Commit 3cb29d6

Browse files
luodaoyiluodaoyi
authored
fix: handle preexisting codex quota identity index during migration (awsl-project#429)
1 parent e8ae197 commit 3cb29d6

File tree

3 files changed

+177
-9
lines changed

3 files changed

+177
-9
lines changed

‎internal/repository/sqlite/db_test.go‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,3 +143,77 @@ func TestNewDBWithDSN_CodexQuotaMigrationPreservesDeletedHistory(t *testing.T) {
143143
t.Fatalf("expected deleted row to be preserved after migration, got %d", deletedCount)
144144
}
145145
}
146+
147+
func TestNewDBWithDSN_CodexQuotaMigrationHandlesPreexistingIdentityIndex(t *testing.T) {
148+
dsn := filepath.Join(t.TempDir(), "codex-quota-preexisting-identity-index.db")
149+
gormDB, err := gorm.Open(gormsqlite.Open(dsn), &gorm.Config{})
150+
if err != nil {
151+
t.Fatalf("open seed db: %v", err)
152+
}
153+
sqlDB, err := gormDB.DB()
154+
if err != nil {
155+
t.Fatalf("get seed sql.DB: %v", err)
156+
}
157+
defer sqlDB.Close()
158+
159+
seedSQL := []string{
160+
`CREATE TABLE codex_quotas (
161+
id TEXT PRIMARY KEY,
162+
created_at INTEGER DEFAULT 0,
163+
updated_at INTEGER DEFAULT 0,
164+
deleted_at INTEGER DEFAULT 0,
165+
tenant_id INTEGER NOT NULL,
166+
identity_key TEXT,
167+
email TEXT,
168+
account_id TEXT,
169+
plan_type TEXT,
170+
is_forbidden INTEGER DEFAULT 0,
171+
primary_window TEXT,
172+
secondary_window TEXT,
173+
code_review_window TEXT
174+
)`,
175+
`CREATE UNIQUE INDEX idx_codex_quotas_tenant_identity ON codex_quotas(tenant_id, identity_key)`,
176+
`CREATE UNIQUE INDEX idx_codex_quotas_email ON codex_quotas(email)`,
177+
`CREATE UNIQUE INDEX idx_codex_quotas_tenant_email ON codex_quotas(tenant_id, email)`,
178+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-8', 1, NULL, 'cnc6n2io9xvfev2mtm5t6hu8@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456445476)`,
179+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-11', 1, NULL, 'fcew8ua8r6u6zekrwqfi60nl@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456444987)`,
180+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-7', 1, NULL, 'puckxqnzu6ktt7k4bcevlw06@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456443639)`,
181+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-9', 1, NULL, 'n7e87dj5hxv2c2m4u0e6l5zo@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456443366)`,
182+
}
183+
for _, sql := range seedSQL {
184+
if err := gormDB.Exec(sql).Error; err != nil {
185+
t.Fatalf("seed fixture: %v", err)
186+
}
187+
}
188+
189+
db, err := NewDBWithDSN("sqlite://" + dsn)
190+
if err != nil {
191+
t.Fatalf("NewDBWithDSN should survive preexisting identity index: %v", err)
192+
}
193+
defer db.Close()
194+
195+
var count int64
196+
if err := db.GormDB().Raw(`
197+
SELECT COUNT(*)
198+
FROM codex_quotas
199+
WHERE tenant_id = 1 AND identity_key = 'account:e94ce011-80f4-490b-b285-d3109db72b0e' AND deleted_at = 0
200+
`).Scan(&count).Error; err != nil {
201+
t.Fatalf("count migrated rows: %v", err)
202+
}
203+
if count != 1 {
204+
t.Fatalf("expected rows to collapse to one active account identity, got %d", count)
205+
}
206+
207+
var rows []struct {
208+
Name string `gorm:"column:name"`
209+
Unique int `gorm:"column:unique"`
210+
}
211+
if err := db.GormDB().Raw(`PRAGMA index_list('codex_quotas')`).Scan(&rows).Error; err != nil {
212+
t.Fatalf("list indexes: %v", err)
213+
}
214+
for _, row := range rows {
215+
if row.Name == "idx_codex_quotas_email" && row.Unique != 0 {
216+
t.Fatalf("expected idx_codex_quotas_email to be recreated as non-unique, got unique=%d", row.Unique)
217+
}
218+
}
219+
}

‎internal/repository/sqlite/migrations.go‎

Lines changed: 30 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,9 @@ func applyCodexQuotaIdentityMigration(db *gorm.DB) error {
247247
if !db.Migrator().HasColumn(&CodexQuota{}, "identity_key") {
248248
return nil
249249
}
250+
if err := dropCodexQuotaIdentityIndexesBeforeBackfill(db); err != nil {
251+
return err
252+
}
250253
if err := db.Exec(codexQuotaIdentityBackfillSQL(db.Dialector.Name())).Error; err != nil {
251254
return err
252255
}
@@ -339,22 +342,19 @@ func codexQuotaIdentityDedupeSQL(dialector string) string {
339342
func ensureCodexQuotaIdentityIndexes(db *gorm.DB) error {
340343
switch db.Dialector.Name() {
341344
case "mysql":
342-
if err := db.Exec("DROP INDEX idx_codex_quotas_tenant_identity ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
343-
return err
344-
}
345345
if err := db.Exec("CREATE UNIQUE INDEX idx_codex_quotas_tenant_identity ON codex_quotas(tenant_id, identity_key, deleted_at)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
346346
return err
347347
}
348348
if err := db.Exec("DROP INDEX idx_codex_quotas_tenant_email ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
349349
return err
350350
}
351-
if err := db.Exec("CREATE INDEX idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
351+
if err := db.Exec("DROP INDEX idx_codex_quotas_email ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
352352
return err
353353
}
354-
case "postgres":
355-
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_identity").Error; err != nil {
354+
if err := db.Exec("CREATE INDEX idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil && !isMySQLDuplicateIndexError(err) {
356355
return err
357356
}
357+
case "postgres":
358358
if err := db.Exec(`
359359
CREATE UNIQUE INDEX IF NOT EXISTS idx_codex_quotas_tenant_identity
360360
ON codex_quotas(tenant_id, identity_key)
@@ -365,13 +365,13 @@ func ensureCodexQuotaIdentityIndexes(db *gorm.DB) error {
365365
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_email").Error; err != nil {
366366
return err
367367
}
368-
if err := db.Exec("CREATE INDEX IF NOT EXISTS idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil {
368+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_email").Error; err != nil {
369369
return err
370370
}
371-
default:
372-
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_identity").Error; err != nil {
371+
if err := db.Exec("CREATE INDEX IF NOT EXISTS idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil {
373372
return err
374373
}
374+
default:
375375
if err := db.Exec(`
376376
CREATE UNIQUE INDEX IF NOT EXISTS idx_codex_quotas_tenant_identity
377377
ON codex_quotas(tenant_id, identity_key)
@@ -382,13 +382,34 @@ func ensureCodexQuotaIdentityIndexes(db *gorm.DB) error {
382382
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_email").Error; err != nil {
383383
return err
384384
}
385+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_email").Error; err != nil {
386+
return err
387+
}
385388
if err := db.Exec("CREATE INDEX IF NOT EXISTS idx_codex_quotas_email ON codex_quotas(email)").Error; err != nil {
386389
return err
387390
}
388391
}
389392
return nil
390393
}
391394

395+
func dropCodexQuotaIdentityIndexesBeforeBackfill(db *gorm.DB) error {
396+
switch db.Dialector.Name() {
397+
case "mysql":
398+
if err := db.Exec("DROP INDEX idx_codex_quotas_tenant_identity ON codex_quotas").Error; err != nil && !isMySQLMissingIndexError(err) {
399+
return err
400+
}
401+
case "postgres":
402+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_identity").Error; err != nil {
403+
return err
404+
}
405+
default:
406+
if err := db.Exec("DROP INDEX IF EXISTS idx_codex_quotas_tenant_identity").Error; err != nil {
407+
return err
408+
}
409+
}
410+
return nil
411+
}
412+
392413
func isMySQLDuplicateIndexError(err error) bool {
393414
if err == nil {
394415
return false

‎internal/repository/sqlite/migrations_test.go‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,33 @@ func TestCodexQuotaIdentityMigrationV8UpPreservesDeletedHistory(t *testing.T) {
103103
}
104104
}
105105

106+
func TestCodexQuotaIdentityMigrationV8UpHandlesPreexistingIdentityIndex(t *testing.T) {
107+
gormDB := openRawSQLiteDB(t)
108+
prepareCodexQuotaMigrationFixtureWithPreexistingIdentityIndex(t, gormDB)
109+
110+
migration := findMigrationByVersion(t, 8)
111+
if err := migration.Up(gormDB); err != nil {
112+
t.Fatalf("run migration v8 up with preexisting identity index: %v", err)
113+
}
114+
115+
var count int64
116+
if err := gormDB.Raw(`
117+
SELECT COUNT(*)
118+
FROM codex_quotas
119+
WHERE tenant_id = 1
120+
AND identity_key = 'account:e94ce011-80f4-490b-b285-d3109db72b0e'
121+
AND deleted_at = 0
122+
`).Scan(&count).Error; err != nil {
123+
t.Fatalf("count migrated rows: %v", err)
124+
}
125+
if count != 1 {
126+
t.Fatalf("expected preexisting index collision rows to collapse to 1, got %d", count)
127+
}
128+
assertIndexExists(t, gormDB, "idx_codex_quotas_tenant_identity", true)
129+
assertIndexExists(t, gormDB, "idx_codex_quotas_email", false)
130+
assertIndexMissing(t, gormDB, "idx_codex_quotas_tenant_email")
131+
}
132+
106133
func TestCodexQuotaIdentityMigrationV8DownReturnsIrreversibleError(t *testing.T) {
107134
gormDB := openRawSQLiteDB(t)
108135
prepareCodexQuotaMigrationFixture(t, gormDB)
@@ -293,6 +320,52 @@ func prepareCodexQuotaMigrationFixtureWithDeletedHistory(t *testing.T, gormDB *g
293320
}
294321
}
295322

323+
func prepareCodexQuotaMigrationFixtureWithPreexistingIdentityIndex(t *testing.T, gormDB *gorm.DB) {
324+
t.Helper()
325+
if err := gormDB.Exec(`DROP TABLE IF EXISTS codex_quotas`).Error; err != nil {
326+
t.Fatalf("drop table: %v", err)
327+
}
328+
if err := gormDB.Exec(`
329+
CREATE TABLE codex_quotas (
330+
id TEXT PRIMARY KEY,
331+
created_at INTEGER DEFAULT 0,
332+
updated_at INTEGER DEFAULT 0,
333+
deleted_at INTEGER DEFAULT 0,
334+
tenant_id INTEGER NOT NULL,
335+
identity_key TEXT,
336+
email TEXT,
337+
account_id TEXT,
338+
plan_type TEXT,
339+
is_forbidden INTEGER DEFAULT 0,
340+
primary_window TEXT,
341+
secondary_window TEXT,
342+
code_review_window TEXT
343+
)
344+
`).Error; err != nil {
345+
t.Fatalf("create table: %v", err)
346+
}
347+
if err := gormDB.Exec(`CREATE UNIQUE INDEX idx_codex_quotas_tenant_identity ON codex_quotas(tenant_id, identity_key)`).Error; err != nil {
348+
t.Fatalf("create legacy identity index: %v", err)
349+
}
350+
if err := gormDB.Exec(`CREATE UNIQUE INDEX idx_codex_quotas_email ON codex_quotas(email)`).Error; err != nil {
351+
t.Fatalf("create broken email index: %v", err)
352+
}
353+
if err := gormDB.Exec(`CREATE UNIQUE INDEX idx_codex_quotas_tenant_email ON codex_quotas(tenant_id, email)`).Error; err != nil {
354+
t.Fatalf("create legacy tenant email index: %v", err)
355+
}
356+
inserts := []string{
357+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-8', 1, NULL, 'cnc6n2io9xvfev2mtm5t6hu8@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456445476)`,
358+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-11', 1, NULL, 'fcew8ua8r6u6zekrwqfi60nl@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456444987)`,
359+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-7', 1, NULL, 'puckxqnzu6ktt7k4bcevlw06@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456443639)`,
360+
`INSERT INTO codex_quotas (id, tenant_id, identity_key, email, account_id, updated_at) VALUES ('row-9', 1, NULL, 'n7e87dj5hxv2c2m4u0e6l5zo@example.com', 'e94ce011-80f4-490b-b285-d3109db72b0e', 1773456443366)`,
361+
}
362+
for _, sql := range inserts {
363+
if err := gormDB.Exec(sql).Error; err != nil {
364+
t.Fatalf("insert fixture: %v", err)
365+
}
366+
}
367+
}
368+
296369
func assertCodexQuotaFixtureCounts(t *testing.T, gormDB *gorm.DB) {
297370
t.Helper()
298371
var duplicateCount int64

0 commit comments

Comments
 (0)