"""Custom alembic renderers that format SQL with pglast."""
from textwrap import indent
from typing import TYPE_CHECKING, Any, Literal
from alembic.autogenerate.render import _repr_type, renderers
from sqlalchemy import types as sa_types
from sqlalchemy.sql.ddl import CreateSchema, DropSchema
from sqlalchemy_declarative_extensions.function.compare import (
CreateFunctionOp,
DropFunctionOp,
UpdateFunctionOp,
)
from sqlalchemy_declarative_extensions.grant.compare import (
GrantPrivilegesOp,
RevokePrivilegesOp,
)
from sqlalchemy_declarative_extensions.procedure.compare import (
CreateProcedureOp,
DropProcedureOp,
UpdateProcedureOp,
)
from sqlalchemy_declarative_extensions.role.compare import (
CreateRoleOp,
DropRoleOp,
UpdateRoleOp,
)
from sqlalchemy_declarative_extensions.schema.compare import (
CreateSchemaOp,
DropSchemaOp,
SchemaOp,
)
from sqlalchemy_declarative_extensions.trigger.compare import (
CreateTriggerOp,
DropTriggerOp,
UpdateTriggerOp,
)
from sqlalchemy_declarative_extensions.view.compare import (
CreateViewOp,
DropViewOp,
UpdateViewOp,
)
from codegen_database.utils.sqlformat import format_sql
if TYPE_CHECKING:
from alembic.autogenerate.api import AutogenContext
# Maximum line width for inline SQL before forcing multi-line.
_MAX_LINE = 80
# ---------------------------------------------------------------------------
# SQL formatting
# ---------------------------------------------------------------------------
def _prettify(
sql: str,
*,
compact_lists_margin: int = 76,
) -> str:
"""Format a SQL statement using pglast."""
return format_sql(
sql,
compact_lists_margin=compact_lists_margin,
).rstrip("\n")
# ---------------------------------------------------------------------------
# Code generation helpers
# ---------------------------------------------------------------------------
def _render_execute(sql: str, *, fstring: bool = False) -> str:
"""Render ``op.execute(\"\"\"...\"\"\")``, going multi-line when needed."""
prefix = "f" if fstring else ""
inline = f'op.execute({prefix}"""{sql}""")'
if "\n" in sql or sql.endswith('"') or len(inline) > _MAX_LINE:
return f'op.execute({prefix}"""\n{indent(sql, " ")}\n""")'
return inline
def _render_execute_text(sql: str) -> str:
"""Render ``op.execute(sa.text(\"\"\"...\"\"\"))``, always multi-line."""
return (
f'op.execute(\n sa.text("""\n{indent(sql, " ")}\n """)\n)'
)
# ---------------------------------------------------------------------------
# Per-op-type renderers
# ---------------------------------------------------------------------------
def _render_sql_op(
autogen_context: AutogenContext,
op: Any, # noqa: ANN401
) -> list[str]:
"""Render ops whose SQL pglast can format (views, etc.)."""
assert autogen_context.connection # noqa: S101
commands = op.to_sql(autogen_context.connection.dialect)
return [_render_execute(_prettify(cmd)) for cmd in commands]
def _render_ddl_op(
_autogen_context: AutogenContext,
op: Any, # noqa: ANN401
) -> list[str]:
"""Render function/procedure ops (pglast-formatted).
Body whitespace isn't round-tripped (alembic indents the nested op text);
the comparison handles that -- see ``dependency._canonical_function_body``.
"""
return [_render_execute(_prettify(cmd)) for cmd in op.to_sql()]
def _render_trigger(
autogen_context: AutogenContext,
op: Any, # noqa: ANN401
) -> list[str]:
"""Render trigger ops."""
assert autogen_context.connection # noqa: S101
commands = op.to_sql(autogen_context.connection)
return [_render_execute(_prettify(cmd)) for cmd in commands]
def _render_schema(
autogen_context: AutogenContext,
op: SchemaOp,
) -> list[str]:
"""Render a schema op, using DDL objects where possible."""
statements = op.to_sql()
cls_names = {
s.__class__.__name__
for s in statements
if isinstance(s, (CreateSchema, DropSchema))
}
if cls_names:
autogen_context.imports.add(
f"from sqlalchemy.sql.ddl import {', '.join(cls_names)}"
)
return [
f'op.execute({cmd.__class__.__name__}("{cmd.element}"))'
if isinstance(cmd, (CreateSchema, DropSchema))
else _render_execute(_prettify(str(cmd)))
for cmd in statements
]
def _render_role(
autogen_context: AutogenContext,
op: Any, # noqa: ANN401
) -> list[str]:
"""Render a role op with pglast-formatted SQL."""
is_dynamic = op.role.is_dynamic
if is_dynamic:
autogen_context.imports.add("import os")
return [
_render_execute(_prettify(cmd), fstring=is_dynamic)
for cmd in op.to_sql(raw=False)
]
def _render_grant(
_autogen_context: AutogenContext,
op: Any, # noqa: ANN401
) -> str:
"""Render a grant/revoke with pglast-formatted SQL."""
return _render_execute_text(
_prettify(str(op.to_sql()), compact_lists_margin=72)
)
_RENDERER_MAP: dict[type, Any] = {
CreateViewOp: _render_sql_op,
UpdateViewOp: _render_sql_op,
DropViewOp: _render_sql_op,
CreateFunctionOp: _render_ddl_op,
UpdateFunctionOp: _render_ddl_op,
DropFunctionOp: _render_ddl_op,
CreateProcedureOp: _render_ddl_op,
UpdateProcedureOp: _render_ddl_op,
DropProcedureOp: _render_ddl_op,
CreateTriggerOp: _render_trigger,
UpdateTriggerOp: _render_trigger,
DropTriggerOp: _render_trigger,
CreateSchemaOp: _render_schema,
DropSchemaOp: _render_schema,
CreateRoleOp: _render_role,
UpdateRoleOp: _render_role,
DropRoleOp: _render_role,
GrantPrivilegesOp: _render_grant,
RevokePrivilegesOp: _render_grant,
}
#: Types whose DDL is *not* their ``impl`` (COORDINATE compiles to a
#: PostGIS ``geography``, STDADDR to a composite) -- these render as
#: themselves, with an import, so the migration recreates the real
#: column type.
_CODEGEN_DATABASE_SELF_RENDERING: frozenset[str] = frozenset(
{
"COORDINATE",
"STDADDR",
}
)
def _in_scope(module: str) -> bool:
"""Whether *module* is one this renderer takes responsibility for.
Only our own types plus ``sqlalchemy_utils``'s encrypted types
(matched by module path so codegen_database needs no runtime
dependency on it) -- a blanket impl-rewrite of *every*
``TypeDecorator`` would silently change how consumers' own
custom types land in their migrations.
"""
return module == "codegen_database.types" or module.startswith(
("codegen_database.types.", "sqlalchemy_utils.types.encrypted"),
)
[docs]
def render_item(
type_: str,
obj: Any, # noqa: ANN401
autogen_context: AutogenContext,
) -> str | Literal[False]:
"""Render custom column types for migrations.
The in-scope types (``codegen_database.types`` plus
``sqlalchemy_utils`` encrypted types, e.g. behind
``fsh_lib.oauth``'s token columns) are ``TypeDecorator`` wrappers
around plain SQL types, carrying Python-side state a migration
can neither render nor use -- an enum class, an encryption-key
callable. Alembic's default rendering emits their dotted class
name with no import, producing a migration that doesn't run.
Migrations only need the database-level type, so this renders
each type's own ``impl`` (``TextEnum`` -> ``sa.Text()``,
``EncryptedText`` -> ``sa.Text()``, ...) instead of maintaining
a parallel name -> DDL table.
Pass this as ``render_item`` to
``context.configure(...)`` in your Alembic ``env.py``.
Returns ``False`` for unrecognized objects so Alembic falls
through to its default rendering.
"""
if type_ != "type":
return False
if not _in_scope(type(obj).__module__):
return False
cls_name = type(obj).__name__
if cls_name in _CODEGEN_DATABASE_SELF_RENDERING:
autogen_context.imports.add(
f"from codegen_database.types import {cls_name}"
)
return f"{cls_name}()"
if isinstance(obj, sa_types.TypeDecorator):
# Delegate to alembic's own type renderer so dialect types
# get their prefix/imports handled exactly as a plain column
# of that type would.
return _repr_type(obj.impl_instance, autogen_context)
return False
[docs]
def register_renderers() -> None:
"""Override the library's renderers with pglast-formatted versions."""
for op_type, renderer in _RENDERER_MAP.items():
renderers.dispatch_for(op_type, replace=True)(renderer)