fix: keep large database startup alive (#3624)

This commit is contained in:
Carlos Florêncio
2026-09-12 16:29:38 +01:00
committed by GitHub
parent 3d3209aa1a
commit b6571a9854
58 changed files with 2781 additions and 439 deletions
@@ -2,6 +2,7 @@
package sqlite
import (
"context"
"database/sql"
"errors"
"fmt"
@@ -391,7 +392,8 @@ func (r *Repository) runMigrations() error {
return nil
}
func (r *Repository) backfillTaskSessionQueueIncarnations() error {
rows, err := r.db.Queryx(`SELECT id FROM task_sessions WHERE queue_incarnation_id = ''`)
ctx := r.migrationContext()
rows, err := r.db.QueryxContext(ctx, `SELECT id FROM task_sessions WHERE queue_incarnation_id = ''`)
if err != nil {
return fmt.Errorf("list sessions missing queue incarnation: %w", err)
}
@@ -411,7 +413,7 @@ func (r *Repository) backfillTaskSessionQueueIncarnations() error {
return fmt.Errorf("iterate sessions missing queue incarnation: %w", err)
}
for _, sessionID := range sessionIDs {
if _, err := r.db.Exec(r.db.Rebind(`
if _, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_sessions
SET queue_incarnation_id = ?
WHERE id = ? AND queue_incarnation_id = ''
@@ -440,7 +442,8 @@ func (r *Repository) backfillPromptSeq() error {
AND (%s < %s OR (%s = %s AND u.id <= task_session_messages.id))
)
WHERE author_type = 'user' AND prompt_seq = 0`, nmU, nmM, nmU, nmM)
if _, err := r.db.Exec(update); err != nil {
ctx := r.migrationContext()
if _, err := r.db.ExecContext(ctx, update); err != nil {
return fmt.Errorf("backfill prompt_seq: %w", err)
}
seed := `
@@ -449,7 +452,7 @@ func (r *Repository) backfillPromptSeq() error {
WHERE author_type = 'user' AND prompt_seq > 0
GROUP BY task_session_id
ON CONFLICT(task_session_id) DO NOTHING`
if _, err := r.db.Exec(seed); err != nil {
if _, err := r.db.ExecContext(ctx, seed); err != nil {
return fmt.Errorf("seed prompt sequence counters: %w", err)
}
return nil
@@ -602,14 +605,14 @@ func jsonBoolToInt(postgres bool, column, parent, key string) string {
// the initial implementation, reconciles legacy bootstrap duplicates, then
// enforces uniqueness only for the two hidden workflows created by this flow.
func (r *Repository) ensureImproveKandevWorkflowTemplateUniqueness() error {
tx, err := r.db.Beginx()
tx, err := r.db.BeginTxx(r.migrationContext(), nil)
if err != nil {
return fmt.Errorf("begin improve kandev workflow migration: %w", err)
}
// Keep this template list synchronized with every hidden Improve Kandev
// workflow bootstrapped by internal/improvekandev. Add new IDs to both this
// reconciliation query and the partial-index predicate below.
if _, err := tx.Exec(`
if _, err := tx.ExecContext(r.migrationContext(), `
UPDATE workflows
SET workflow_template_id = ''
WHERE id IN (
@@ -631,11 +634,11 @@ func (r *Repository) ensureImproveKandevWorkflowTemplateUniqueness() error {
_ = tx.Rollback()
return fmt.Errorf("reconcile improve kandev workflow duplicates: %w", err)
}
if _, err := tx.Exec(`DROP INDEX IF EXISTS uniq_workflows_workspace_template_hidden`); err != nil {
if _, err := tx.ExecContext(r.migrationContext(), `DROP INDEX IF EXISTS uniq_workflows_workspace_template_hidden`); err != nil {
_ = tx.Rollback()
return fmt.Errorf("drop broad workflow template index: %w", err)
}
if _, err := tx.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS uniq_improve_kandev_workflows
if _, err := tx.ExecContext(r.migrationContext(), `CREATE UNIQUE INDEX IF NOT EXISTS uniq_improve_kandev_workflows
ON workflows(workspace_id, workflow_template_id, hidden)
WHERE workflow_template_id IN ('improve-kandev', 'report-kandev-issue')`); err != nil {
_ = tx.Rollback()
@@ -651,7 +654,7 @@ func (r *Repository) ensureImproveKandevWorkflowTemplateUniqueness() error {
// folder attachments existed. CREATE TABLE/INDEX IF NOT EXISTS is replay-safe
// on SQLite and Postgres.
func (r *Repository) ensureTaskWorkspaceFoldersSchema() error {
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
CREATE TABLE IF NOT EXISTS task_workspace_folders (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL,
@@ -676,7 +679,7 @@ func (r *Repository) ensureTaskWorkspaceFoldersSchema() error {
// existed. It replays the same DDL as the schema-init step; CREATE TABLE/INDEX
// IF NOT EXISTS is replay-safe on SQLite and Postgres.
func (r *Repository) ensureRepositorySetsSchema() error {
if _, err := r.db.Exec(repositorySetsSchemaDDL); err != nil {
if _, err := r.db.ExecContext(r.migrationContext(), repositorySetsSchemaDDL); err != nil {
return fmt.Errorf("create repository sets schema: %w", err)
}
return nil
@@ -685,7 +688,7 @@ func (r *Repository) ensureRepositorySetsSchema() error {
// ensureRepositoryBranchPoliciesSchema replays the policy DDL for databases
// created before repository branch policies existed.
func (r *Repository) ensureRepositoryBranchPoliciesSchema() error {
if _, err := r.db.Exec(repositoryBranchPoliciesSchemaDDL); err != nil {
if _, err := r.db.ExecContext(r.migrationContext(), repositoryBranchPoliciesSchemaDDL); err != nil {
return fmt.Errorf("create repository branch policies schema: %w", err)
}
return nil
@@ -705,7 +708,7 @@ func (r *Repository) recreateTable(tableName, triggerPhrase string, statements [
}
var tableSql string
err := r.db.QueryRow(`SELECT sql FROM sqlite_master WHERE type='table' AND name=?`, tableName).Scan(&tableSql)
err := r.db.QueryRowContext(r.migrationContext(), `SELECT sql FROM sqlite_master WHERE type='table' AND name=?`, tableName).Scan(&tableSql)
if errors.Is(err, sql.ErrNoRows) {
return false, nil // Table doesn't exist yet; migration not applicable
}
@@ -716,19 +719,21 @@ func (r *Repository) recreateTable(tableName, triggerPhrase string, statements [
return false, nil // Trigger phrase absent; migration already applied or not needed
}
if _, err := r.db.Exec(`PRAGMA foreign_keys=OFF`); err != nil {
if _, err := r.db.ExecContext(r.migrationContext(), `PRAGMA foreign_keys=OFF`); err != nil {
return false, fmt.Errorf("disable foreign keys: %w", err)
}
defer func() { _, _ = r.db.Exec(`PRAGMA foreign_keys=ON`) }()
// This restoration is cleanup, so it must not inherit a canceled startup
// context and leave the shared SQLite writer with foreign-key checks off.
defer func() { _, _ = r.db.ExecContext(context.Background(), `PRAGMA foreign_keys=ON`) }()
tx, err := r.db.Beginx()
tx, err := r.db.BeginTxx(r.migrationContext(), nil)
if err != nil {
return false, fmt.Errorf("begin migration transaction: %w", err)
}
defer func() { _ = tx.Rollback() }()
for _, stmt := range statements {
if _, err := tx.Exec(stmt); err != nil {
if _, err := tx.ExecContext(r.migrationContext(), stmt); err != nil {
return false, fmt.Errorf("migration %s failed: %w", tableName, err)
}
}
@@ -806,7 +811,7 @@ func (r *Repository) backfillExecutorsRunningFromTaskSessions() error {
// Check whether task_sessions still has the column. If migration already ran,
// the column is gone and there's nothing to backfill.
var tableSql string
if err := r.db.QueryRow(`SELECT sql FROM sqlite_master WHERE type='table' AND name='task_sessions'`).Scan(&tableSql); err != nil {
if err := r.db.QueryRowContext(r.migrationContext(), `SELECT sql FROM sqlite_master WHERE type='table' AND name='task_sessions'`).Scan(&tableSql); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
@@ -820,7 +825,7 @@ func (r *Repository) backfillExecutorsRunningFromTaskSessions() error {
// SELECT … LEFT JOIN to find sessions with execution data but no executors_running row.
// Insert with the minimum field set; runtime/status are best-effort defaults
// (subsequent Launch / Resume will overwrite via the lifecycle manager's persistence).
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
INSERT INTO executors_running (
id, session_id, task_id, executor_id, runtime, status, resumable,
resume_token, last_message_uuid, agent_execution_id, container_id,
@@ -887,7 +892,7 @@ const commitCaptureActivatedAtMetaKey = "commit_capture_activated_at"
// exist, so a failure here must abort boot (propagated below) rather than
// leave every future commit insert failing silently forever.
func (r *Repository) migrateSessionCommitsDedupeAndActivation() error {
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
DELETE FROM task_session_commits
WHERE id NOT IN (
SELECT id FROM (
@@ -903,7 +908,7 @@ func (r *Repository) migrateSessionCommitsDedupeAndActivation() error {
`); err != nil {
return fmt.Errorf("dedupe task_session_commits: %w", err)
}
if _, err := r.db.Exec(
if _, err := r.db.ExecContext(r.migrationContext(),
`CREATE UNIQUE INDEX IF NOT EXISTS uniq_session_commits_session_sha ON task_session_commits(session_id, commit_sha)`,
); err != nil {
return fmt.Errorf("create uniq_session_commits_session_sha: %w", err)
@@ -916,7 +921,7 @@ func (r *Repository) migrateSessionCommitsDedupeAndActivation() error {
// production (persistence.Provide creates it before opening any
// repository), but repo-level tests build a bare DB via NewWithDB where
// it does not, so recreate it defensively.
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
CREATE TABLE IF NOT EXISTS kandev_meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL DEFAULT ''
@@ -924,7 +929,7 @@ func (r *Repository) migrateSessionCommitsDedupeAndActivation() error {
return fmt.Errorf("ensure kandev_meta: %w", err)
}
now := time.Now().UTC().Format(time.RFC3339Nano)
if _, err := r.db.Exec(r.db.Rebind(`
if _, err := r.db.ExecContext(r.migrationContext(), r.db.Rebind(`
INSERT INTO kandev_meta (key, value) VALUES (?, ?)
ON CONFLICT(key) DO NOTHING
`), commitCaptureActivatedAtMetaKey, now); err != nil {
@@ -1081,10 +1086,10 @@ func (r *Repository) migrateTaskEnvironmentReposAllowMultiBranch() error {
}
func (r *Repository) migrateTaskEnvironmentReposAllowMultiBranchPostgres() error {
if _, err := r.db.Exec(`ALTER TABLE task_environment_repos ADD COLUMN IF NOT EXISTS branch_slug TEXT NOT NULL DEFAULT ''`); err != nil {
if _, err := r.db.ExecContext(r.migrationContext(), `ALTER TABLE task_environment_repos ADD COLUMN IF NOT EXISTS branch_slug TEXT NOT NULL DEFAULT ''`); err != nil {
return fmt.Errorf("add task_environment_repos.branch_slug: %w", err)
}
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
DO $$
DECLARE
old_constraint_name text;
@@ -1254,13 +1259,13 @@ func (r *Repository) migrateSessionsRemoveWorkflowStepID() error {
// Runs before ensureTaskEnvironmentTaskUniqueIndex so the unique constraint
// can be added cleanly. Idempotent — a no-op once the data is healed.
func (r *Repository) healDuplicateTaskEnvironments() error {
tx, err := r.db.Begin()
tx, err := r.db.BeginTx(r.migrationContext(), nil)
if err != nil {
return fmt.Errorf("heal duplicate envs: begin tx: %w", err)
}
defer tx.Rollback() //nolint:errcheck
rows, err := tx.Query(`
rows, err := tx.QueryContext(r.migrationContext(), `
SELECT task_id
FROM task_environments
GROUP BY task_id
@@ -1285,7 +1290,7 @@ func (r *Repository) healDuplicateTaskEnvironments() error {
_ = rows.Close()
for _, taskID := range taskIDs {
if err := healDuplicateTaskEnvForTask(tx, taskID); err != nil {
if err := healDuplicateTaskEnvForTask(r.migrationContext(), tx, taskID); err != nil {
return err
}
}
@@ -1294,9 +1299,9 @@ func (r *Repository) healDuplicateTaskEnvironments() error {
// healDuplicateTaskEnvForTask keeps the most recently updated env for a task,
// re-points sessions on the loser rows to the winner, then deletes losers.
func healDuplicateTaskEnvForTask(tx *sql.Tx, taskID string) error {
func healDuplicateTaskEnvForTask(ctx context.Context, tx *sql.Tx, taskID string) error {
var winnerID string
if err := tx.QueryRow(`
if err := tx.QueryRowContext(ctx, `
SELECT id FROM task_environments
WHERE task_id = ?
ORDER BY updated_at DESC, created_at DESC
@@ -1305,7 +1310,7 @@ func healDuplicateTaskEnvForTask(tx *sql.Tx, taskID string) error {
return fmt.Errorf("heal duplicate envs: find winner for task %s: %w", taskID, err)
}
if _, err := tx.Exec(`
if _, err := tx.ExecContext(ctx, `
UPDATE task_sessions
SET task_environment_id = ?
WHERE task_id = ?
@@ -1314,7 +1319,7 @@ func healDuplicateTaskEnvForTask(tx *sql.Tx, taskID string) error {
return fmt.Errorf("heal duplicate envs: relink sessions for task %s: %w", taskID, err)
}
if _, err := tx.Exec(`
if _, err := tx.ExecContext(ctx, `
DELETE FROM task_environments
WHERE task_id = ?
AND id != ?
@@ -1329,7 +1334,7 @@ func healDuplicateTaskEnvForTask(tx *sql.Tx, taskID string) error {
// instead of silently producing two rows for the same task. Must run AFTER
// healDuplicateTaskEnvironments, which collapses any pre-existing duplicates.
func (r *Repository) ensureTaskEnvironmentTaskUniqueIndex() error {
_, err := r.db.Exec(`
_, err := r.db.ExecContext(r.migrationContext(), `
CREATE UNIQUE INDEX IF NOT EXISTS uniq_task_environments_task_id
ON task_environments(task_id)
`)
@@ -1351,7 +1356,7 @@ func (r *Repository) healSessionTaskEnvironmentIDs() error {
// ensureTaskEnvironmentTaskUniqueIndex guarantees ≤1 row per task at
// runtime, but the SQL reads as non-deterministic in isolation. Belt
// and suspenders.
if _, err := r.db.Exec(`
if _, err := r.db.ExecContext(r.migrationContext(), `
UPDATE task_sessions
SET task_environment_id = (
SELECT te.id FROM task_environments te WHERE te.task_id = task_sessions.task_id LIMIT 1