SurfSense/surfsense_backend/alembic/env.py

import asyncio
import logging
import os
import sys
from logging.config import fileConfig

import sqlalchemy as sa
from alembic.script import ScriptDirectory
from sqlalchemy import pool
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config

from alembic import context

# Ensure the app directory is in the Python path
# This allows Alembic to find your models
sys.path.insert(0, os.path.realpath(os.path.join(os.path.dirname(__file__), "..")))

# Import your models base
from app.db import Base  # Assuming your Base is defined in app.db

# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config

# Override SQLAlchemy URL from environment variables when available
database_url = os.getenv("DATABASE_URL")
if database_url:
    config.set_main_option("sqlalchemy.url", database_url)

# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
    fileConfig(config.config_file_name)

# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = Base.metadata

MIGRATION_ADVISORY_LOCK_NAMESPACE = "surfsense"
MIGRATION_ADVISORY_LOCK_NAME = "alembic_migrations"

# Migration 170 renamed searchspaces -> workspaces, so a ``workspaces`` table
# can only exist once the schema is at revision >= 170. If it exists while the
# recorded revision is missing or still pre-170, the schema did not come from
# this migration history at all -- it was created by the startup bootstrap
# (``Base.metadata.create_all`` in ``app.db.create_db_and_tables``), which
# always builds the *current* model shape. Replaying history against such a
# schema fails (e.g. migration 5's ``ALTER COLUMN ... TYPE`` is rejected
# because the column already sits in zero_publication's column list), so the
# schema is adopted by stamping head instead.
BOOTSTRAP_MARKER_TABLE = "workspaces"
RENAME_REVISION = "170"


def _stamp_head(connection: Connection, script: ScriptDirectory) -> None:
    context.get_context().stamp(script, script.get_current_head())
    if connection.in_transaction():
        # The outer begin_transaction() is a no-op under
        # transaction_per_migration, so commit explicitly.
        connection.commit()


def _fast_forward_fresh_db(connection: Connection) -> bool:
    """Build a fresh (empty) DB at head via create_all instead of replaying.

    Historical migrations were written against the pre-workspace-rename
    schema (``searchspaces``, ``search_space_id``), while migration 0's
    ``create_all`` builds the *current* models -- so replaying the chain on a
    fresh DB crashes as soon as a migration touches a renamed object (first
    at migration 18). A fresh DB needs no history: create the head-shape
    schema directly, mirror migration 0's indexes, create the Zero
    publication, and stamp head. Replay remains only for legacy DBs that
    genuinely contain the old objects.

    ponytail: seed-data migrations (114/128 default prompts) are skipped on
    this path, same as always for create_all-bootstrapped DBs; the app copes
    with missing seeds. If seeds ever become mandatory, add a runtime seeding
    step rather than resurrecting the replay.
    """
    for table in ("documents", "searchspaces", BOOTSTRAP_MARKER_TABLE):
        if connection.execute(
            sa.text("SELECT to_regclass(:t)"), {"t": table}
        ).scalar():
            return False
    if connection.execute(sa.text("SELECT to_regclass('alembic_version')")).scalar():
        current = connection.execute(
            sa.text("SELECT version_num FROM alembic_version")
        ).scalar()
        if current:
            return False

    logging.getLogger("alembic.env").info(
        "Fresh database detected: creating head-shape schema via create_all "
        "and stamping head instead of replaying migration history."
    )
    connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS vector"))
    connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
    Base.metadata.create_all(bind=connection)
    # Same core indexes migration 0 created (runtime setup_indexes() adds the
    # rest concurrently on app boot).
    connection.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS document_vector_index ON documents "
            "USING hnsw (embedding public.vector_cosine_ops)"
        )
    )
    connection.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS document_search_index ON documents "
            "USING gin (to_tsvector('english', content))"
        )
    )
    connection.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS chucks_vector_index ON chunks "
            "USING hnsw (embedding public.vector_cosine_ops)"
        )
    )
    connection.execute(
        sa.text(
            "CREATE INDEX IF NOT EXISTS chucks_search_index ON chunks "
            "USING gin (to_tsvector('english', content))"
        )
    )

    from app.zero_publication import ensure_publication

    ensure_publication(connection)

    _stamp_head(connection, ScriptDirectory.from_config(config))
    return True


def _adopt_bootstrapped_schema(connection: Connection) -> bool:
    """Stamp head instead of replaying history on a create_all-created DB.

    Returns True when the schema was adopted (migrations must then be
    skipped for this run).

    ponytail: assumes the bootstrapped schema matches the checked-out models
    (true whenever the backend booted on this checkout, since create_all runs
    on every startup). If the checkout moved ahead without a backend boot,
    column-level drift from the skipped migrations is possible; the upgrade
    path is re-bootstrapping (boot the backend once) before stamping.
    """
    marker = connection.execute(
        sa.text("SELECT to_regclass(:t)"), {"t": BOOTSTRAP_MARKER_TABLE}
    ).scalar()
    if marker is None:
        return False

    # Guard against a legacy-shape DB that merely had missing tables filled in
    # by a later create_all: adoption requires the core tables to be in the
    # current (post-rename) shape too, not just the marker table to exist.
    documents_renamed = connection.execute(
        sa.text(
            "SELECT 1 FROM information_schema.columns "
            "WHERE table_schema = current_schema() "
            "AND table_name = 'documents' AND column_name = 'workspace_id'"
        )
    ).scalar()
    if not documents_renamed:
        return False

    current = None
    if connection.execute(sa.text("SELECT to_regclass('alembic_version')")).scalar():
        current = connection.execute(
            sa.text("SELECT version_num FROM alembic_version")
        ).scalar()

    script = ScriptDirectory.from_config(config)
    pre_rename_revisions = {
        rev.revision for rev in script.iterate_revisions(RENAME_REVISION, "base")
    } - {RENAME_REVISION}
    if current is not None and current not in pre_rename_revisions:
        # Genuinely migration-managed at >= 170; run migrations normally.
        return False

    logging.getLogger("alembic.env").info(
        "Adopting bootstrap-created schema (%r exists, recorded revision %r "
        "predates the workspace rename): stamping %s instead of replaying "
        "migration history.",
        BOOTSTRAP_MARKER_TABLE,
        current,
        script.get_current_head(),
    )
    _stamp_head(connection, script)
    return True


# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.


def run_migrations_offline() -> None:
    """Run migrations in 'offline' mode.

    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.

    Calls to context.execute() here emit the given string to the
    script output.

    """
    url = config.get_main_option("sqlalchemy.url")
    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )

    with context.begin_transaction():
        context.run_migrations()


def do_run_migrations(connection: Connection) -> None:
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        transaction_per_migration=True,
    )

    lock_params = {
        "namespace": MIGRATION_ADVISORY_LOCK_NAMESPACE,
        "name": MIGRATION_ADVISORY_LOCK_NAME,
    }
    connection.execute(
        sa.text("SELECT pg_advisory_lock(hashtext(:namespace), hashtext(:name))"),
        lock_params,
    )
    try:
        if not _fast_forward_fresh_db(connection) and not _adopt_bootstrapped_schema(
            connection
        ):
            with context.begin_transaction():
                context.run_migrations()
    finally:
        connection.execute(
            sa.text("SELECT pg_advisory_unlock(hashtext(:namespace), hashtext(:name))"),
            lock_params,
        )


async def run_async_migrations() -> None:
    """In this scenario we need to create an Engine
    and associate a connection with the context.

    """

    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."""

    asyncio.run(run_async_migrations())


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()