diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index e9acd448de..e758d8fca2 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -445,24 +445,37 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: if int(row[5] or 0) > 0 ] if claim_pk != ["thread_id"]: - conn.execute("ALTER TABLE research_thread_claims RENAME TO research_thread_claims_legacy") - conn.execute( - """ - CREATE TABLE research_thread_claims ( - owner_subject TEXT NOT NULL, - thread_id TEXT NOT NULL PRIMARY KEY REFERENCES chat_threads(id) ON DELETE CASCADE, - created_at INTEGER NOT NULL - ) WITHOUT ROWID - """ - ) - conn.execute( - """INSERT OR IGNORE INTO research_thread_claims - (owner_subject, thread_id, created_at) - SELECT owner_subject, thread_id, created_at - FROM research_thread_claims_legacy - ORDER BY created_at, owner_subject""" - ) - conn.execute("DROP TABLE research_thread_claims_legacy") + # Rebuild the claims table (legacy owner_subject+thread_id PK -> thread_id PK) + # atomically. Without an explicit transaction the RENAME/CREATE/INSERT/DROP run + # in autocommit, so an interruption after CREATE left the new table empty and the + # rows orphaned in _legacy, and the migration never re-triggered. + conn.commit() + conn.execute("BEGIN IMMEDIATE") + try: + conn.execute( + "ALTER TABLE research_thread_claims RENAME TO research_thread_claims_legacy" + ) + conn.execute( + """ + CREATE TABLE research_thread_claims ( + owner_subject TEXT NOT NULL, + thread_id TEXT NOT NULL PRIMARY KEY REFERENCES chat_threads(id) ON DELETE CASCADE, + created_at INTEGER NOT NULL + ) WITHOUT ROWID + """ + ) + conn.execute( + """INSERT OR IGNORE INTO research_thread_claims + (owner_subject, thread_id, created_at) + SELECT owner_subject, thread_id, created_at + FROM research_thread_claims_legacy + ORDER BY created_at, owner_subject""" + ) + conn.execute("DROP TABLE research_thread_claims_legacy") + conn.commit() + except Exception: + conn.rollback() + raise conn.execute( """INSERT OR IGNORE INTO research_thread_claims (owner_subject, thread_id, created_at) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index 3baf17e550..a87d785355 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -423,6 +423,81 @@ def test_owner_scoped_claim_schema_migrates_to_global(tmp_path, monkeypatch): assert research_db.claim_next("migration-worker")["id"] == "alice-run" +def test_owner_scoped_claim_migration_rolls_back_on_interruption(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(studio_db, "_schema_ready", False) + studio_db.upsert_chat_thread( + { + "id": "shared-thread", + "title": "Shared", + "modelType": "base", + "modelId": "model", + "createdAt": 1, + } + ) + conn = studio_db.get_connection() + try: + conn.execute("DROP TABLE research_thread_claims") + conn.execute( + """CREATE TABLE research_thread_claims ( + owner_subject TEXT NOT NULL, + thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE, + created_at INTEGER NOT NULL, + PRIMARY KEY(owner_subject, thread_id) + ) WITHOUT ROWID""" + ) + conn.execute( + "INSERT INTO research_thread_claims VALUES ('alice', 'shared-thread', 10)" + ) + conn.commit() + finally: + conn.close() + + # Simulate a crash midway through the migration (after RENAME/CREATE/INSERT, + # right before DROP). With the atomic transaction the whole rebuild must roll + # back, leaving the legacy owner-scoped table and its data intact. + real_connect = studio_db.sqlite3.connect + + class _FailingConnection(studio_db.sqlite3.Connection): + def execute(self, sql, *args, **kwargs): + if "DROP TABLE research_thread_claims_legacy" in sql: + raise RuntimeError("simulated crash during migration") + return super().execute(sql, *args, **kwargs) + + def _failing_connect(path, *args, **kwargs): + kwargs["factory"] = _FailingConnection + return real_connect(path, *args, **kwargs) + + monkeypatch.setattr(studio_db.sqlite3, "connect", _failing_connect) + studio_db._schema_ready = False + with pytest.raises(RuntimeError, match = "simulated crash"): + studio_db.get_connection() + + # Recover: the interrupted migration left nothing half-applied, so a clean boot + # completes the migration and preserves the original claim exactly once. + monkeypatch.setattr(studio_db.sqlite3, "connect", real_connect) + studio_db._schema_ready = False + conn = studio_db.get_connection() + try: + primary_key = [ + row["name"] + for row in conn.execute("PRAGMA table_info(research_thread_claims)").fetchall() + if row["pk"] + ] + claims = conn.execute( + "SELECT owner_subject, thread_id FROM research_thread_claims" + ).fetchall() + legacy = conn.execute( + "SELECT name FROM sqlite_master WHERE name = 'research_thread_claims_legacy'" + ).fetchall() + finally: + conn.close() + + assert primary_key == ["thread_id"] + assert [tuple(row) for row in claims] == [("alice", "shared-thread")] + assert legacy == [] + + def test_pruning_messages_preserves_runs_whose_user_message_survives(research_home): _create() studio_db.upsert_chat_message(