from __future__ import annotations
import logging
import re
from dataclasses import dataclass, replace
from graphlib import TopologicalSorter
from typing import Any, Literal
import pglast
from alembic.operations import ops as alembic_ops
from pglast.error import Error as PglastError
from pglast.visitors import Visitor
from sqlalchemy import MetaData # noqa: TC002
from sqlalchemy_declarative_extensions.alembic.function import (
CreateFunctionOp,
DropFunctionOp,
UpdateFunctionOp,
)
from sqlalchemy_declarative_extensions.alembic.procedure import (
CreateProcedureOp,
DropProcedureOp,
UpdateProcedureOp,
)
from sqlalchemy_declarative_extensions.alembic.schema import (
CreateSchemaOp,
DropSchemaOp,
)
from sqlalchemy_declarative_extensions.alembic.trigger import (
CreateTriggerOp,
DropTriggerOp,
UpdateTriggerOp,
)
from sqlalchemy_declarative_extensions.alembic.view import (
CreateViewOp,
DropViewOp,
UpdateViewOp,
)
from sqlalchemy_declarative_extensions.dialects.postgresql.function import (
type_map as _pg_type_map,
)
from sqlalchemy_declarative_extensions.dialects.postgresql.grant import (
DefaultGrantStatement,
GrantStatement,
)
from sqlalchemy_declarative_extensions.grant.compare import (
GrantPrivilegesOp,
RevokePrivilegesOp,
)
from sqlalchemy_declarative_extensions.op import MigrateOp
from sqlalchemy_declarative_extensions.role.compare import (
CreateRoleOp,
DropRoleOp,
)
from sqlalchemy_declarative_extensions.role.generic import Role
from codegen_database.pg_extension import CreateExtensionOp
logger = logging.getLogger(__name__)
# Union of alembic's built-in ops and sqlalchemy-declarative-extensions ops,
# which don't share a common base class.
type AnyOp = alembic_ops.MigrateOperation | MigrateOp
Phase = Literal["drop", "create"]
# Op types grouped by direction.
_CREATE_OPS = (
CreateRoleOp,
CreateSchemaOp,
alembic_ops.CreateTableOp,
alembic_ops.CreateIndexOp,
CreateViewOp,
CreateFunctionOp,
CreateProcedureOp,
CreateTriggerOp,
GrantPrivilegesOp,
)
_DROP_OPS = (
DropRoleOp,
DropSchemaOp,
alembic_ops.DropTableOp,
alembic_ops.DropIndexOp,
DropViewOp,
DropFunctionOp,
DropProcedureOp,
DropTriggerOp,
RevokePrivilegesOp,
)
_UPDATE_OPS = (
UpdateViewOp,
UpdateFunctionOp,
UpdateProcedureOp,
UpdateTriggerOp,
)
[docs]
@dataclass(frozen=True)
class EntityIdentifier:
"""Identifies a database entity within a migration's dependency graph.
``name`` is ``None`` for schema-level entities (i.e. the entity *is*
the schema). For tables, views, and functions, ``name`` holds the
unqualified object name and ``schema`` holds its containing schema.
``phase`` distinguishes drop and create ops for the same entity when
an ``Update*Op`` has been expanded. ``"drop"`` ops are ordered before
``"create"`` ops for the same entity.
"""
schema: str = "public"
name: str | None = None
phase: Phase | None = None
def _op_phase(op: AnyOp) -> Phase | None:
"""Return ``"drop"`` or ``"create"`` for ops that have a direction.
Returns ``None`` for ``ModifyTableOps`` (column adds/drops) and any
``Update*Op`` that wasn't split by ``expand_update_ops``.
"""
if isinstance(op, _DROP_OPS):
return "drop"
if isinstance(op, _CREATE_OPS):
return "create"
return None
def _entity_schema(op: AnyOp) -> str | None:
"""Extract the schema name from an entity op."""
if isinstance(op, (CreateSchemaOp, DropSchemaOp)):
return op.schema.name
if isinstance(op, _CREATE_OPS + _DROP_OPS):
for attr in ("view", "function", "procedure"):
entity = getattr(op, attr, None)
if entity is not None:
return entity.schema or "public"
trigger = getattr(op, "trigger", None)
if trigger is not None:
# PostgreSQL triggers store the target as "schema.table"
# in the `on` attribute; they have no `schema` field.
on_schema, _ = _parse_qualified_name(trigger.on)
return f"__triggers__{on_schema}"
return None
def _entity_name(op: AnyOp) -> str | None:
"""Extract the entity name from a declarative-extensions op."""
for attr in ("view", "function", "procedure"):
entity = getattr(op, attr, None)
if entity is not None:
return entity.name
trigger = getattr(op, "trigger", None)
if trigger is not None:
return trigger.name
return None
def _entity_definition(op: AnyOp) -> str | None:
"""Extract a parseable SQL definition (views only).
Function/procedure bodies contain PL/pgSQL which needs special handling;
use :func:`_plpgsql_table_refs` for those.
"""
view = getattr(op, "view", None)
if view is not None and hasattr(view, "definition"):
defn = view.definition
if isinstance(defn, str):
return defn
return None
def _plpgsql_queries(obj: object) -> list[str]:
"""Recursively collect SQL query strings from a ``parse_plpgsql`` tree."""
queries: list[str] = []
if isinstance(obj, dict):
if "PLpgSQL_expr" in obj:
query = obj["PLpgSQL_expr"].get("query", "")
# Skip trivial expressions like NEW/OLD.
if query and query.upper() not in ("NEW", "OLD"):
queries.append(query)
else:
for value in obj.values():
queries.extend(_plpgsql_queries(value))
elif isinstance(obj, list):
for item in obj:
queries.extend(_plpgsql_queries(item))
return queries
def _fn_setof_ref(func: object) -> tuple[str, str] | None:
"""Return ``(schema, name)`` if *func* has a ``SETOF schema.name`` return.
Accepts both raw strings (``"SETOF private.inventory"``) and
normalised ``FunctionReturn`` objects produced by
``Function.normalize()`` during autogenerate comparison. Returns
``None`` for non-SETOF or unqualified returns.
"""
returns = getattr(func, "returns", None)
if not returns:
return None
# SDE normalises functions before creating ``CreateFunctionOp``, so
# ``returns`` may be a ``FunctionReturn`` object rather than a plain string.
if isinstance(returns, str):
stripped = returns.strip()
elif hasattr(returns, "value") and isinstance(returns.value, str):
stripped = returns.value.strip()
else:
return None
if not stripped.upper().startswith("SETOF "):
return None
ref = stripped[6:].strip()
schema, sep, name = ref.partition(".")
if sep:
return (schema.lower(), name.lower())
return None
def _function_return_refs(op: AnyOp) -> set[tuple[str, str]]:
"""Extract ``(schema, name)`` from a function op's ``SETOF`` return type.
Delegates to :func:`_fn_setof_ref` for the actual parsing. Returns
an empty set for non-function ops or non-SETOF returns.
"""
func = getattr(op, "function", None)
if func is None:
return set()
ref = _fn_setof_ref(func)
return {ref} if ref is not None else set()
def _sql_function_body(op: AnyOp) -> str | None:
"""Return the body of a LANGUAGE sql function op, or None.
Returns ``None`` for non-function ops, functions in another
language, or functions with an empty body.
"""
func = getattr(op, "function", None)
if func is None:
return None
defn = getattr(func, "definition", None)
language = getattr(func, "language", "")
if not defn or language.lower() != "sql":
return None
return defn
def _plpgsql_function_queries(op: AnyOp) -> list[str]:
"""Return embedded SQL queries from a PL/pgSQL function op.
Returns an empty list for non-function ops, functions in another
language, or if parsing fails.
"""
func = getattr(op, "function", None)
if func is None:
return []
defn = getattr(func, "definition", None)
language = getattr(func, "language", "")
if not defn or language.lower() != "plpgsql":
return []
# pglast.parse_plpgsql requires a full CREATE FUNCTION statement.
schema_part = f"{func.schema}." if func.schema else ""
wrapper = (
f"CREATE FUNCTION {schema_part}__cave_parse_helper()"
f" RETURNS trigger LANGUAGE plpgsql AS $${defn}$$;"
)
try:
tree = pglast.parse_plpgsql(wrapper)
except PglastError:
logger.debug(
"Could not parse PL/pgSQL body for %s",
func.name,
)
return []
return _plpgsql_queries(tree)
def _sql_function_table_refs(
op: AnyOp,
) -> set[tuple[str, str]]:
"""Extract ``(schema, table)`` pairs from a LANGUAGE sql function.
Uses ``pglast.parse_sql`` to find table references in the
function body.
Returns an empty set for non-sql-language ops or if parsing
fails.
"""
defn = _sql_function_body(op)
if defn is None:
return set()
return _view_table_refs(defn)
def _plpgsql_table_refs(
op: AnyOp,
) -> set[tuple[str, str]]:
"""Extract ``(schema, table)`` pairs from a PL/pgSQL function body.
Uses ``pglast.parse_plpgsql`` to extract embedded SQL statements from
the function body, then ``pglast.parse_sql`` to find table references
within those statements.
Returns an empty set for non-function ops or if parsing fails.
"""
refs: set[tuple[str, str]] = set()
for query in _plpgsql_function_queries(op):
refs |= _view_table_refs(query)
return refs
def _role_name(member: Role | str) -> str:
"""Extract a role name from a Role object or string."""
if isinstance(member, Role):
return member.name
return member
def _id_for_declarative_op(
op: AnyOp,
phase: Phase | None,
) -> EntityIdentifier | None:
"""Return identifier for view/function/procedure/trigger ops."""
name = _entity_name(op)
if name is not None:
schema = _entity_schema(op) or "public"
return EntityIdentifier(
schema=schema.lower(),
name=name.lower(),
phase=phase,
)
return None
def _entity_identifier( # noqa: PLR0911
op: AnyOp,
) -> EntityIdentifier | None:
"""Return the identifier for the entity this op acts on.
Returns ``None`` for unrecognised op types (logged as a warning).
"""
phase = _op_phase(op)
if isinstance(op, (CreateSchemaOp, DropSchemaOp)):
return EntityIdentifier(
schema=op.schema.name.lower(),
phase=phase,
)
if isinstance(op, (alembic_ops.CreateTableOp, alembic_ops.DropTableOp)):
return EntityIdentifier(
schema=(op.schema or "public").lower(),
name=op.table_name.lower(),
phase=phase,
)
if isinstance(op, alembic_ops.ModifyTableOps):
return EntityIdentifier(
schema=(op.schema or "public").lower(),
name=op.table_name.lower(),
)
if isinstance(op, (alembic_ops.CreateIndexOp, alembic_ops.DropIndexOp)):
return EntityIdentifier(
schema=(op.schema or "public").lower(),
name=f"__index__{(op.index_name or '').lower()}",
phase=phase,
)
if isinstance(op, (CreateRoleOp, DropRoleOp)):
return EntityIdentifier(
schema="__roles__",
name=op.role.name.lower(),
phase=phase,
)
if isinstance(op, (GrantPrivilegesOp, RevokePrivilegesOp)):
return EntityIdentifier(
schema="__grants__",
name=str(op.to_sql()).lower(),
phase=phase,
)
result = _id_for_declarative_op(op, phase)
if result is not None:
return result
logger.warning(
"Unhandled op type %s; ordering is unconstrained",
type(op).__name__,
)
return None
def _refs_for_role(
op: CreateRoleOp | DropRoleOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Return references for a role op.
Picks up: parent roles declared via ``in_roles`` (e.g.
``authenticator`` inheriting ``anon``). The parent role
must exist before the child role can be created.
"""
if not op.role.in_roles:
return set()
return {
EntityIdentifier(
schema="__roles__",
name=_role_name(member).lower(),
phase=phase,
)
for member in op.role.in_roles
}
def _refs_for_grant(
op: GrantPrivilegesOp | RevokePrivilegesOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Return references for a grant/revoke op.
Picks up:
- The target role the privilege is granted to.
- For ``GrantStatement``: the schema and object being granted
on (e.g. ``GRANT SELECT ON reporting.students TO analyst``
depends on the ``reporting`` schema and the ``students`` view).
- For ``DefaultGrantStatement``: the schemas the default
grant applies to.
"""
grant_obj = op.grant
refs: set[EntityIdentifier] = set()
# Depend on the target role.
refs.add(
EntityIdentifier(
schema="__roles__",
name=grant_obj.grant.target_role.lower(),
phase=phase,
)
)
# Depend on referenced schemas and target objects.
if isinstance(grant_obj, GrantStatement):
for target_name in grant_obj.targets:
schema_part, sep, obj_name = target_name.partition(".")
if sep:
refs.add(
EntityIdentifier(
schema=schema_part.lower(),
phase=phase,
)
)
refs.add(
EntityIdentifier(
schema=schema_part.lower(),
name=obj_name.lower(),
phase=phase,
)
)
else:
refs.add(
EntityIdentifier(
schema=target_name.lower(),
phase=phase,
)
)
elif isinstance(grant_obj, DefaultGrantStatement):
for schema_name in grant_obj.default_grant.in_schemas:
refs.add(
EntityIdentifier(
schema=schema_name.lower(),
phase=phase,
)
)
return refs
def _parse_qualified_name(
qualified: str,
) -> tuple[str, str]:
"""Split ``schema.name`` into ``(schema, name)``."""
parts = qualified.split(".", 1)
if len(parts) > 1:
return parts[0].lower(), parts[1].lower()
return "public", parts[0].lower()
def _refs_for_trigger(
op: AnyOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Return refs for trigger ops.
Picks up:
- The target table/view the trigger is attached to
(``ON schema.table``).
- The function the trigger executes (``EXECUTE
FUNCTION schema.func()``).
"""
trigger = getattr(op, "trigger", None)
if trigger is None:
return set()
refs: set[EntityIdentifier] = set()
# Depend on the target view/table (``on``).
on_schema, on_name = _parse_qualified_name(trigger.on)
refs.add(
EntityIdentifier(
schema=on_schema,
name=on_name,
phase=phase,
)
)
refs.add(EntityIdentifier(schema=on_schema, phase=phase))
# Depend on the executed function.
fn_schema, fn_name = _parse_qualified_name(trigger.execute)
refs.add(
EntityIdentifier(
schema=fn_schema,
name=fn_name,
phase=phase,
)
)
return refs
def _add_table_refs(
refs: set[EntityIdentifier],
table_pairs: set[tuple[str, str]],
phase: Phase | None,
) -> None:
"""Add ``EntityIdentifier`` entries for ``(schema, name)`` pairs."""
for ref_schema, ref_name in table_pairs:
refs.add(
EntityIdentifier(
schema=ref_schema,
name=ref_name,
phase=phase,
)
)
if phase is not None:
refs.add(
EntityIdentifier(
schema=ref_schema,
name=ref_name,
)
)
def _view_table_refs(definition: str) -> set[tuple[str, str]]:
"""Extract ``(schema, table)`` pairs from a view SQL definition."""
refs: set[tuple[str, str]] = set()
try:
parsed = pglast.parse_sql(definition)
class _TableFinder(Visitor):
def visit_RangeVar( # noqa: N802
self,
_ancestors: object,
node: object,
) -> None:
if name := getattr(node, "relname", None):
schema = getattr(node, "schemaname", "public") or "public"
refs.add((schema.lower(), name.lower()))
_TableFinder()(parsed)
except PglastError:
logger.debug(
"Could not parse view definition: %s",
definition[:80],
)
return refs
def _sql_function_function_refs(op: AnyOp) -> set[tuple[str, str]]:
"""Extract ``(schema, name)`` function calls from a LANGUAGE sql function.
Returns an empty set for non-sql-language ops or if parsing fails.
"""
defn = _sql_function_body(op)
if defn is None:
return set()
return _view_function_refs(defn)
def _plpgsql_function_refs(op: AnyOp) -> set[tuple[str, str]]:
"""Extract ``(schema, name)`` function calls from a PL/pgSQL function body.
Returns an empty set for non-function ops or if parsing fails.
"""
queries = _plpgsql_function_queries(op)
if not queries:
return set()
refs: set[tuple[str, str]] = set()
for query in queries:
refs |= _view_function_refs(query)
return refs
def _refs_from_definitions(
op: AnyOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Extract entity refs from view SQL and function bodies/returns."""
refs: set[EntityIdentifier] = set()
# View definitions: parse with pglast.
definition = _entity_definition(op)
if definition is not None:
_add_table_refs(refs, _view_table_refs(definition), phase)
# Functions called from a view body.
_add_table_refs(refs, _view_function_refs(definition), phase)
# PL/pgSQL function bodies: parse with pglast.
_add_table_refs(refs, _plpgsql_table_refs(op), phase)
_add_table_refs(refs, _plpgsql_function_refs(op), phase)
# LANGUAGE sql function bodies: parse with pglast.
_add_table_refs(refs, _sql_function_table_refs(op), phase)
_add_table_refs(refs, _sql_function_function_refs(op), phase)
# Function return types: SETOF schema.view.
_add_table_refs(refs, _function_return_refs(op), phase)
return refs
def _refs_for_declarative_op(
op: AnyOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Return references for view/function/procedure/trigger ops.
Picks up:
- Triggers: delegated to ``_refs_for_trigger`` (target
table + executed function).
- For update ops that were split into drop+create: the
create phase depends on the drop phase of the same
entity (drop old definition before creating new one).
- The containing schema (view/function must be created
after its schema).
- Tables/views referenced in the SQL body, extracted by
parsing view definitions, PL/pgSQL bodies, SQL function
bodies, and ``SETOF`` return types.
"""
# Triggers have their own reference logic via the `on` field.
trigger_refs = _refs_for_trigger(op, phase)
if trigger_refs:
return trigger_refs
name = _entity_name(op)
if name is None:
return set()
schema = (_entity_schema(op) or "public").lower()
self_id = EntityIdentifier(
schema=schema,
name=name.lower(),
phase=phase,
)
refs: set[EntityIdentifier] = set()
# For split update ops: create must wait for drop.
if phase == "create":
refs.add(
EntityIdentifier(
schema=schema,
name=name.lower(),
phase="drop",
)
)
# Depend on the containing schema: created after its schema on
# create, and dropped before its schema on drop (the drop-phase
# edge direction reverses this so DROP SCHEMA waits for the
# function -- e.g. the ``codegen_database`` utility schema can't be
# dropped while its ``codegen_database_date_bin`` chart polyfill
# still lives in it).
refs.add(EntityIdentifier(schema=schema, phase=phase))
# Tables/views referenced in SQL definitions and bodies.
refs |= _refs_from_definitions(op, phase)
refs.discard(self_id)
return refs
def _entity_references(
op: AnyOp,
) -> set[EntityIdentifier]:
"""Return identifiers of entities this op references.
The sort loop decides edge direction based on the op's phase:
create-phase ops need their references to exist first (normal edge),
drop-phase ops need their dependents dropped first (reversed edge).
Only identifiers present among the current migration's ops will
produce dependency edges; references to already-existing entities
are filtered out in ``sort_migration_ops``.
"""
phase = _op_phase(op)
# Tables depend on their containing schema. FK deps are
# handled separately via fk_graph in sort_migration_ops.
if isinstance(op, (alembic_ops.CreateTableOp, alembic_ops.DropTableOp)):
return {
EntityIdentifier(
schema=(op.schema or "public").lower(),
phase=phase,
)
}
# Roles depend on parent roles (in_roles).
if isinstance(op, (CreateRoleOp, DropRoleOp)):
return _refs_for_role(op, phase)
# Grants depend on the target role, schema, and object.
if isinstance(op, (GrantPrivilegesOp, RevokePrivilegesOp)):
return _refs_for_grant(op, phase)
# Views, functions, procedures, triggers: schema + SQL refs.
return _refs_for_declarative_op(op, phase)
def _op_label(op: AnyOp) -> str:
"""Return a compact label for an op, used in log messages.
Format is ``OpType(schema.name)`` e.g. ``CreateViewOp(private.students)``,
or ``OpType(schema)`` for schema-level ops.
"""
identifier = _entity_identifier(op)
if identifier is None:
entity = "?"
elif identifier.name is None:
entity = identifier.schema
else:
entity = f"{identifier.schema}.{identifier.name}"
return f"{type(op).__name__}({entity})"
_QUALIFIED_FUNCNAME_MIN_PARTS = 2
_KIND_FUNCTION = "function"
_KIND_VIEW = "view"
def _view_function_refs(definition: str) -> set[tuple[str, str]]:
"""Extract schema-qualified ``(schema, name)`` function calls from SQL.
Uses ``pglast.parse_sql`` to walk ``FuncCall`` nodes. Only calls
with an explicit schema qualifier are returned (unqualified calls
cannot be resolved unambiguously).
"""
refs: set[tuple[str, str]] = set()
try:
parsed = pglast.parse_sql(definition)
class _FuncFinder(Visitor):
def visit_FuncCall( # noqa: N802
self,
_ancestors: object,
node: object,
) -> None:
funcname = getattr(node, "funcname", None)
if funcname is None:
return
if len(funcname) < _QUALIFIED_FUNCNAME_MIN_PARTS:
return
schema = getattr(funcname[-2], "sval", None)
name = getattr(funcname[-1], "sval", None)
if schema and name:
refs.add((schema.lower(), name.lower()))
_FuncFinder()(parsed)
except PglastError:
logger.debug(
"Could not parse view definition for function refs: %s",
definition[:80],
)
return refs
def _updated_function_names(
migration_ops: list[AnyOp],
) -> set[tuple[str, str]]:
"""Return ``(schema, name)`` for every function/procedure being updated."""
updated: set[tuple[str, str]] = set()
for op in migration_ops:
if isinstance(op, (UpdateFunctionOp, UpdateProcedureOp)):
fn = getattr(op, "function", None) or getattr(op, "procedure", None)
if fn is not None:
updated.add(((fn.schema or "public").lower(), fn.name.lower()))
return updated
def _existing_trigger_keys(
migration_ops: list[AnyOp],
) -> set[tuple[str, str]]:
"""Return ``(schema, trigger_name)`` for triggers already in the op list."""
keys: set[tuple[str, str]] = set()
for op in migration_ops:
trigger = getattr(op, "trigger", None)
if trigger is not None:
on_schema, _ = _parse_qualified_name(trigger.on)
keys.add((on_schema, trigger.name.lower()))
return keys
def _existing_view_keys(migration_ops: list[AnyOp]) -> set[tuple[str, str]]:
"""Return ``(schema, view_name)`` for views already in the op list."""
keys: set[tuple[str, str]] = set()
for op in migration_ops:
view = getattr(op, "view", None)
if view is not None:
keys.add(((view.schema or "public").lower(), view.name.lower()))
return keys
def _updated_view_names(
migration_ops: list[AnyOp],
) -> set[tuple[str, str]]:
"""``(schema, name)`` of every view being recreated (``UpdateViewOp``)."""
updated: set[tuple[str, str]] = set()
for op in migration_ops:
if isinstance(op, UpdateViewOp):
view = op.view
updated.add(((view.schema or "public").lower(), view.name.lower()))
return updated
def _column_altered_table_keys(
migration_ops: list[AnyOp],
) -> set[tuple[str, str]]:
"""``(schema, table)`` for tables with a column TYPE change.
Postgres rejects ``ALTER COLUMN ... TYPE`` while a view depends on the
column, so dependent views must drop/recreate around it.
"""
def has_type_change(op: alembic_ops.ModifyTableOps) -> bool:
return any(
isinstance(sub, alembic_ops.AlterColumnOp)
and sub.modify_type is not None
for sub in op.ops
)
return {
((op.schema or "public").lower(), op.table_name.lower())
for op in migration_ops
if isinstance(op, alembic_ops.ModifyTableOps) and has_type_change(op)
}
def _canonical_function_body(definition: str) -> str:
"""Function body with per-line whitespace stripped, for equality only.
Alembic indents the nested op text, so the stored/reflected body is always
re-indented vs the model definition (and ``BEGIN`` sticks to ``AS $$``, so
``dedent`` can't undo it). Trailing whitespace also drifts: a body
compiled from a SQLAlchemy ``select()`` (e.g. a ledger chart function)
carries end-of-line spaces the created/reflected form drops. SQL /
PL/pgSQL is whitespace-insensitive, so strip both ends of every line and
compare loosely; the op keeps its real definition for rendering.
"""
return "\n".join(line.strip() for line in definition.splitlines()).strip()
def _canonical_type(type_: Any) -> str: # noqa: ANN401
"""Type without modifiers / default schema, aliased -- for equality only.
Postgres normalizes function parameter and return types on storage: it
drops length / precision modifiers (``varchar(500)`` -> ``varchar``,
``numeric(10, 2)`` -> ``numeric``) and omits the default ``public.``
schema (``setof public.inventory_ledger`` -> ``setof inventory_ledger``).
Normalize both sides the same way so a function isn't flagged as changed
on every regenerate.
"""
canon = re.sub(r"\s*\([^)]*\)", "", str(type_)).strip().lower()
canon = re.sub(r"\bpublic\.", "", canon)
return _pg_type_map.get(canon, canon)
def _param_identities(parameters: Any) -> Any: # noqa: ANN401
"""Parameters as ``(name, type, mode)``, for equality only.
Drops the default text, which Postgres rewrites on storage and can't
round-trip, e.g. ``'2000-01-03'::timestamptz`` ->
``'2000-01-03 00:00:00-05'::timestamp with time zone``. Canonicalizes
the type too, since Postgres strips length modifiers on parameters
(``varchar(500)`` -> ``varchar``).
"""
if not parameters:
return parameters
return [
p
if isinstance(p, str)
else (
p.name.lower() if p.name is not None else None,
_canonical_type(p.type),
p.mode,
)
for p in parameters
]
_TABLE_RETURN_RE = re.compile(r"(?is)^\s*table\s*\((.*)\)\s*$")
def _split_top_level(cols: str) -> list[str]:
"""Split a comma list, ignoring commas inside parens (``numeric(10,2)``)."""
parts: list[str] = []
depth = 0
start = 0
for i, ch in enumerate(cols):
if ch == "(":
depth += 1
elif ch == ")":
depth -= 1
elif ch == "," and depth == 0:
parts.append(cols[start:i])
start = i + 1
parts.append(cols[start:])
return [p.strip() for p in parts if p.strip()]
def _return_identity(returns: Any) -> Any: # noqa: ANN401
"""Canonical identity for a function return type, for equality only.
Postgres round-trips ``RETURNS TABLE (...)`` into a structured column
list (``FunctionReturn(table=[(name, type), ...])``), while the model
carries it as a raw ``"table (...)"`` string in ``FunctionReturn.value``
(SDE only parses the space-less ``table(`` form with table-mode params).
Normalize both to ``("table", ((name, type), ...))`` -- with SDE's type
aliases (``numeric`` -> ``decimal``) -- so a chart function isn't flagged
as a spurious update on every regenerate.
"""
if returns is None:
return returns
table = getattr(returns, "table", None)
value = getattr(returns, "value", None)
if not table and isinstance(value, str):
match = _TABLE_RETURN_RE.match(value)
if match:
table = [
col.split(maxsplit=1)
for col in _split_top_level(match.group(1))
]
if table:
return ("table", tuple(_col_identity(col) for col in table))
if isinstance(value, str):
return ("scalar", _canonical_type(value))
return returns
def _col_identity(col: Any) -> tuple[str, str]: # noqa: ANN401
"""``(name, type)`` for one table-return column, types aliased."""
if isinstance(col, str):
name, _, type_ = col.strip().partition(" ")
else:
name, type_ = col[0], col[1]
return str(name).lower(), _canonical_type(type_)
def _qualify_default_schema(name: str | None) -> str | None:
"""Qualify an unqualified name with ``public`` (reflection drops it).
e.g. ``users_raw`` -> ``public.users_raw``; ``private.x`` is left as-is.
"""
if name is not None and "." not in name:
return f"public.{name}"
return name
def _function_update_is_spurious(op: UpdateFunctionOp) -> bool:
"""True when a function Update is only schema/whitespace/default noise."""
def canon(fn: Any) -> Any: # noqa: ANN401
return replace(
fn,
schema=fn.schema or "public",
definition=_canonical_function_body(fn.definition),
parameters=_param_identities(fn.parameters),
returns=_return_identity(fn.returns),
)
return canon(op.from_function) == canon(op.function)
def _view_update_is_spurious(op: UpdateViewOp) -> bool:
"""True when a view Update differs only by the default schema."""
def canon(view: Any) -> Any: # noqa: ANN401
return replace(view, schema=view.schema or "public")
return canon(op.from_view) == canon(op.view)
def _trigger_update_is_spurious(op: UpdateTriggerOp) -> bool:
"""True when a trigger Update differs only by ``public.`` qualification.
Triggers have no ``schema`` field; it lives in the ``on``/``execute`` names.
"""
def canon(trigger: Any) -> Any: # noqa: ANN401
return replace(
trigger,
on=_qualify_default_schema(trigger.on),
execute=_qualify_default_schema(trigger.execute),
)
return canon(op.from_trigger) == canon(op.trigger)
[docs]
def drop_spurious_declarative_updates(ops: list[AnyOp]) -> list[AnyOp]:
"""Drop no-op function/view/trigger ``Update`` ops (default-schema and
renderer-whitespace churn). Real changes survive; dependent recreation a
real change or column alter needs is re-added by
``_inject_dependent_update_ops``.
"""
return [
op
for op in ops
if not (
(
isinstance(op, UpdateFunctionOp)
and _function_update_is_spurious(op)
)
or (isinstance(op, UpdateViewOp) and _view_update_is_spurious(op))
or (
isinstance(op, UpdateTriggerOp)
and _trigger_update_is_spurious(op)
)
)
]
def _inject_dependent_update_ops( # noqa: C901, PLR0915
migration_ops: list[AnyOp],
metadata: MetaData | None,
) -> list[AnyOp]:
"""Inject update ops for all transitive dependents of updated functions.
When a function body changes, SDE only emits ``UpdateFunctionOp``
(or ``UpdateProcedureOp``). Because ``expand_update_ops`` will
split these into ``DropFunctionOp`` + ``CreateFunctionOp``,
PostgreSQL will refuse the drop if anything still references the
function.
This function performs a BFS over the dependency graph starting from
the updated functions. At each step it looks up triggers, views, and
functions that depend on the current node and injects self-referential
update ops for any not already in the op list:
- *function* → dependent triggers (``EXECUTE FUNCTION``) and views
(schema-qualified call in SQL body) are enqueued.
- *view* → dependent views (referenced in ``FROM``/``JOIN``) and
functions (``SETOF view``) are enqueued.
The injected ops are expanded by ``expand_update_ops`` in the normal
way so the topological sort can interleave drops before the function
drop and creates after the function create.
Args:
migration_ops: Current list of ops (may already contain some
trigger/view/function update ops).
metadata: SQLAlchemy ``MetaData`` holding SDE-registered objects;
pass ``None`` to skip injection.
Returns:
A new list with the injected ops appended.
"""
if metadata is None:
return list(migration_ops)
seen_fns = _updated_function_names(migration_ops)
updated_views = _updated_view_names(migration_ops)
altered_tables = _column_altered_table_keys(migration_ops)
if not (seen_fns or updated_views or altered_tables):
return list(migration_ops)
all_views = list(metadata.info.get("views") or [])
all_fns = list(metadata.info.get("functions") or [])
views_by_key: dict[tuple[str, str], Any] = {
((v.schema or "public").lower(), v.name.lower()): v for v in all_views
}
fns_by_key: dict[tuple[str, str], Any] = {
((f.schema or "public").lower(), f.name.lower()): f for f in all_fns
}
all_triggers = list(metadata.info.get("triggers") or [])
view_fn_refs = {
k: _view_function_refs(getattr(v, "definition", "") or "")
for k, v in views_by_key.items()
}
view_tbl_refs = {
k: _view_table_refs(getattr(v, "definition", "") or "")
for k, v in views_by_key.items()
}
fn_setof_refs = {
k: ref
for k, fn in fns_by_key.items()
if (ref := _fn_setof_ref(fn)) is not None
}
seen_triggers = _existing_trigger_keys(migration_ops)
seen_views = _existing_view_keys(migration_ops)
extra: list[AnyOp] = []
queue: list[tuple[str, str, str]] = [
(s, n, _KIND_FUNCTION) for s, n in seen_fns
]
# Recreated views drop their INSTEAD OF triggers; revisit to re-add them.
queue += [(s, n, _KIND_VIEW) for s, n in updated_views]
visited: set[tuple[str, str, str]] = set(queue)
def _enqueue_view(vkey: tuple[str, str]) -> None:
seen_views.add(vkey)
view = views_by_key[vkey]
extra.append(UpdateViewOp(view, view))
entry = (vkey[0], vkey[1], _KIND_VIEW)
if entry not in visited:
visited.add(entry)
queue.append(entry)
def _enqueue_trigger(trigger: Any) -> None: # noqa: ANN401
"""Enqueue an ``UpdateTriggerOp`` for *trigger* if not already seen."""
on_s, _ = _parse_qualified_name(trigger.on)
tkey = (on_s, trigger.name.lower())
if tkey not in seen_triggers:
seen_triggers.add(tkey)
extra.append(UpdateTriggerOp(trigger, trigger))
def _enqueue_fn(fn_key: tuple[str, str]) -> None:
seen_fns.add(fn_key)
fn = fns_by_key[fn_key]
extra.append(UpdateFunctionOp(fn, fn))
entry = (fn_key[0], fn_key[1], _KIND_FUNCTION)
if entry not in visited:
visited.add(entry)
queue.append(entry)
def _visit_fn(fn_key: tuple[str, str]) -> None:
"""Enqueue triggers and views that depend on *fn_key*."""
for trigger in all_triggers:
if _parse_qualified_name(trigger.execute) == fn_key:
_enqueue_trigger(trigger)
for vkey, fn_refs in view_fn_refs.items():
if fn_key in fn_refs and vkey not in seen_views:
_enqueue_view(vkey)
def _visit_view(view_key: tuple[str, str]) -> None:
"""Enqueue triggers, views and functions that depend on *view_key*."""
# INSTEAD OF triggers drop with the view; recreate them.
for trigger in all_triggers:
if _parse_qualified_name(trigger.on) == view_key:
_enqueue_trigger(trigger)
for vkey, tbl_refs in view_tbl_refs.items():
if view_key in tbl_refs and vkey not in seen_views:
_enqueue_view(vkey)
for fn_key, ret_ref in fn_setof_refs.items():
if ret_ref == view_key and fn_key not in seen_fns:
_enqueue_fn(fn_key)
# A column TYPE change forces dependent views to drop/recreate around it.
for tkey in altered_tables:
for vkey, tbl_refs in view_tbl_refs.items():
if tkey in tbl_refs and vkey not in seen_views:
_enqueue_view(vkey)
while queue:
schema, name, kind = queue.pop(0)
if kind == _KIND_FUNCTION:
_visit_fn((schema, name))
else:
_visit_view((schema, name))
return list(migration_ops) + extra
[docs]
def expand_update_ops(
migration_ops: list[AnyOp],
) -> list[AnyOp]:
"""Split ``Update*Op`` into ``Drop*Op`` + ``Create*Op``.
All update op types are split so the topological sort can interleave
drops and creates correctly:
- **Views**: ``CREATE OR REPLACE VIEW`` fails when a dependent view
has an incompatible column list. Splitting lets dependents be
dropped before the dependency is dropped and recreated.
- **Functions/procedures**: ``DROP FUNCTION`` fails when dependent
triggers or views still reference the function. Splitting lets
the sort order drops before the function drop and creates after the
function create. ``_inject_dependent_update_ops`` ensures
that dependent objects are present in the op list before this
expansion runs.
- **Triggers**: ``CREATE OR REPLACE TRIGGER`` is not available in
the PostgreSQL versions we target.
"""
result: list[AnyOp] = []
for op in migration_ops:
if isinstance(op, UpdateViewOp):
result.append(DropViewOp(op.from_view))
result.append(CreateViewOp(op.view))
elif isinstance(op, UpdateFunctionOp):
result.append(DropFunctionOp(op.from_function))
result.append(CreateFunctionOp(op.function))
elif isinstance(op, UpdateProcedureOp):
result.append(DropProcedureOp(op.from_procedure))
result.append(CreateProcedureOp(op.procedure))
elif isinstance(op, UpdateTriggerOp):
result.append(DropTriggerOp(op.from_trigger))
result.append(CreateTriggerOp(op.trigger))
else:
result.append(op)
return result
def _partition_extension_ops(
ops: list[AnyOp],
) -> tuple[list[AnyOp], list[AnyOp]]:
"""Split *ops* into ``(extension_ops, rest)``.
``CreateExtensionOp`` instances must always run first because they
install operator classes (e.g. ``btree_gist``) that ``CREATE TABLE``
statements may depend on.
"""
ext: list[AnyOp] = []
rest: list[AnyOp] = []
for op in ops:
if isinstance(op, CreateExtensionOp):
ext.append(op)
else:
rest.append(op)
return ext, rest
def _index_and_modify_refs(
op: AnyOp,
phase: Phase | None,
) -> set[EntityIdentifier]:
"""Edges from index ops + ``ModifyTableOps`` to their parent table.
- Create-index / drop-index point at the corresponding CREATE-
/ DROP-phase table so the table is created first / dropped
last.
- ``ModifyTableOps`` (the wrapper autogenerate emits for new-
table batches like CreateIndexOp) always pins to the
CREATE-phase table id so the wrapper runs after the
CreateTableOp.
"""
if isinstance(
op,
(alembic_ops.CreateIndexOp, alembic_ops.DropIndexOp),
):
return {
EntityIdentifier(
schema=(op.schema or "public").lower(),
name=(op.table_name or "").lower(),
phase=phase,
)
}
if isinstance(op, alembic_ops.ModifyTableOps):
return {
EntityIdentifier(
schema=(op.schema or "public").lower(),
name=op.table_name.lower(),
phase="create",
)
}
return set()
[docs]
def prune_redundant_index_drops[T: AnyOp](ops: list[T]) -> list[T]:
"""Drop redundant ``DROP INDEX`` ops for tables dropped in this migration.
Postgres drops a table's indexes when the table is dropped, so a
separate ``DROP INDEX`` for a table that also has a ``DropTableOp`` in
the same migration is redundant -- and, when it lands after the table
drop (e.g. wrapped in a ``ModifyTableOps`` that pins to the create
phase), fails outright with "index does not exist". Strip such index
drops, whether top-level or nested inside a ``ModifyTableOps`` (drop
the wrapper if it empties out).
"""
dropped_tables = {
((op.schema or "public").lower(), op.table_name.lower())
for op in ops
if isinstance(op, alembic_ops.DropTableOp)
}
if not dropped_tables:
return ops
def _index_table_dropped(op: alembic_ops.DropIndexOp) -> bool:
return (
(op.schema or "public").lower(),
(op.table_name or "").lower(),
) in dropped_tables
result: list[T] = []
for op in ops:
if isinstance(op, alembic_ops.DropIndexOp) and _index_table_dropped(op):
continue
if isinstance(op, alembic_ops.ModifyTableOps):
op.ops[:] = [
sub
for sub in op.ops
if not (
isinstance(sub, alembic_ops.DropIndexOp)
and _index_table_dropped(sub)
)
]
if not op.ops:
continue
result.append(op)
return result
[docs]
def sort_migration_ops(
migration_ops: list[AnyOp],
*,
fk_graph: dict[tuple[str, str], set[tuple[str, str]]],
) -> list[AnyOp]:
"""Return *migration_ops* topologically sorted by entity dependencies.
Dependency edges are derived from the ops themselves:
- A table depends on its schema.
- A table depends on tables it references via foreign keys.
- A replaceable entity (view, function, ...) depends on its schema
and on every schema-qualified table or view referenced in its SQL
definition.
Only dependencies between ops in the current migration produce edges;
references to already-existing objects are ignored.
Edge direction is determined per-op by phase: drop-phase ops reverse
their edges (dependents dropped first), create-phase ops use normal
direction (dependencies created first).
:param migration_ops: Operations from a single migration script.
:param fk_graph: FK dependency map (table -> referenced tables),
built from :func:`build_fk_graph_from_metadata`.
:returns: A new list containing the same ops in dependency order.
"""
logger.debug(
"Sorting %d ops: %s",
len(migration_ops),
[_op_label(op) for op in migration_ops],
)
# Extensions have no dependencies and must precede everything else
# (they provide operator classes etc. required by CREATE TABLE).
extension_ops, migration_ops = _partition_extension_ops(migration_ops)
# Build a mapping from entity identifier to op. Ops without a
# recognised identifier are appended at the end in original order.
op_by_entity: dict[EntityIdentifier, AnyOp] = {}
unkeyed_ops: list[AnyOp] = []
for op in migration_ops:
entity = _entity_identifier(op)
if entity is not None:
op_by_entity[entity] = op
else:
unkeyed_ops.append(op)
# Use EntityIdentifier (frozen dataclass, hashable) as graph nodes.
sorter: TopologicalSorter[EntityIdentifier] = TopologicalSorter()
for current_id, current_op in op_by_entity.items():
sorter.add(current_id)
phase = _op_phase(current_op)
refs = _entity_references(current_op)
# Add FK-based refs for both create and drop table ops.
# The fk_graph (built from metadata) is used because
# neither op type carries FK info at this point.
if isinstance(
current_op,
(
alembic_ops.CreateTableOp,
alembic_ops.DropTableOp,
),
):
table_key = (
(current_op.schema or "public").lower(),
current_op.table_name.lower(),
)
for ref_schema, ref_table in fk_graph.get(table_key, set()):
refs.add(
EntityIdentifier(
schema=ref_schema,
name=ref_table,
phase=phase,
)
)
refs |= _index_and_modify_refs(current_op, phase)
# Replace-in-place across object kinds: when this migration both
# drops and creates a relation under the same (schema, name) --
# e.g. a dimension flipping from a view-backed model to a
# passthrough table -- the create must wait for the drop, or
# Postgres rejects CREATE with "relation already exists".
# ``_refs_for_declarative_op`` adds this drop->create edge only
# for replaceable (view/function) create ops; generalise it to
# any create op (notably ``CreateTableOp``) whose name a drop op
# in this same migration also targets.
if phase == "create" and current_id.name is not None:
refs.add(
EntityIdentifier(
schema=current_id.schema,
name=current_id.name,
phase="drop",
)
)
for ref_id in refs:
# Only add edges for references that resolve to an op in
# this migration. Unresolved references point to entities
# that already exist in the database.
if ref_id in op_by_entity:
# Drop-phase: reverse the edge — dependents must be
# dropped first. Create-phase/unphased: normal
# direction — the referenced entity must exist before
# we use it.
if phase == "drop":
node, prerequisite = ref_id, current_id
else:
node, prerequisite = current_id, ref_id
logger.debug(
"Edge: %s before %s",
_op_label(op_by_entity[prerequisite]),
_op_label(op_by_entity[node]),
)
sorter.add(node, prerequisite)
sorted_ops = (
extension_ops
+ [op_by_entity[eid] for eid in sorter.static_order()]
+ unkeyed_ops
)
logger.debug(
"Sorted order: %s",
[_op_label(op) for op in sorted_ops],
)
return sorted_ops