"""Rewrite autogenerated revisions so every scoped table is protected in the same transaction.
The hook reads the compiled metadata only. For each table it emits one
sequence under the guard's hatch: create then protect, unprotect then drop, or
unprotect, change, protect when a scope column changes. A partitioned table is
dropped after its partitions, which the guard unprotects, detaches and drops;
partitions themselves are never autogenerated, a revision creates them with
``op.ensure_range_partitions``. A downgrade never protects with the head
scopes: it unprotects and stops on a ``NotImplementedError`` the developer
replaces with the target revision's protect. The revision carries loom's guard
operations, never SQL.
"""
from __future__ import annotations
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any
from alembic.operations import ops
from loom.core.backend.scoped_ddl import privilege_names, scope_documents
from loom.core.model.introspection import declared_privileges, get_table_name, is_row_scoped
from loom.core.model.privilege import Privilege
from loom.core.model.scoped import ScopedTable
from loom.core.repository.sqlalchemy.migrations.operations import (
DropPartitionsOp,
GrantTableOp,
HandWrittenProtectOp,
OpenHatchOp,
ProtectScopedTableOp,
UnprotectScopedTableOp,
)
if TYPE_CHECKING:
from loom.core.locator import Application
RevisionHook = Callable[[Any, Any, Sequence[ops.MigrationScript]], None]
_Seal = Callable[[str], ops.MigrateOperation]
_COLUMN_OPS = (ops.AlterColumnOp, ops.DropColumnOp)
[docs]
def scope_protection_hook(application: Application) -> RevisionHook:
"""Return the ``process_revision_directives`` callable for ``application``."""
rewriter = _Rewriter(application)
def hook(_context: Any, _revision: Any, directives: Sequence[ops.MigrationScript]) -> None:
for script in directives:
if script.upgrade_ops is not None:
script.upgrade_ops.ops[:] = rewriter.rewrite(script.upgrade_ops.ops)
if script.downgrade_ops is not None:
script.downgrade_ops.ops[:] = rewriter.rewrite(
script.downgrade_ops.ops, downgrade=True
)
return hook
class _Rewriter:
def __init__(self, application: Application) -> None:
self._scoped: Mapping[str, ScopedTable] = {
table.name: table for table in application.scoped.values()
}
self._privileges: dict[str, Mapping[str, frozenset[Privilege]]] = {
get_table_name(model): declared_privileges(model)
for model in application.models
if not is_row_scoped(model) and declared_privileges(model)
}
def rewrite(
self, operations: Iterable[ops.MigrateOperation], *, downgrade: bool = False
) -> list[ops.MigrateOperation]:
seal = self._hand_written if downgrade else self._protect
rewritten: list[ops.MigrateOperation] = []
for op in operations:
rewritten.extend(self._sequence_for(op, seal))
return rewritten
def _sequence_for(self, op: ops.MigrateOperation, seal: _Seal) -> list[ops.MigrateOperation]:
if isinstance(op, ops.CreateTableOp):
return self._create(op, seal)
if isinstance(op, ops.DropTableOp) and op.table_name in self._scoped:
return [
self._hatch(),
*self._partitions(op.table_name),
self._unprotect(op.table_name),
op,
]
if isinstance(op, ops.ModifyTableOps) and self._touches_scope(op):
return [self._hatch(), self._unprotect(op.table_name), op, seal(op.table_name)]
return [op]
def _create(self, op: ops.CreateTableOp, seal: _Seal) -> list[ops.MigrateOperation]:
if op.table_name in self._scoped:
return [self._hatch(), op, seal(op.table_name)]
privileges = self._privileges.get(op.table_name)
if privileges is None:
return [op]
readers = privilege_names(privileges.get("readers", frozenset()))
writers = privilege_names(privileges.get("writers", frozenset()))
return [op, GrantTableOp(op.table_name, readers, writers)]
def _touches_scope(self, op: ops.ModifyTableOps) -> bool:
scoped = self._scoped.get(op.table_name)
if scoped is None:
return False
columns = {column.column for column in scoped.scopes}
return any(_column_name(inner) in columns for inner in op.ops)
def _partitions(self, table_name: str) -> list[ops.MigrateOperation]:
if self._scoped[table_name].partition_by is None:
return []
return [DropPartitionsOp(table_name)]
def _hatch(self) -> OpenHatchOp:
return OpenHatchOp()
def _protect(self, table_name: str) -> ProtectScopedTableOp:
scoped = self._scoped.get(table_name)
if scoped is None:
raise ValueError(f"{table_name!r} is not a scoped table of this application")
return ProtectScopedTableOp(
table_name, scope_documents(scoped), privilege_names(scoped.privileges)
)
def _hand_written(self, table_name: str) -> HandWrittenProtectOp:
return HandWrittenProtectOp(table_name)
def _unprotect(self, table_name: str) -> UnprotectScopedTableOp:
return UnprotectScopedTableOp(table_name)
def _column_name(op: ops.MigrateOperation) -> str | None:
if isinstance(op, ops.AddColumnOp):
return op.column.name
if isinstance(op, _COLUMN_OPS):
return op.column_name
return None