diff --git a/sdk/typescript/_bundled_plugin/scripts/workbench_db.py b/sdk/typescript/_bundled_plugin/scripts/workbench_db.py index c4f26854..1593dc6a 100644 --- a/sdk/typescript/_bundled_plugin/scripts/workbench_db.py +++ b/sdk/typescript/_bundled_plugin/scripts/workbench_db.py @@ -91,9 +91,12 @@ ) from workbench_schema import ( MIGRATIONS, - normalize_pre_release_migrations, - repair_deep_scan_migration, - sql_statements, +) +from workbench_schema import ( + apply_migrations as apply_schema_migrations, +) +from workbench_schema import ( + sql_statements as sql_statements, ) from workbench_source_excerpt import finding_source_excerpt from workbench_target import ( @@ -277,39 +280,7 @@ def disable_setup_ui(connection: sqlite3.Connection, args: argparse.Namespace) - def apply_migrations(connection: sqlite3.Connection) -> None: - connection.commit() - connection.execute("BEGIN IMMEDIATE") - try: - connection.execute( - """ - CREATE TABLE IF NOT EXISTS schema_migrations ( - version INTEGER PRIMARY KEY, - name TEXT NOT NULL, - applied_at TEXT NOT NULL - ) - """ - ) - normalize_pre_release_migrations(connection, now()) - applied = { - row["version"] for row in connection.execute("SELECT version FROM schema_migrations") - } - for version, name, sql in MIGRATIONS: - if version in applied: - if version == 11: - repair_deep_scan_migration(connection) - continue - for statement in sql_statements(sql): - connection.execute(statement) - connection.execute( - "INSERT INTO schema_migrations (version, name, applied_at) VALUES (?, ?, ?)", - (version, name, now()), - ) - if 16 not in applied: - backfill_security_targets(connection) - connection.commit() - except BaseException: - connection.rollback() - raise + apply_schema_migrations(connection, MIGRATIONS, now, backfill_security_targets) def require_target(value: str) -> Path: diff --git a/sdk/typescript/_bundled_plugin/scripts/workbench_schema.py b/sdk/typescript/_bundled_plugin/scripts/workbench_schema.py index 75b81a26..8bed89cf 100644 --- a/sdk/typescript/_bundled_plugin/scripts/workbench_schema.py +++ b/sdk/typescript/_bundled_plugin/scripts/workbench_schema.py @@ -2,6 +2,7 @@ import argparse import sqlite3 +from collections.abc import Callable MIGRATIONS = ( ( @@ -637,110 +638,77 @@ ) -def normalize_pre_release_migrations( - connection: sqlite3.Connection, timestamp: str +def apply_migrations( + connection: sqlite3.Connection, + migrations: tuple[tuple[int, str, str], ...], + now: Callable[[], str], + backfill_security_targets: Callable[[sqlite3.Connection], None], ) -> None: - execution_migrations = { - row["version"]: row["name"] - for row in connection.execute( - "SELECT version, name FROM schema_migrations WHERE version IN (11, 12)" - ) - } - supported_execution_migrations = { - 11: {"deep scan orchestration state", "scan execution profiles"}, - 12: { - "scan continuation threads", - "dynamic scan execution profiles", - "phase-specific scan progress", - }, - } - if any( - name not in supported_execution_migrations[version] - for version, name in execution_migrations.items() - ): - raise SystemExit( - "The Codex Security database has an unsupported execution-profile migration history." - ) - - phase_progress_migration = connection.execute( - "SELECT name FROM schema_migrations WHERE version = 12" - ).fetchone() - if ( - phase_progress_migration is not None - and phase_progress_migration["name"] == "phase-specific scan progress" - ): - target_migration = connection.execute( - "SELECT name FROM schema_migrations WHERE version = 20" - ).fetchone() - if target_migration is not None: - raise SystemExit( - "The Codex Security database has an unsupported pre-release migration history." - ) + connection.commit() + connection.execute("BEGIN IMMEDIATE") + try: connection.execute( - "UPDATE schema_migrations SET version = 20 WHERE version = 12 AND name = ?", - ("phase-specific scan progress",), - ) - - normalize_pre_release_execution_profile_migrations(connection, timestamp) - - preflight_progress_migration = connection.execute( - "SELECT name FROM schema_migrations WHERE version = 13" - ).fetchone() - if ( - preflight_progress_migration is not None - and preflight_progress_migration["name"] == "current scan preflight state" - ): - target_migration = connection.execute( - "SELECT name FROM schema_migrations WHERE version = 21" - ).fetchone() - if target_migration is not None: - raise SystemExit( - "The Codex Security database has an unsupported pre-release migration history." + """ + CREATE TABLE IF NOT EXISTS schema_migrations ( + version INTEGER PRIMARY KEY, + name TEXT NOT NULL, + applied_at TEXT NOT NULL ) - connection.execute( - "UPDATE schema_migrations SET version = 21 WHERE version = 13 AND name = ?", - ("current scan preflight state",), - ) - - migration = connection.execute( - "SELECT name FROM schema_migrations WHERE version = 2" - ).fetchone() - if migration is None or migration["name"] != "finding management schema": - return - - legacy_versions = { - row["version"]: row["name"] - for row in connection.execute( - "SELECT version, name FROM schema_migrations WHERE version BETWEEN 2 AND 5" + """ ) - } - expected = { - 2: "finding management schema", - 3: "scan handoff delivery claims", - 4: "finding remediation action claims", - 5: "scan target snapshot digests", - } - for version, name in legacy_versions.items(): - if expected.get(version) != name: - raise SystemExit( - "The Codex Security database has an unsupported pre-release migration history." + normalize_pre_release_migrations(connection, now()) + applied = { + row["version"] for row in connection.execute("SELECT version FROM schema_migrations") + } + should_backfill_targets = False + for version, name, sql in migrations: + if version in applied: + if version == 2: + add_column_if_missing( + connection, "workspaces", "capability_preflight_json", "TEXT" + ) + elif version == 6: + repair_thread_scoped_workspaces_migration(connection) + elif version == 11: + repair_deep_scan_migration(connection) + elif version == 12: + add_column_if_missing(connection, "scans", "continuation_thread_id", "TEXT") + elif version == 13: + add_column_if_missing( + connection, + "scan_progress", + "scope_file_count", + "INTEGER CHECK (scope_file_count >= 0)", + ) + elif version == 16: + should_backfill_targets = repair_stable_targets_migration(connection) + elif version == 26: + add_column_if_missing( + connection, + "scans", + "completion_warnings_json", + "TEXT NOT NULL DEFAULT '[]'", + ) + continue + if version == 6: + repair_thread_scoped_workspaces_migration(connection) + elif version == 16: + should_backfill_targets = repair_stable_targets_migration(connection) + else: + for statement in sql_statements(sql): + connection.execute(statement) + connection.execute( + "INSERT INTO schema_migrations (version, name, applied_at) VALUES (?, ?, ?)", + (version, name, now()), ) - - connection.execute( - "DELETE FROM schema_migrations WHERE version = 5 AND name = ?", - (expected[5],), - ) - for old_version, new_version in ((4, 5), (3, 4), (2, 3)): - connection.execute( - "UPDATE schema_migrations SET version = ? WHERE version = ? AND name = ?", - (new_version, old_version, expected[old_version]), - ) - add_column_if_missing(connection, "workspaces", "capability_preflight_json", "TEXT") - add_column_if_missing(connection, "scans", "target_snapshot_digest", "TEXT") - connection.execute( - "INSERT INTO schema_migrations (version, name, applied_at) VALUES (?, ?, ?)", - (2, "persist capability preflight summaries", timestamp), - ) + if 27 in applied: + repair_deep_scan_failure_counter_migration(connection) + if should_backfill_targets: + backfill_security_targets(connection) + connection.commit() + except BaseException: + connection.rollback() + raise def normalize_pre_release_execution_profile_migrations( @@ -762,6 +730,12 @@ def normalize_pre_release_execution_profile_migrations( } model_migration_name = "persist scan model settings" warnings_migration_name = "persist scan completion warnings" + if execution_migrations.get(25) == "dynamic scan execution profiles": + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 25 AND name = ?", + (model_migration_name, "dynamic scan execution profiles"), + ) + execution_migrations[25] = model_migration_name has_legacy_profile_history = execution_migrations.get(11) == legacy_names[11] has_public_warnings_history = ( execution_migrations.get(25) == warnings_migration_name @@ -872,6 +846,233 @@ def normalize_pre_release_execution_profile_migrations( ) +def normalize_pre_release_migrations(connection: sqlite3.Connection, timestamp: str) -> None: + completion_warning_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 25" + ).fetchone() + if ( + completion_warning_migration is not None + and completion_warning_migration["name"] == "persist scan completion warnings" + ): + if ( + connection.execute("SELECT 1 FROM schema_migrations WHERE version = 26").fetchone() + is not None + ): + raise SystemExit( + "The Codex Security database has an unsupported pre-release migration history." + ) + connection.execute( + "UPDATE schema_migrations SET version = 26 WHERE version = 25 AND name = ?", + ("persist scan completion warnings",), + ) + + phase_progress_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 12" + ).fetchone() + if ( + phase_progress_migration is not None + and phase_progress_migration["name"] == "phase-specific scan progress" + ): + target_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 20" + ).fetchone() + if target_migration is not None: + raise SystemExit( + "The Codex Security database has an unsupported pre-release migration history." + ) + connection.execute( + "UPDATE schema_migrations SET version = 20 WHERE version = 12 AND name = ?", + ("phase-specific scan progress",), + ) + + normalize_pre_release_execution_profile_migrations(connection, timestamp) + + preflight_progress_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 13" + ).fetchone() + if ( + preflight_progress_migration is not None + and preflight_progress_migration["name"] == "current scan preflight state" + ): + target_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 21" + ).fetchone() + if target_migration is not None: + raise SystemExit( + "The Codex Security database has an unsupported pre-release migration history." + ) + connection.execute( + "UPDATE schema_migrations SET version = 21 WHERE version = 13 AND name = ?", + ("current scan preflight state",), + ) + + delivered_claim_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 18" + ).fetchone() + if ( + delivered_claim_migration is not None + and delivered_claim_migration["name"] == "scan target summaries" + ): + connection.execute( + "UPDATE scans SET handoff_claimed_at = NULL, handoff_claim_token = NULL " + "WHERE handoff_status = 'delivered'" + ) + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 18 AND name = ?", + ("clear legacy delivered handoff claims", "scan target summaries"), + ) + + setup_preferences_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 19" + ).fetchone() + legacy_setup_preferences_migrations = { + "structured scan guidance context", + "idempotent scan lifecycle requests", + } + if ( + setup_preferences_migration is not None + and setup_preferences_migration["name"] in legacy_setup_preferences_migrations + ): + migration_sql = next(sql for version, _, sql in MIGRATIONS if version == 19) + for statement in sql_statements(migration_sql): + connection.execute(statement.replace("CREATE TABLE ", "CREATE TABLE IF NOT EXISTS ", 1)) + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 19 AND name = ?", + ("persist setup workspace preference", setup_preferences_migration["name"]), + ) + + phase_progress_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 20" + ).fetchone() + legacy_phase_progress_migrations = { + "retain superseded scan lifecycle requests", + "threat model publication receipts", + } + if ( + phase_progress_migration is not None + and phase_progress_migration["name"] in legacy_phase_progress_migrations + ): + add_column_if_missing( + connection, + "scan_progress", + "phase_items_total", + "INTEGER NOT NULL DEFAULT 0 CHECK (phase_items_total >= 0)", + ) + add_column_if_missing( + connection, + "scan_progress", + "phase_items_completed", + "INTEGER NOT NULL DEFAULT 0 " + "CHECK (phase_items_completed >= 0 AND phase_items_completed <= phase_items_total)", + ) + add_column_if_missing( + connection, + "scan_progress", + "phase_progress_unit", + "TEXT CHECK (phase_progress_unit IS NULL OR phase_progress_unit IN (" + "'checks', 'threat_surfaces', 'review_receipts', 'candidate_findings', " + "'validated_findings', 'report_artifacts'))", + ) + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 20 AND name = ?", + ("phase-specific scan progress", phase_progress_migration["name"]), + ) + + preflight_progress_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 21" + ).fetchone() + legacy_preflight_progress_migrations = { + "scan progress projection and activity", + "deep coordinator manifest receipts", + } + if ( + preflight_progress_migration is not None + and preflight_progress_migration["name"] in legacy_preflight_progress_migrations + ): + add_column_if_missing( + connection, + "scan_progress", + "preflight_issues_json", + "TEXT NOT NULL DEFAULT '[]'", + ) + add_column_if_missing( + connection, + "scan_progress", + "preflight_checks_total", + "INTEGER NOT NULL DEFAULT 0 CHECK (preflight_checks_total >= 0)", + ) + add_column_if_missing( + connection, + "scan_progress", + "preflight_checks_completed", + "INTEGER NOT NULL DEFAULT 0 CHECK (preflight_checks_completed >= 0 " + "AND preflight_checks_completed <= preflight_checks_total)", + ) + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 21 AND name = ?", + ("current scan preflight state", preflight_progress_migration["name"]), + ) + + scan_recipe_migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 22" + ).fetchone() + if ( + scan_recipe_migration is not None + and scan_recipe_migration["name"] == "dynamic scan execution profiles" + ): + add_column_if_missing(connection, "scans", "recipe_json", "TEXT") + add_column_if_missing( + connection, + "scans", + "parent_scan_id", + "TEXT REFERENCES scans(id) ON DELETE SET NULL", + ) + connection.execute( + "UPDATE schema_migrations SET name = ? WHERE version = 22 AND name = ?", + ("replayable scan launch recipes", "dynamic scan execution profiles"), + ) + + migration = connection.execute( + "SELECT name FROM schema_migrations WHERE version = 2" + ).fetchone() + if migration is None or migration["name"] != "finding management schema": + return + + legacy_versions = { + row["version"]: row["name"] + for row in connection.execute( + "SELECT version, name FROM schema_migrations WHERE version BETWEEN 2 AND 5" + ) + } + expected = { + 2: "finding management schema", + 3: "scan handoff delivery claims", + 4: "finding remediation action claims", + 5: "scan target snapshot digests", + } + for version, name in legacy_versions.items(): + if expected.get(version) != name: + raise SystemExit( + "The Codex Security database has an unsupported pre-release migration history." + ) + + connection.execute( + "DELETE FROM schema_migrations WHERE version = 5 AND name = ?", + (expected[5],), + ) + for old_version, new_version in ((4, 5), (3, 4), (2, 3)): + connection.execute( + "UPDATE schema_migrations SET version = ? WHERE version = ? AND name = ?", + (new_version, old_version, expected[old_version]), + ) + add_column_if_missing(connection, "workspaces", "capability_preflight_json", "TEXT") + add_column_if_missing(connection, "scans", "target_snapshot_digest", "TEXT") + connection.execute( + "INSERT INTO schema_migrations (version, name, applied_at) VALUES (?, ?, ?)", + (2, "persist capability preflight summaries", timestamp), + ) + + def repair_deep_scan_migration(connection: sqlite3.Connection) -> None: scan_columns = {row["name"] for row in connection.execute("PRAGMA table_info(scans)")} owner_column_missing = "deep_scan_owner_thread_id" not in scan_columns @@ -906,6 +1107,91 @@ def repair_deep_scan_migration(connection: sqlite3.Connection) -> None: statement = statement.replace(prefix, f"{prefix}IF NOT EXISTS ", 1) break connection.execute(statement) + if statement.startswith("UPDATE scans") and "continuation_thread_id" in scan_columns: + connection.execute( + "UPDATE scans SET deep_scan_owner_thread_id = continuation_thread_id " + "WHERE mode = 'deep' AND status = 'running' " + "AND continuation_thread_id IS NOT NULL" + ) + + +def repair_deep_scan_failure_counter_migration(connection: sqlite3.Connection) -> None: + columns = {row["name"] for row in connection.execute("PRAGMA table_info(deep_scan_runs)")} + threshold_missing = "stop_after_consecutive_errors" not in columns + add_column_if_missing( + connection, + "deep_scan_runs", + "stop_after_consecutive_errors", + "INTEGER NOT NULL DEFAULT 1 CHECK (stop_after_consecutive_errors >= 1)", + ) + if threshold_missing: + connection.execute( + "UPDATE deep_scan_runs SET stop_after_consecutive_errors = stop_after_no_new" + ) + add_column_if_missing( + connection, + "deep_scan_runs", + "consecutive_errors", + "INTEGER NOT NULL DEFAULT 0 CHECK (consecutive_errors >= 0)", + ) + + +def repair_thread_scoped_workspaces_migration(connection: sqlite3.Connection) -> None: + add_column_if_missing(connection, "workspaces", "thread_id", "TEXT") + connection.execute( + "CREATE INDEX IF NOT EXISTS workspaces_by_thread_and_updated_at " + "ON workspaces(thread_id, updated_at DESC)" + ) + + +def repair_stable_targets_migration(connection: sqlite3.Connection) -> bool: + workspace_columns = {row["name"] for row in connection.execute("PRAGMA table_info(workspaces)")} + scan_columns = {row["name"] for row in connection.execute("PRAGMA table_info(scans)")} + existing_objects = { + row["name"] + for row in connection.execute( + "SELECT name FROM sqlite_master WHERE name IN ('security_targets', 'scans_by_target')" + ) + } + if ( + "target_id" in workspace_columns + and "target_id" in scan_columns + and existing_objects == {"security_targets", "scans_by_target"} + ): + return False + + migration_sql = next(sql for version, _, sql in MIGRATIONS if version == 16) + for statement in sql_statements(migration_sql): + if statement.startswith("ALTER TABLE workspaces"): + add_column_if_missing( + connection, + "workspaces", + "target_id", + "TEXT REFERENCES security_targets(id)", + ) + continue + if statement.startswith("ALTER TABLE scans"): + add_column_if_missing( + connection, + "scans", + "target_id", + "TEXT REFERENCES security_targets(id)", + ) + continue + statement = statement.replace("CREATE TABLE ", "CREATE TABLE IF NOT EXISTS ", 1) + statement = statement.replace("CREATE INDEX ", "CREATE INDEX IF NOT EXISTS ", 1) + connection.execute(statement) + connection.execute( + """ + UPDATE scans + SET target_id = NULL + WHERE target_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 FROM security_targets WHERE security_targets.id = scans.target_id + ) + """ + ) + return True def add_column_if_missing(