diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 48651aa..2eb235f 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -33,7 +33,8 @@ tests/ pytest — async via pytest-asyncio - **SQLAlchemy 2.x async syntax** — `select(Model).where(...)`, `await session.scalars(...)`. No legacy `Query` API. - **Sessions come from `async_session()` in `db.py`** — use `async with async_session() as session:` and commit explicitly. -- **Migrations are not yet wired.** When you add a column, also add a startup migration in `db.py:init_db()` until we adopt Alembic. Don't skip this and ship — it'll break prod. +- **Migrations are Alembic.** When you change `models.py`, generate a revision (`cd backend && alembic revision --autogenerate -m "..."`), read it, and commit it under `alembic/versions/`. `init_db()` in `db.py` runs `alembic upgrade head` on every boot (in a worker thread, since Alembic's `env.py` drives its own async engine — see the comment on `_run_alembic_sync`); a legacy DB that already has the full schema but no `alembic_version` table gets `stamp head` instead (see `_pre_alembic_schema_present`). Don't hand-write `ALTER TABLE` in `db.py` anymore — that's what `_run_migrations` used to be (kept only as dead code / rollback reference, see its docstring). +- **Boot-time migration assumes a single app instance** (which is what docker-compose runs). The stamp-vs-upgrade decision and the `alembic upgrade head` that follows are not atomic: two instances booting simultaneously against the same empty Postgres can both pick the upgrade path and race on `CREATE TABLE` ("relation already exists"). If this app is ever scaled to multiple replicas, run migrations as a one-off init container/job (`alembic upgrade head`) before starting replicas, instead of relying on boot-time migration. - **No raw SQL strings without a comment explaining why.** ORM first. ## Errors @@ -74,7 +75,7 @@ tests/ pytest — async via pytest-asyncio - [ ] `ruff check . && ruff format .` clean - [ ] `pytest tests/ -v` passes -- [ ] New columns / tables migrated in `db.py:init_db()` (until Alembic lands) +- [ ] New columns / tables: `alembic revision --autogenerate`, reviewed and committed under `alembic/versions/` - [ ] Logger used; no `print` - [ ] Error path tested - [ ] If you added a route, also added a Pydantic schema for body/response diff --git a/backend/alembic.ini b/backend/alembic.ini new file mode 100644 index 0000000..a1a16fa --- /dev/null +++ b/backend/alembic.ini @@ -0,0 +1,56 @@ +# Alembic configuration for the trainable backend. +# +# The actual DB URL is NOT read from here — env.py overrides it at runtime +# from `config.settings.database_url` (same source the app's async engine +# uses), so DATABASE_URL / .env stays the single source of truth. The +# sqlalchemy.url below is only a placeholder to satisfy tools that expect +# the key to exist. + +[alembic] +script_location = %(here)s/alembic +prepend_sys_path = . +path_separator = os +sqlalchemy.url = driver://user:pass@localhost/dbname + + +[post_write_hooks] +# See https://alembic.sqlalchemy.org/en/latest/autogenerate.html#post-write-hooks +# hooks = ruff +# ruff.type = exec +# ruff.executable = ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Logging configuration. +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/backend/alembic/env.py b/backend/alembic/env.py new file mode 100644 index 0000000..c731a26 --- /dev/null +++ b/backend/alembic/env.py @@ -0,0 +1,112 @@ +"""Alembic environment script. + +Wires Alembic to the app's own settings/models instead of a second, +hand-maintained config: + + - the DB URL always comes from ``config.settings.database_url`` (so + DATABASE_URL / .env stays the single source of truth — nothing new to + keep in sync in alembic.ini), + - ``target_metadata`` is ``db.Base.metadata`` after importing every model + module, so ``alembic revision --autogenerate`` sees the full schema. + +The app uses an async engine (asyncpg / aiosqlite), so migrations run +through SQLAlchemy's async engine too, following the pattern from the +Alembic cookbook for async applications: +https://alembic.sqlalchemy.org/en/latest/cookbook.html#using-asyncio-with-alembic +""" + +import asyncio +from logging.config import fileConfig + +from sqlalchemy import pool +from sqlalchemy.engine import Connection +from sqlalchemy.ext.asyncio import async_engine_from_config + +from alembic import context + +# Import the app's settings + models so Base.metadata is fully populated +# and the DB URL matches exactly what the running app uses. +from config import settings +from db import Base +import models # noqa: F401 (populates Base.metadata as a side effect) + +# this is the Alembic Config object, which provides access to the values +# within the .ini file in use. +config = context.config + +# Interpret the config file for Python logging. +# +# disable_existing_loggers=False is load-bearing: this env.py also runs +# in-process at app startup (init_db() -> _run_alembic_sync), after main.py +# and every module-level `logging.getLogger(__name__)` have already been +# created. fileConfig's default (True) would silently disable all of those +# app loggers the moment migrations run — killing app logging after boot +# (and breaking any caplog-based test that runs after an Alembic test). +if config.config_file_name is not None: + fileConfig(config.config_file_name, disable_existing_loggers=False) + +# The app's own metadata — autogenerate diffs against this. +target_metadata = Base.metadata + +# Always drive the URL from the app's settings, not alembic.ini. +config.set_main_option("sqlalchemy.url", settings.database_url) + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode (emits SQL, no DB connection).""" + url = settings.database_url + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + # Batch mode so future ALTER-style migrations work on SQLite (which + # cannot ALTER COLUMN / DROP COLUMN natively); a no-op on Postgres. + render_as_batch=True, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def do_run_migrations(connection: Connection) -> None: + context.configure( + connection=connection, + target_metadata=target_metadata, + # Batch mode so future ALTER-style migrations work on SQLite (which + # cannot ALTER COLUMN / DROP COLUMN natively); a no-op on Postgres. + render_as_batch=True, + ) + + with context.begin_transaction(): + context.run_migrations() + + +async def run_async_migrations() -> None: + """Run migrations against the async engine built from settings.""" + connectable = async_engine_from_config( + config.get_section(config.config_ini_section, {}), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + async with connectable.connect() as connection: + await connection.run_sync(do_run_migrations) + + await connectable.dispose() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + Called both from the `alembic` CLI (its own fresh event loop) and + programmatically from `db.py` via `asyncio.to_thread` (also a fresh + thread with no running loop) — `asyncio.run()` is safe in both cases. + """ + asyncio.run(run_async_migrations()) + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/backend/alembic/script.py.mako b/backend/alembic/script.py.mako new file mode 100644 index 0000000..fbc4b07 --- /dev/null +++ b/backend/alembic/script.py.mako @@ -0,0 +1,26 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/backend/alembic/versions/0bae8e0f765d_initial_schema.py b/backend/alembic/versions/0bae8e0f765d_initial_schema.py new file mode 100644 index 0000000..7d1e022 --- /dev/null +++ b/backend/alembic/versions/0bae8e0f765d_initial_schema.py @@ -0,0 +1,514 @@ +"""initial schema + +Revision ID: 0bae8e0f765d +Revises: +Create Date: 2026-07-18 23:25:19.238561 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = "0bae8e0f765d" +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "projects", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.Column("sandbox_config", sa.JSON(), nullable=True), + sa.Column("updated_at", sa.String(), nullable=True), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "experiments", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("project_id", sa.String(length=36), nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=True), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("hypothesis", sa.Text(), nullable=True), + sa.Column("state", sa.String(length=20), nullable=True), + sa.Column("started_at", sa.String(), nullable=True), + sa.Column("completed_at", sa.String(), nullable=True), + sa.Column("dataset_ref", sa.String(length=512), nullable=True), + sa.Column("instructions", sa.Text(), nullable=True), + sa.Column("tags", sa.JSON(), nullable=True), + sa.Column("pinned", sa.Boolean(), nullable=True), + sa.Column("archived", sa.Boolean(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.Column("updated_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_experiments_project_id"), "experiments", ["project_id"], unique=False + ) + op.create_index( + op.f("ix_experiments_session_id"), "experiments", ["session_id"], unique=False + ) + op.create_index( + op.f("ix_experiments_state"), "experiments", ["state"], unique=False + ) + op.create_table( + "dataset_versions", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("project_id", sa.String(length=36), nullable=False), + sa.Column("kind", sa.String(length=20), nullable=False), + sa.Column("name", sa.String(length=255), nullable=True), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("hash", sa.String(length=64), nullable=False), + sa.Column("path", sa.String(length=512), nullable=False), + sa.Column("size_bytes", sa.Integer(), nullable=True), + sa.Column("parent_id", sa.Integer(), nullable=True), + sa.Column("parent_hash", sa.String(length=64), nullable=True), + sa.Column("source_session_id", sa.String(length=36), nullable=True), + sa.Column("source_experiment_id", sa.String(length=36), nullable=True), + sa.Column("metadata", sa.JSON(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["parent_id"], + ["dataset_versions.id"], + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + ), + sa.ForeignKeyConstraint( + ["source_experiment_id"], + ["experiments.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_dataset_versions_hash"), "dataset_versions", ["hash"], unique=False + ) + op.create_index( + op.f("ix_dataset_versions_kind"), "dataset_versions", ["kind"], unique=False + ) + op.create_index( + op.f("ix_dataset_versions_parent_id"), + "dataset_versions", + ["parent_id"], + unique=False, + ) + op.create_index( + op.f("ix_dataset_versions_project_id"), + "dataset_versions", + ["project_id"], + unique=False, + ) + op.create_index( + op.f("ix_dataset_versions_source_experiment_id"), + "dataset_versions", + ["source_experiment_id"], + unique=False, + ) + op.create_index( + op.f("ix_dataset_versions_source_session_id"), + "dataset_versions", + ["source_session_id"], + unique=False, + ) + op.create_table( + "registered_models", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("project_id", sa.String(length=36), nullable=False), + sa.Column("experiment_id", sa.String(length=36), nullable=True), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("version", sa.Integer(), nullable=False), + sa.Column("source_session_id", sa.String(length=36), nullable=True), + sa.Column("artifact_uri", sa.String(length=512), nullable=False), + sa.Column("artifact_size_bytes", sa.Integer(), nullable=True), + sa.Column("metrics_summary", sa.JSON(), nullable=True), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("hyperparams", sa.JSON(), nullable=True), + sa.Column("dataset_refs", sa.JSON(), nullable=True), + sa.Column("metrics_history", sa.JSON(), nullable=True), + sa.Column("serving_app_path", sa.String(length=512), nullable=True), + sa.Column("api_key", sa.String(length=64), nullable=True), + sa.Column("framework", sa.String(length=50), nullable=True), + sa.Column("status", sa.String(length=20), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["experiment_id"], + ["experiments.id"], + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_registered_models_experiment_id"), + "registered_models", + ["experiment_id"], + unique=False, + ) + op.create_index( + op.f("ix_registered_models_project_id"), + "registered_models", + ["project_id"], + unique=False, + ) + op.create_index( + op.f("ix_registered_models_source_session_id"), + "registered_models", + ["source_session_id"], + unique=False, + ) + op.create_table( + "sessions", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("project_id", sa.String(length=36), nullable=True), + sa.Column("name", sa.String(length=255), nullable=True), + sa.Column("experiment_id", sa.String(length=36), nullable=True), + sa.Column("state", sa.String(length=50), nullable=True), + sa.Column("model", sa.String(length=100), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.Column("updated_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["experiment_id"], + ["experiments.id"], + ), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_sessions_experiment_id"), "sessions", ["experiment_id"], unique=False + ) + op.create_index( + op.f("ix_sessions_project_id"), "sessions", ["project_id"], unique=False + ) + op.create_table( + "artifacts", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("stage", sa.String(length=50), nullable=False), + sa.Column("artifact_type", sa.String(length=50), nullable=False), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column("path", sa.String(length=512), nullable=False), + sa.Column("s3_path", sa.String(length=512), nullable=True), + sa.Column("metadata", sa.JSON(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_artifacts_session_id"), "artifacts", ["session_id"], unique=False + ) + op.create_table( + "deployments", + sa.Column("id", sa.String(length=36), nullable=False), + sa.Column("model_id", sa.String(length=36), nullable=False), + sa.Column("endpoint_url", sa.String(length=512), nullable=True), + sa.Column("status", sa.String(length=20), nullable=True), + sa.Column("error", sa.Text(), nullable=True), + sa.Column("modal_app", sa.String(length=255), nullable=True), + sa.Column("modal_function", sa.String(length=255), nullable=True), + sa.Column("compute", sa.String(length=20), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.Column("updated_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["model_id"], + ["registered_models.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_deployments_model_id"), "deployments", ["model_id"], unique=False + ) + op.create_table( + "experiment_datasets", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("experiment_id", sa.String(length=36), nullable=False), + sa.Column("dataset_version_id", sa.Integer(), nullable=False), + sa.Column("role", sa.String(length=20), nullable=False), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["dataset_version_id"], ["dataset_versions.id"], ondelete="CASCADE" + ), + sa.ForeignKeyConstraint( + ["experiment_id"], ["experiments.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_experiment_datasets_dataset_version_id"), + "experiment_datasets", + ["dataset_version_id"], + unique=False, + ) + op.create_index( + op.f("ix_experiment_datasets_experiment_id"), + "experiment_datasets", + ["experiment_id"], + unique=False, + ) + op.create_table( + "log_events", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("stage", sa.String(length=50), nullable=False), + sa.Column("step", sa.Integer(), nullable=False), + sa.Column("key", sa.String(length=255), nullable=False), + sa.Column("type", sa.String(length=30), nullable=False), + sa.Column("run_tag", sa.String(length=100), nullable=True), + sa.Column("payload", sa.JSON(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_log_events_session_id"), "log_events", ["session_id"], unique=False + ) + op.create_table( + "messages", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("role", sa.String(length=50), nullable=False), + sa.Column("content", sa.Text(), nullable=False), + sa.Column("metadata", sa.JSON(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_messages_session_id"), "messages", ["session_id"], unique=False + ) + op.create_table( + "metrics", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("stage", sa.String(length=50), nullable=False), + sa.Column("step", sa.Integer(), nullable=False), + sa.Column("name", sa.String(length=100), nullable=False), + sa.Column("value", sa.Float(), nullable=False), + sa.Column("run_tag", sa.String(length=100), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_metrics_session_id"), "metrics", ["session_id"], unique=False + ) + op.create_table( + "processed_dataset_meta", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("experiment_id", sa.String(length=36), nullable=False), + sa.Column("columns", sa.JSON(), nullable=False), + sa.Column("feature_columns", sa.JSON(), nullable=True), + sa.Column("target_column", sa.String(length=255), nullable=True), + sa.Column("total_rows", sa.Integer(), nullable=False), + sa.Column("train_rows", sa.Integer(), nullable=True), + sa.Column("val_rows", sa.Integer(), nullable=True), + sa.Column("test_rows", sa.Integer(), nullable=True), + sa.Column("quality_stats", sa.JSON(), nullable=True), + sa.Column("source_files", sa.JSON(), nullable=True), + sa.Column("output_files", sa.JSON(), nullable=True), + sa.Column("s3_synced", sa.String(length=10), nullable=True), + sa.Column("s3_prefix", sa.String(length=512), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["experiment_id"], + ["experiments.id"], + ), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_processed_dataset_meta_experiment_id"), + "processed_dataset_meta", + ["experiment_id"], + unique=False, + ) + op.create_index( + op.f("ix_processed_dataset_meta_session_id"), + "processed_dataset_meta", + ["session_id"], + unique=False, + ) + op.create_table( + "run_snapshots", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("experiment_id", sa.String(length=36), nullable=True), + sa.Column("session_id", sa.String(length=36), nullable=True), + sa.Column("dataset_hash", sa.String(length=64), nullable=True), + sa.Column("code_hash", sa.String(length=64), nullable=True), + sa.Column("hyperparams", sa.JSON(), nullable=True), + sa.Column("env_lockfile", sa.Text(), nullable=True), + sa.Column("manifest_uri", sa.String(length=512), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["experiment_id"], + ["experiments.id"], + ), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_run_snapshots_experiment_id"), + "run_snapshots", + ["experiment_id"], + unique=False, + ) + op.create_index( + op.f("ix_run_snapshots_session_id"), + "run_snapshots", + ["session_id"], + unique=False, + ) + op.create_table( + "tasks", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("subject", sa.String(length=255), nullable=False), + sa.Column("active_form", sa.String(length=255), nullable=True), + sa.Column("short_description", sa.Text(), nullable=True), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("status", sa.String(length=20), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.Column("updated_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["session_id"], + ["sessions.id"], + ), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index(op.f("ix_tasks_session_id"), "tasks", ["session_id"], unique=False) + op.create_table( + "usage_events", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("session_id", sa.String(length=36), nullable=False), + sa.Column("project_id", sa.String(length=36), nullable=True), + sa.Column("kind", sa.String(length=20), nullable=False), + sa.Column("agent_type", sa.String(length=50), nullable=True), + sa.Column("agent_id", sa.String(length=100), nullable=True), + sa.Column("provider", sa.String(length=50), nullable=True), + sa.Column("model", sa.String(length=100), nullable=True), + sa.Column("input_tokens", sa.Integer(), nullable=True), + sa.Column("output_tokens", sa.Integer(), nullable=True), + sa.Column("cache_read_input_tokens", sa.Integer(), nullable=True), + sa.Column("cache_creation_input_tokens", sa.Integer(), nullable=True), + sa.Column("sandbox_seconds", sa.Float(), nullable=True), + sa.Column("gpu_type", sa.String(length=50), nullable=True), + sa.Column("cost_usd", sa.Float(), nullable=True), + sa.Column("is_error", sa.Boolean(), nullable=True), + sa.Column("extra", sa.JSON(), nullable=True), + sa.Column("created_at", sa.String(), nullable=True), + sa.ForeignKeyConstraint(["session_id"], ["sessions.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_usage_events_project_id"), "usage_events", ["project_id"], unique=False + ) + op.create_index( + op.f("ix_usage_events_session_id"), "usage_events", ["session_id"], unique=False + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index(op.f("ix_usage_events_session_id"), table_name="usage_events") + op.drop_index(op.f("ix_usage_events_project_id"), table_name="usage_events") + op.drop_table("usage_events") + op.drop_index(op.f("ix_tasks_session_id"), table_name="tasks") + op.drop_table("tasks") + op.drop_index(op.f("ix_run_snapshots_session_id"), table_name="run_snapshots") + op.drop_index(op.f("ix_run_snapshots_experiment_id"), table_name="run_snapshots") + op.drop_table("run_snapshots") + op.drop_index( + op.f("ix_processed_dataset_meta_session_id"), + table_name="processed_dataset_meta", + ) + op.drop_index( + op.f("ix_processed_dataset_meta_experiment_id"), + table_name="processed_dataset_meta", + ) + op.drop_table("processed_dataset_meta") + op.drop_index(op.f("ix_metrics_session_id"), table_name="metrics") + op.drop_table("metrics") + op.drop_index(op.f("ix_messages_session_id"), table_name="messages") + op.drop_table("messages") + op.drop_index(op.f("ix_log_events_session_id"), table_name="log_events") + op.drop_table("log_events") + op.drop_index( + op.f("ix_experiment_datasets_experiment_id"), table_name="experiment_datasets" + ) + op.drop_index( + op.f("ix_experiment_datasets_dataset_version_id"), + table_name="experiment_datasets", + ) + op.drop_table("experiment_datasets") + op.drop_index(op.f("ix_deployments_model_id"), table_name="deployments") + op.drop_table("deployments") + op.drop_index(op.f("ix_artifacts_session_id"), table_name="artifacts") + op.drop_table("artifacts") + op.drop_index(op.f("ix_sessions_project_id"), table_name="sessions") + op.drop_index(op.f("ix_sessions_experiment_id"), table_name="sessions") + op.drop_table("sessions") + op.drop_index( + op.f("ix_registered_models_source_session_id"), table_name="registered_models" + ) + op.drop_index( + op.f("ix_registered_models_project_id"), table_name="registered_models" + ) + op.drop_index( + op.f("ix_registered_models_experiment_id"), table_name="registered_models" + ) + op.drop_table("registered_models") + op.drop_index( + op.f("ix_dataset_versions_source_session_id"), table_name="dataset_versions" + ) + op.drop_index( + op.f("ix_dataset_versions_source_experiment_id"), table_name="dataset_versions" + ) + op.drop_index(op.f("ix_dataset_versions_project_id"), table_name="dataset_versions") + op.drop_index(op.f("ix_dataset_versions_parent_id"), table_name="dataset_versions") + op.drop_index(op.f("ix_dataset_versions_kind"), table_name="dataset_versions") + op.drop_index(op.f("ix_dataset_versions_hash"), table_name="dataset_versions") + op.drop_table("dataset_versions") + op.drop_index(op.f("ix_experiments_state"), table_name="experiments") + op.drop_index(op.f("ix_experiments_session_id"), table_name="experiments") + op.drop_index(op.f("ix_experiments_project_id"), table_name="experiments") + op.drop_table("experiments") + op.drop_table("projects") + # ### end Alembic commands ### diff --git a/backend/db.py b/backend/db.py index c57c0c9..6fb9510 100644 --- a/backend/db.py +++ b/backend/db.py @@ -1,6 +1,8 @@ """PostgreSQL database setup with async SQLAlchemy.""" +import asyncio import logging +from pathlib import Path from sqlalchemy import event, inspect, text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine @@ -41,7 +43,28 @@ class Base(DeclarativeBase): def _run_migrations(connection): - """Add columns/tables that create_all won't add to existing tables.""" + """DEAD CODE — retained for history/rollback reference only, no longer + called from `init_db`. See PR #121 (introduce Alembic migrations). + + This ~430-line hand-maintained sequence of guarded ALTER TABLE / + destructive DELETE / backfill statements used to re-run on every boot + with no versioning or rollback. It has been superseded by the Alembic + revision chain in `alembic/versions/` — the initial revision was + generated by autogenerating against `models.py` (which, by the time + this was cut over, already reflected every change this function used + to apply) and verified schema-equivalent against a DB built by + `Base.metadata.create_all` + this function, on both SQLite and + Postgres. The only discovered difference: this function creates two + redundant duplicate indexes (`ix_dataset_versions_source_experiment`, + `ix_registered_models_source_session` — note the missing `_id` suffix) + that shadow the index SQLAlchemy already creates from the columns' + `index=True` in `models.py`; the Alembic revision does not recreate + that duplication. + + If Alembic ever needs to be rolled back, restore the call to this + function from `init_db` (see git history) — do not delete it outright + without checking git blame/PR #121 first. + """ insp = inspect(connection) # Add s3_path to artifacts if missing @@ -477,6 +500,76 @@ def _run_migrations(connection): logger.debug("Index %s create skipped: %s", idx_name, e) +# Tables representative of the fully-migrated pre-Alembic schema. The +# initial Alembic revision (alembic/versions/0bae8e0f765d_initial_schema.py) +# was autogenerated against this exact end-state, so a DB that already has +# all of these is assumed current and only needs `alembic stamp head` (mark +# as up to date, no DDL) rather than `alembic upgrade head` (which would +# try — and fail — to (re)create tables that already exist). This assumption +# holds for any DB that has been through a boot of the old `_run_migrations` +# path, since that function ran on every startup and brought the schema +# fully current before this cutover. A DB with only *some* of these tables +# (older than assumed, never fully migrated) isn't handled here and falls +# through to `upgrade head`, which will surface a loud DDL error rather than +# silently stamping an incomplete schema. +_LEGACY_SCHEMA_MARKER_TABLES = ( + "projects", + "sessions", + "experiments", + "registered_models", + "dataset_versions", +) + + +def _pre_alembic_schema_present(connection) -> bool: + """True if this DB already has the full current schema but hasn't been + stamped with an Alembic version yet — see `_LEGACY_SCHEMA_MARKER_TABLES`.""" + insp = inspect(connection) + if insp.has_table("alembic_version"): + return False + return all(insp.has_table(t) for t in _LEGACY_SCHEMA_MARKER_TABLES) + + +def _alembic_config(): + from alembic.config import Config + + backend_dir = Path(__file__).resolve().parent + cfg = Config(str(backend_dir / "alembic.ini")) + cfg.set_main_option("script_location", str(backend_dir / "alembic")) + # Absolute, not the ini's literal "." — makes `env.py`'s `from config + # import settings` / `from db import Base` / `import models` resolve + # regardless of the caller's current working directory. + cfg.set_main_option("prepend_sys_path", str(backend_dir)) + # env.py also reads settings.database_url directly, but setting it here + # too keeps `alembic.config.Config` usable standalone (e.g. if some + # future caller inspects cfg.get_main_option("sqlalchemy.url")). + cfg.set_main_option("sqlalchemy.url", settings.database_url) + return cfg + + +def _run_alembic_sync(*, stamp_only: bool) -> None: + """Runs Alembic's `upgrade head` (or `stamp head`) synchronously. + + Must be called via `asyncio.to_thread` from async code: Alembic's + env.py drives the app's *async* engine internally with its own + `asyncio.run(...)` (see alembic/env.py), which raises if invoked from + a thread that already has a running event loop. A fresh worker thread + has no running loop, so it's safe there regardless of caller. + """ + from alembic import command + + cfg = _alembic_config() + if stamp_only: + command.stamp(cfg, "head") + logger.info( + "[DB] Existing pre-Alembic schema detected — stamped alembic_version" + " at head without altering schema" + ) + else: + command.upgrade(cfg, "head") + logger.info("[DB] Alembic migrations applied up to head") + + async def init_db(): from models import ( # noqa: F401 Artifact, @@ -497,8 +590,20 @@ async def init_db(): ) async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) - await conn.run_sync(_run_migrations) + stamp_only = await conn.run_sync(_pre_alembic_schema_present) + + # Alembic owns schema DDL now — `Base.metadata.create_all` / + # `_run_migrations` are no longer called here. See `_run_migrations`'s + # docstring and PR #121 for the schema-equivalence verification. + # + # NOTE: single-instance assumption. The stamp-vs-upgrade check above and + # the Alembic run below use separate connections and are not atomic, so + # two app instances booting concurrently against the same empty DB could + # both take the upgrade path and race on CREATE TABLE. Fine for the + # docker-compose deployment (one backend container); if this ever runs + # with multiple replicas, apply migrations via a one-off init job before + # starting the app instead (see backend/AGENTS.md, "Database"). + await asyncio.to_thread(_run_alembic_sync, stamp_only=stamp_only) async def get_db() -> AsyncSession: diff --git a/backend/requirements.txt b/backend/requirements.txt index 37850d1..cf2a4c1 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -3,6 +3,7 @@ uvicorn[standard]==0.30.0 sqlalchemy[asyncio]==2.0.35 asyncpg>=0.30.0 aiosqlite==0.20.0 +alembic>=1.13.0,<2.0.0 anthropic>=0.42.0 claude-agent-sdk>=0.1.51 openai>=1.54.0 diff --git a/backend/tests/test_alembic_migrations.py b/backend/tests/test_alembic_migrations.py new file mode 100644 index 0000000..a478ac9 --- /dev/null +++ b/backend/tests/test_alembic_migrations.py @@ -0,0 +1,112 @@ +"""Tests for the Alembic boot-time migration path in db.py. + +Covers the stamp-vs-upgrade heuristic (`_pre_alembic_schema_present`) and +schema-sanity checks that `alembic upgrade head` / `alembic stamp head` on a +fresh SQLite DB behave as init_db() expects (regression guard for PR #121 and +the review finding on #153). + +These are single-connection tests; the concurrent multi-instance boot race is +intentionally NOT simulated here — boot-time migration assumes one app +instance (see the NOTE in `init_db()` and backend/AGENTS.md, "Database"). +""" + +import pytest +from sqlalchemy import create_engine, inspect, text + +from db import ( + _LEGACY_SCHEMA_MARKER_TABLES, + _pre_alembic_schema_present, + _run_alembic_sync, +) + + +def _make_sync_sqlite_engine(): + return create_engine("sqlite://") + + +def _create_tables(conn, names): + for name in names: + conn.execute(text(f'CREATE TABLE "{name}" (id INTEGER PRIMARY KEY)')) + + +def test_heuristic_empty_db_takes_upgrade_path(): + engine = _make_sync_sqlite_engine() + with engine.begin() as conn: + assert _pre_alembic_schema_present(conn) is False + + +def test_heuristic_full_legacy_schema_takes_stamp_path(): + engine = _make_sync_sqlite_engine() + with engine.begin() as conn: + _create_tables(conn, _LEGACY_SCHEMA_MARKER_TABLES) + assert _pre_alembic_schema_present(conn) is True + + +def test_heuristic_already_stamped_db_takes_upgrade_path(): + engine = _make_sync_sqlite_engine() + with engine.begin() as conn: + _create_tables(conn, (*_LEGACY_SCHEMA_MARKER_TABLES, "alembic_version")) + assert _pre_alembic_schema_present(conn) is False + + +def test_heuristic_partial_legacy_schema_falls_through_to_upgrade(): + # A DB with only *some* marker tables must NOT be stamped — upgrade head + # should run and surface a loud DDL error rather than silently stamping + # an incomplete schema (see the comment on _LEGACY_SCHEMA_MARKER_TABLES). + engine = _make_sync_sqlite_engine() + with engine.begin() as conn: + _create_tables(conn, _LEGACY_SCHEMA_MARKER_TABLES[:2]) + assert _pre_alembic_schema_present(conn) is False + + +@pytest.fixture() +def _sqlite_file_db(tmp_path, monkeypatch): + """Point settings.database_url at a fresh file-backed SQLite DB. + + Both db._alembic_config() and alembic/env.py read + config.settings.database_url at call time, so patching the attribute is + enough to redirect the whole Alembic run. + """ + from config import settings + + db_path = tmp_path / "alembic_test.db" + monkeypatch.setattr(settings, "database_url", f"sqlite+aiosqlite:///{db_path}") + return db_path + + +def test_upgrade_head_on_fresh_db_creates_full_schema(_sqlite_file_db): + # Sync test on purpose: _run_alembic_sync runs asyncio.run() internally + # (via alembic/env.py) and must be called with no running event loop. + _run_alembic_sync(stamp_only=False) + + engine = create_engine(f"sqlite:///{_sqlite_file_db}") + with engine.connect() as conn: + insp = inspect(conn) + for table in _LEGACY_SCHEMA_MARKER_TABLES: + assert insp.has_table(table), f"upgrade head did not create {table!r}" + assert insp.has_table("alembic_version") + # Next boot must take the upgrade path (a no-op at head), not stamp. + assert _pre_alembic_schema_present(conn) is False + + +def test_stamp_head_records_version_without_touching_schema(_sqlite_file_db): + # Simulate the legacy-deployment boot: schema exists (markers suffice for + # the stamp command itself), no alembic_version yet. + engine = create_engine(f"sqlite:///{_sqlite_file_db}") + with engine.begin() as conn: + _create_tables(conn, _LEGACY_SCHEMA_MARKER_TABLES) + + _run_alembic_sync(stamp_only=True) + + with engine.connect() as conn: + insp = inspect(conn) + assert insp.has_table("alembic_version") + version = conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + assert version, "stamp head left alembic_version empty" + # stamp must not create any migration-managed tables beyond the + # markers we made + alembic_version itself. + assert set(insp.get_table_names()) == { + *_LEGACY_SCHEMA_MARKER_TABLES, + "alembic_version", + } + assert _pre_alembic_schema_present(conn) is False