Repositories

def create_async_engine(datasource: str | URL) -> AsyncEngine:
    """Build the asyncpg engine an `AsyncDBRepository`'s session is bound to.

    Mirrors [`create_engine`][xmas_core.repositories.engine.create_engine]; the schema
    migrations still need a synchronous engine, so an application creates both.

    Args:
        datasource: A connection uri or a `sqlalchemy.URL`.

    Returns:
        A PostgreSQL async engine bound to the configured schema.

    Raises:
        DatasourceError: The URL could not be parsed, or names a database other than
            PostgreSQL.
    """
    url = _postgresql_url(datasource)
    kwargs, options = _engine_kwargs()
    engine = _sa_create_async_engine(
        url.set(drivername="postgresql+asyncpg"),
        connect_args={"timeout": 5.0},
        execution_options=options,
        **kwargs,
    )
    return engine

def create_engine(datasource: str | URL) -> Engine:
    """Build the synchronous engine a repository's session is bound to.

    The engine owns a connection pool for as long as it lives; the caller that built it
    disposes it. `CORETABLE_SCHEMA` is bound to the configured schema, which
    [`for_schema`][xmas_core.repositories.engine.for_schema] rebinds per tenant.

    Args:
        datasource: A connection uri or a `sqlalchemy.URL`.

    Returns:
        A PostgreSQL engine bound to the configured schema.

    Raises:
        DatasourceError: The URL could not be parsed, or names a database other than
            PostgreSQL.
    """
    url = _postgresql_url(datasource)
    kwargs, options = _engine_kwargs()
    engine = _sa_create_engine(
        url.set(drivername="postgresql+psycopg"),
        connect_args={"connect_timeout": 5},
        execution_options=options,
        **kwargs,
    )
    return engine

def drop_schema(engine: Engine, schema: str | None = None) -> None:
    """Delete the coretable and its related tables and functions from one schema.

    A schema with no alembic revision is left alone: there is nothing this created to take
    back. A database created before 2.0 carries objects this revision knows nothing about and
    is dropped at the database level (`DROP SCHEMA <schema> CASCADE`) instead.

    Args:
        engine: A PostgreSQL engine, as built by
            [`create_engine`][xmas_core.repositories.engine.create_engine].
        schema: The schema to drop from. Defaults as in
            [`ensure_schema`][xmas_core.repositories.engine.ensure_schema].

    Raises:
        RuntimeError: The schema is at a revision this release does not know.
    """
    logger.debug("deleting tables")
    schema = (
        schema or _schema_map(engine).get(CORETABLE_SCHEMA) or get_settings().db_schema
    )
    # A downgrade creates the version table first, so an empty schema would be left with one
    # and a missing schema would fail; this reads without creating anything. On a connection
    # of its own: the read autobegins a transaction, and alembic leaves a transaction it did
    # not begin uncommitted, so the downgrade would be rolled back when the block closes.
    with engine.connect() as conn:
        schema = _resolve_schema(conn, schema)
        revision = MigrationContext.configure(
            conn, opts={"version_table_schema": schema}
        ).get_current_revision()
    if revision is None:
        logger.debug("no revision in schema %s, nothing to drop", schema)
        return
    alembic_cfg = _alembic_config(schema)
    _refuse_unknown_revisions(
        engine, script.ScriptDirectory.from_config(alembic_cfg), schema, [revision]
    )
    with engine.connect() as conn:
        alembic_cfg.attributes["connection"] = conn
        command.downgrade(alembic_cfg, "base")

def ensure_schema(
    engine: Engine, schema: str | None = None, *, bootstrap: bool = False
) -> None:
    """Check that one schema behind `engine` is in the state the repository assumes.

    Checks that the schema is accessible, at the current revision, and that the stored SRID
    matches the configured one. Only with `bootstrap=True` does it change anything: it creates
    a missing schema and runs the alembic migrations - creating the tables if the schema is
    empty - which makes it the only call a tenant's provisioning has to make, and the only one
    needing a role privileged enough to create a schema. Without it a missing, empty or
    outdated schema is an error: a mistyped name, or an unrelated existing schema, must not
    become a seeded tenant just because something tried to read from it. Call it once per
    schema, at startup; it is not something a repository does on its own, so a short-lived
    repository never pays for it. Several schemas are brought to head one after another, not from threads
    of one process: alembic keeps its migration context in module globals.

    Args:
        engine: A PostgreSQL engine, as built by
            [`create_engine`][xmas_core.repositories.engine.create_engine].
        schema: The schema to bring to head. Defaults to the one `engine` is bound to, so
            `ensure_schema(for_schema(engine, tenant))` migrates that tenant rather than the
            process default - then to the configured one, and with none configured to the
            connection's `current_schema()`, which is what the engine renders for it.
        bootstrap: Create the schema if it does not exist, and bring it to head.

    Raises:
        RuntimeError: The schema is missing, empty or at an older revision and `bootstrap`
            is false, the current user may not create it, the user lacks the required
            privilege on the schema, the stored SRID differs from the
            configured one, a coretable was found with no alembic revision, or the schema
            is at a revision this release does not know (a 1.x or pre-2.0 database).
    """
    settings = get_settings()
    schema = schema or _schema_map(engine).get(CORETABLE_SCHEMA) or settings.db_schema

    def _check_schema_accessibility(privilege: str) -> None:
        """Raises an exception if the schema does not exist or is not accessible to the current user."""
        user, allowed = conn.execute(
            text("SELECT current_user, has_schema_privilege(:schema, :privilege)"),
            {"schema": schema, "privilege": privilege},
        ).one()
        if not allowed:
            raise RuntimeError(f"User {user} lacks {privilege} on schema '{schema}'")

    def _create_schema_if_missing(name: str) -> None:
        """Create schema `name` when nothing is there yet and `bootstrap` allows it."""
        # matched against the stored name rather than parsed as an identifier, which is
        # what `to_regnamespace` would do - it answers NULL for an existing `MixedCase`
        if conn.execute(
            text("SELECT 1 FROM pg_namespace WHERE nspname = :schema"),
            {"schema": name},
        ).scalar():
            return
        if not bootstrap:
            raise RuntimeError(
                f"Schema '{name}' does not exist; create it with "
                "`xmas db create-schema` or `ensure_schema(..., bootstrap=True)`"
            )
        # only when it is genuinely missing: `CREATE SCHEMA IF NOT EXISTS` wants CREATE on
        # the *database* and is refused for lacking it even when the schema is already
        # there, which is the usual grant for a tenant that owns its own schema and nothing
        # else. `IF NOT EXISTS` is still on for two processes bootstrapping the same tenant.
        try:
            conn.execute(CreateSchema(name, if_not_exists=True))
            conn.commit()
        except ProgrammingError as e:
            raise RuntimeError(
                f"Schema '{name}' does not exist and could not be created"
            ) from e

    def _check_db_srid() -> None:
        """Raises an exception if the DB SRID differs from the configured one."""
        # unqualified, and deliberately not in `CORETABLE_SCHEMA`: PostGIS installs
        # `geometry_columns` wherever the extension went, which is not a tenant's schema
        geometry_columns = Table(
            "geometry_columns",
            MetaData(),
            Column("f_table_name"),
            Column("srid"),
            Column("f_table_schema"),
        )
        stmt = select(geometry_columns.c.srid).where(
            geometry_columns.c.f_table_name == "coretable",
            geometry_columns.c.f_table_schema == schema,
        )
        srid = conn.execute(stmt).scalar_one()
        if srid != settings.db_srid:
            raise RuntimeError(
                f"DB SRID '{srid}' and configured SRID '{settings.db_srid}' must identical"
            )

    # test for tables and revision
    with engine.connect() as conn:
        schema = _resolve_schema(conn, schema)
        alembic_cfg = _alembic_config(schema)
        script_dir = script.ScriptDirectory.from_config(alembic_cfg)
        current_version = script_dir.get_heads()
        _create_schema_if_missing(schema)
        _check_schema_accessibility("USAGE")
        inspector = inspect(conn)
        tables = inspector.get_table_names(schema=schema)
        # a partial schema counts as existing data: `refs` alone missing used to
        # route a populated database into the creation branch below
        existing_coretable = {"coretable", "refs"} & set(tables)
        if "alembic_version" in tables:
            alembic_table = Table(
                "alembic_version",
                MetaData(schema=schema),
                Column("version_num"),
            )
            stmt = select(alembic_table.c.version_num)
            db_version = conn.execute(stmt).scalars().all()
        else:
            db_version = []
        _refuse_unknown_revisions(engine, script_dir, schema, db_version)
        if db_version:
            _check_db_srid()
        is_current_version = set(db_version) == set(current_version)
        if is_current_version:
            logger.info("Database is at current revision")
            return
        if bootstrap:
            _check_schema_accessibility("CREATE")
    # handle schema upgrade or table creation
    if existing_coretable and not db_version:
        e = RuntimeError("Coretable with no revision found in database")
        e.add_note(
            "it is likely that the database was set up with an older version of this library which didn't use revisions yet"
        )
        e.add_note(
            "please set up a new database or add a revision corresponding to the current model manually"
        )
        e.add_note(f"found table(s): {', '.join(sorted(existing_coretable))}")
        if existing_coretable != {"coretable", "refs"}:
            e.add_note(
                "the schema is incomplete and is never completed automatically; note "
                "that an ogr2ogr/QGIS export of a coretable is not a coretable "
                "database - it carries no refs table and renames the geometry column"
            )
        raise e
    if not bootstrap:
        raise RuntimeError(
            f"Schema {schema!r} is not at the current revision "
            f"({', '.join(db_version) or 'empty'}); bring it to head with "
            "`xmas db create-schema` or `ensure_schema(..., bootstrap=True)`"
        )
    # bring the schema to head, creating the tables if it is empty
    logger.info(
        "Running database migrations" if db_version else "Creating new database schema"
    )
    with engine.connect() as conn:
        alembic_cfg.attributes["connection"] = conn
        command.upgrade(alembic_cfg, "head")

def for_schema(engine: Engine | AsyncEngine, schema: str) -> Engine | AsyncEngine:
    """Bind an engine to one schema, sharing its pool.

    This is how one process serves several tenants: the returned engine is a shallow view over
    the same pool, differing only in which schema `CORETABLE_SCHEMA` resolves to. The compiled
    statement is identical for every tenant - SQLAlchemy substitutes the schema after
    compilation - so the statement cache is shared and cannot leak between them.

    Args:
        engine: An engine from `create_engine` or `create_async_engine`.
        schema: The schema to address. SQLAlchemy quotes it where it substitutes it.

    Returns:
        An engine of the same kind, bound to `schema`.
    """
    return engine.execution_options(schema_translate_map={CORETABLE_SCHEMA: schema})

def add_navigable_role(
    source_featuretype: str,
    navigable_role: str,
    target_featuretype: str,
    appschema: str,
    appschema_version: str,
    rel_direction: Literal["forward", "inverse"],
    dependent_part: Literal["source", "target"] | None = None,
) -> Function[bool]:
    """Build the ``add_navigable_role`` call, which idempotently registers a navigable role.

    Args:
        source_featuretype: Owning/source feature type name.
        navigable_role: Role (association) name linking source to target.
        target_featuretype: Referenced/target feature type name.
        appschema: Appschema prefix (e.g. ``"xplan"``).
        appschema_version: Appschema version (e.g. ``"6.1"``).
        rel_direction: Whether the refs edge runs ``"forward"`` or ``"inverse"``.
        dependent_part: Ownership marker; ``"target"`` means source existentially owns target,
            ``"source"`` means target existentially owns source, ``None`` for a non-ownership role.

    Returns:
        The function's ``BOOLEAN`` result as a column element; the function is scalar, so it belongs
        in a select list rather than a ``FROM``. Execute with
        ``select(add_navigable_role(...))`` and ``scalar_one()``.
    """
    return _CoretableFunction(
        "add_navigable_role",
        source_featuretype,
        navigable_role,
        target_featuretype,
        appschema,
        appschema_version,
        rel_direction,
        dependent_part,
        type_=Boolean(),
    )