"""Ledger-backed state machine factory.
:class:`CodegenDatabaseStateMachine` creates a validated transition log for
any entity that moves through a set of named states. The valid
transitions are compiled directly into a ``BEFORE INSERT`` trigger
function, making them schema-level constraints that Alembic tracks
like any other database object.
Structure
---------
For a state machine named ``order_status`` in schema ``app``:
- ``app.order_status_raw`` -- append-only transition log.
Columns: ``id`` (PK), ``subject_id`` (FK to subject table),
``from_state`` TEXT nullable (NULL for the initial transition),
``to_state`` TEXT not null, ``created_at``, plus any
*extra_columns*.
- ``app.order_status`` -- view showing the full transition log.
- ``app.order_status_current`` -- view showing the current state per
subject via ``DISTINCT ON (subject_id) ORDER BY created_at DESC``.
Usage::
order_status = CodegenDatabaseStateMachine(
name="order_status",
schemaname="app",
metadata=metadata,
subject_column=Column(
"order_id", Integer,
CodegenDatabaseForeignKey("orders.id"), nullable=False
),
states=["pending", "confirmed", "shipped", "delivered", "cancelled"],
valid_transitions=[
(None, "pending"), # initial state
("pending", "confirmed"),
("pending", "cancelled"),
("confirmed", "shipped"),
("shipped", "delivered"),
("shipped", "cancelled"),
],
extra_columns=[
Column("changed_by", String, nullable=False),
Column("notes", Text, nullable=True),
],
)
# Readable proxies
order_status.table # the raw transitions table
order_status.current_view # current state per subject
# Query builders
current_q = construct_state_machine_current_query(order_status)
history_q = construct_state_machine_history_query(order_status)
"""
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any
from sqlalchemy import (
CheckConstraint,
Column,
DateTime,
MetaData,
Select,
Table,
Text,
Uuid,
bindparam,
func,
insert,
select,
text,
union_all,
)
from sqlalchemy.dialects import postgresql as pg_dialect
from sqlalchemy_declarative_extensions import (
View,
register_function,
register_trigger,
register_view,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionParam,
FunctionSecurity,
Trigger,
)
from codegen_database.declarative import _coerce_arg
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.factory.context import FactoryContext
from codegen_database.fk import (
CodegenDatabaseForeignKey,
DimensionRef,
register_dimension,
)
from codegen_database.plugin import (
Plugin,
PluginCollection,
produces,
requires,
)
from codegen_database.plugins.column_name import construct_column_name_plugin
from codegen_database.plugins.fk import TableFKPlugin
from codegen_database.plugins.pk import SerialPKPlugin
from codegen_database.types import TextEnum
from codegen_database.utils.naming import resolve_name
from codegen_database.utils.query import compile_query
from codegen_database.utils.template import load_template
from codegen_database.validator import reject_frozen_function_defaults
if TYPE_CHECKING:
import enum
from collections.abc import Iterable, Iterator
_TEMPLATES = Path(__file__).resolve().parent / "templates"
_NAMING_DEFAULTS = {
"sm_raw_table": "%(name)s_raw",
"sm_validate_function": "%(schema)s_%(name)s_validate_transition",
"sm_validate_trigger": "%(schema)s_%(name)s_validate_transition",
"sm_current_view": "%(name)s_current",
"sm_transition_to_function": "%(schema)s_%(name)s_transition_to",
"sm_sync_function": "%(schema)s_%(name)s_sync_%(column)s",
"sm_sync_trigger": "%(schema)s_%(name)s_sync_%(column)s",
}
def _index_from_states_by_target(
valid_transitions: list[tuple[str | None, str]],
) -> dict[str, list[str | None]]:
"""Build ``to_state -> [legal from_states]`` from the SM's transitions."""
by_target: dict[str, list[str | None]] = {}
for from_s, to_s in valid_transitions:
by_target.setdefault(to_s, []).append(from_s)
return by_target
[docs]
class CodegenDatabaseStateMachine:
r"""Ledger-backed state machine with transition validation.
Creates an append-only transitions table, a full history view, and
a current-state view. A ``BEFORE INSERT`` trigger validates that
each new transition is:
1. Declared in *valid_transitions*.
2. Consistent with the subject's actual current state.
The valid transitions are compiled into the trigger function body
at migration time, so changing them requires a migration.
Args:
name: Base name for all generated objects.
schemaname: PostgreSQL schema for all objects.
metadata: SQLAlchemy ``MetaData`` to register on.
subject_column: ``Column`` that identifies the entity being
tracked. Use ``CodegenDatabaseForeignKey`` to reference the
subject's table. Must be ``nullable=False``.
states: List of valid state strings.
valid_transitions: List of ``(from_state, to_state)`` pairs.
Use ``None`` as *from_state* for the initial transition.
extra_columns: Additional columns on the raw table (e.g.
``changed_by``, ``notes``).
allow_initial: If ``True`` (default), a row with
``from_state = NULL`` is allowed as the first transition
for a subject. Set to ``False`` to require an explicit
initial state defined in *valid_transitions*.
Raises:
CodegenDatabaseValidationError: If *states* or *valid_transitions* is
empty, or if a transition references an undeclared state.
"""
def __init__( # noqa: PLR0913, PLR0915
self,
name: str,
schemaname: str,
metadata: MetaData,
subject_column: Column,
states: list[str],
valid_transitions: list[tuple[str | None, str]],
extra_columns: list[Column] | None = None,
*,
actor_column: str | None = None,
actor_fk: str | None = None,
allow_initial: bool = True,
) -> None:
"""Build the state machine and register all objects."""
self._validate(states, valid_transitions)
self.name = name
self.schemaname = schemaname
self.metadata = metadata
raw_name = resolve_name(
metadata,
"sm_raw_table",
{"name": name, "schema": schemaname},
_NAMING_DEFAULTS,
)
# PK plugin
pk_plugin = SerialPKPlugin()
ctx = FactoryContext(
tablename=name,
schemaname=schemaname,
metadata=metadata,
schema_items=[subject_column, *(extra_columns or [])],
plugins=[],
)
pk_plugin.run(ctx)
construct_column_name_plugin("created_at_column", "created_at").run(ctx)
pk_columns = ctx["pk_columns"]
valid_state_set = set(states)
# State check expression.
state_values = ", ".join(f"'{s}'" for s in states)
state_check = CheckConstraint(
f"to_state IN ({state_values})",
name=f"{raw_name}_to_state_ck",
)
from_state_check = CheckConstraint(
f"from_state IS NULL OR from_state IN ({state_values})",
name=f"{raw_name}_from_state_ck",
)
_actor_col: Column | None = None
if actor_column is not None and actor_fk is not None:
_actor_col = Column(
actor_column,
Uuid,
CodegenDatabaseForeignKey(actor_fk),
server_default=text(
"current_setting('app.current_token', true)::uuid"
),
nullable=True,
)
raw_table = Table(
raw_name,
metadata,
*pk_columns,
Column("from_state", Text, nullable=True),
Column("to_state", Text, nullable=False),
Column(
"created_at",
DateTime(timezone=True),
server_default=func.now(),
nullable=False,
),
subject_column,
*(extra_columns or []),
*([_actor_col] if _actor_col is not None else []),
state_check,
from_state_check,
schema=schemaname,
)
reject_frozen_function_defaults([raw_table])
fk_ctx = FactoryContext(
tablename=name,
schemaname=schemaname,
metadata=metadata,
schema_items=[
subject_column,
*(extra_columns or []),
*([_actor_col] if _actor_col is not None else []),
],
plugins=[],
)
fk_ctx["raw_table"] = raw_table
TableFKPlugin("raw_table").run(fk_ctx)
self.table = raw_table
raw_fullname = f"{schemaname}.{raw_name}"
subject_col_name = subject_column.name
# Validate trigger.
validate_fn_name = resolve_name(
metadata,
"sm_validate_function",
{"name": name, "schema": schemaname},
_NAMING_DEFAULTS,
)
validate_trig_name = resolve_name(
metadata,
"sm_validate_trigger",
{"name": name, "schema": schemaname},
_NAMING_DEFAULTS,
)
template = load_template(
_TEMPLATES / "validate_transition.plpgsql.mako"
)
body = template.render(
name=name,
transitions_table=raw_fullname,
subject_col=subject_col_name,
valid_transitions=valid_transitions,
allow_initial=allow_initial,
valid_state_set=valid_state_set,
)
register_function(
metadata,
Function(
validate_fn_name,
body,
returns="trigger",
language="plpgsql",
schema=schemaname,
security=FunctionSecurity.definer,
),
)
register_trigger(
metadata,
Trigger.before(
"insert",
on=raw_fullname,
execute=f"{schemaname}.{validate_fn_name}",
name=validate_trig_name,
).for_each_row(),
)
# transition_to(subject, to_state, ...extras) helper. Locks the
# latest row per subject for the duration of the transaction so
# that concurrent callers serialize through PG, then inserts the
# new row with from_state = actual current state.
transition_fn_name = resolve_name(
metadata,
"sm_transition_to_function",
{"name": name, "schema": schemaname},
_NAMING_DEFAULTS,
)
extra_col_names = [c.name for c in (extra_columns or [])]
transition_template = load_template(
_TEMPLATES / "transition_to.plpgsql.mako"
)
transition_body = transition_template.render(
transitions_table=raw_fullname,
subject_col=subject_col_name,
extra_columns=extra_col_names,
)
subject_pg_type = subject_column.type.compile(
dialect=pg_dialect.dialect()
)
transition_params = [
FunctionParam.input("p_subject", subject_pg_type),
FunctionParam.input("p_to_state", "text"),
*(
FunctionParam.input(
f"p_{col.name}",
col.type.compile(dialect=pg_dialect.dialect()),
)
for col in (extra_columns or [])
),
]
register_function(
metadata,
Function(
transition_fn_name,
transition_body,
returns="text",
language="plpgsql",
schema=schemaname,
parameters=transition_params,
security=FunctionSecurity.definer,
),
)
# History view (full log ordered by subject, then time).
history_sql = compile_query(
select(raw_table).order_by(
raw_table.c[subject_col_name],
raw_table.c["created_at"],
)
)
register_view(
metadata,
View(name, history_sql, schema=schemaname),
)
self._history_proxy = Table(
name,
MetaData(),
*[
Column(c.name, c.type, primary_key=c.primary_key)
for c in raw_table.columns
],
schema=schemaname,
)
# Current-state view (latest transition per subject).
current_view_name = resolve_name(
metadata,
"sm_current_view",
{"name": name, "schema": schemaname},
_NAMING_DEFAULTS,
)
subject_col_obj = raw_table.c[subject_col_name]
pk_col_name = pk_columns.first_key
current_sql = compile_query(
select(raw_table)
.distinct(subject_col_obj)
.order_by(
subject_col_obj,
raw_table.c["created_at"].desc(),
raw_table.c[pk_col_name].desc(),
)
)
register_view(
metadata,
View(current_view_name, current_sql, schema=schemaname),
)
self.current_view = Table(
current_view_name,
MetaData(),
*[
Column(c.name, c.type, primary_key=c.primary_key)
for c in raw_table.columns
],
schema=schemaname,
)
# Register for FK resolution.
register_dimension(
metadata,
name,
DimensionRef(schema=schemaname, table=raw_name),
)
# Stash the bits ``transition_to`` needs at call time.
self._transition_fn_name = transition_fn_name
self._subject_col_name = subject_col_name
self._extra_col_names = list(extra_col_names)
self._from_states_for = _index_from_states_by_target(valid_transitions)
[docs]
def transition_to(
self,
db: Any, # noqa: ANN401
**kwargs: Any, # noqa: ANN401
) -> Any: # noqa: ANN401
"""Invoke the SM's generated ``transition_to`` SQL function.
Builds ``select(func.<schema>.<schema>_<name>_transition_to(
...))`` and dispatches via ``db.execute(...)``. The SQL
function locks the latest raw row for the subject, reads
the actual current state, and inserts a new row with
``from_state = <current>`` -- the BEFORE INSERT trigger
then validates the ``(from_state, to_state)`` pair.
Kwargs are keyed by the subject column's name (e.g.
``line_item_id`` for a state machine whose subject is a
line item), ``to_state``, plus any extra-column names
declared on the SM. :class:`enum.Enum` values are coerced
to ``.value`` so callers can pass typed enums directly.
Args:
db: Any object with an ``execute(stmt)`` method --
``Session``, ``AsyncSession``, or ``Connection``.
codegen_database does no awaiting itself; if ``execute``
returns an awaitable, ``await`` the return value
at the call site.
**kwargs: ``<subject_col_name>``, ``to_state``, and
any extra-column kwargs declared on the SM.
Returns:
Whatever ``db.execute(stmt)`` returns -- a ``Result``
for sync sessions, an ``Awaitable[Result]`` for async
ones.
Raises:
KeyError: If the subject or ``to_state`` kwarg is
missing.
"""
args = [
_coerce_arg(kwargs[self._subject_col_name]),
_coerce_arg(kwargs["to_state"]),
*(
_coerce_arg(kwargs.get(col_name))
for col_name in self._extra_col_names
),
]
fn_ref = getattr(
getattr(func, self.schemaname),
self._transition_fn_name,
)
return db.execute(select(fn_ref(*args)))
[docs]
def can_transition_to(
self,
from_state: Any, # noqa: ANN401
to_state: Any, # noqa: ANN401
) -> bool:
"""Return whether ``from_state -> to_state`` is a legal transition.
Args:
from_state: The subject's current state, or ``None`` for
a subject that hasn't transitioned yet.
to_state: The state being checked.
Returns:
``True`` if ``(from_state, to_state)`` is in
*valid_transitions*.
"""
from_value = _coerce_arg(from_state)
to_value = _coerce_arg(to_state)
return from_value in self._from_states_for.get(to_value, [])
[docs]
async def transition_many(
self,
db: Any, # noqa: ANN401
*,
subjects: Select | Iterable[Any],
to_state: Any, # noqa: ANN401
**extras: Any, # noqa: ANN401
) -> Any: # noqa: ANN401
"""Bulk-transition every subject in *subjects* to *to_state*.
Issues a single ``INSERT INTO raw (subject, from_state,
to_state, ...) SELECT ... FROM current_view`` after locking
the latest raw row per affected subject -- mirroring what
:meth:`transition_to` does for a single subject so concurrent
single-row and bulk callers serialize against each other.
The set of legal current states is **auto-derived from the
SM's declared transitions**: only subjects whose current
state has a registered ``(from, to_state)`` edge are
advanced. Subjects in any other state are silently skipped.
This is the same set the BEFORE INSERT trigger would accept,
applied as a ``WHERE`` filter so the trigger never has to
reject rows.
The BEFORE INSERT trigger validates each emitted row, so a
misconfigured ``to_state`` (one that's not a valid target
from any state) results in zero rows inserted -- not an
error.
**AsyncSession only.** Unlike :meth:`transition_to` and
:meth:`~codegen_database.declarative.CodegenDatabaseFunctionMixin.call`,
this method issues two statements (the ``FOR UPDATE`` lock
and the ``INSERT ... SELECT``) and awaits both internally
-- so it cannot dispatch transparently through a sync
``Session``. If you need a sync variant, run the two
statements manually with the helpers exposed on
``self.table`` and ``self.current_view``.
Args:
db: An :class:`~sqlalchemy.ext.asyncio.AsyncSession`.
subjects: A SQLAlchemy :class:`~sqlalchemy.Select`
that yields subject ids, or any iterable of ids
(UUIDs, ints, etc.). Iterables are bound as a
``VALUES`` list; a ``Select`` is correlated as a
subquery so callers can compose against other
tables (``select(LineItem.id).where(...)``).
to_state: Target state. :class:`enum.Enum` values are
coerced via ``.value``.
**extras: Values for the SM's extra columns
(``actor_id``, ``note``, ...). Each value is
applied to every row. ``Enum`` values are coerced.
Returns:
The :class:`~sqlalchemy.engine.Result` from the
``INSERT ... RETURNING`` -- iterate ``.scalars()`` for
the subject ids that actually transitioned.
"""
raw = self.table
current = self.current_view
subject_col_name = self._subject_col_name
to_state_value = _coerce_arg(to_state)
valid_from_states = self._from_states_for.get(to_state_value, [])
if isinstance(subjects, Select):
subjects_select = subjects
else:
subjects_select = union_all(
*(
select(
bindparam(
None, subject, type_=raw.c[subject_col_name].type
)
)
for subject in subjects
)
)
# Lock latest raw row per affected subject so single-row and
# bulk transitions serialize on the same rows. MAX(id) picks
# the latest reliably because the PK is a serial -- the same
# tie-breaker ``current_view`` uses after ``created_at DESC``.
# DISTINCT ON would be cleaner but Postgres forbids
# ``FOR UPDATE`` on a SELECT that carries it.
latest_id_per_subject = (
select(func.max(raw.c.id).label("max_id"))
.where(raw.c[subject_col_name].in_(subjects_select))
.group_by(raw.c[subject_col_name])
.subquery()
)
await db.execute(
select(raw.c.id)
.where(raw.c.id.in_(select(latest_id_per_subject.c.max_id)))
.with_for_update(),
)
select_cols: list = [
current.c[subject_col_name],
current.c.to_state, # becomes from_state on the new row
bindparam(
"p_to_state",
to_state_value,
type_=raw.c.to_state.type,
),
]
insert_cols: list[str] = [subject_col_name, "from_state", "to_state"]
for col_name in self._extra_col_names:
insert_cols.append(col_name)
select_cols.append(
bindparam(
f"p_{col_name}",
_coerce_arg(extras.get(col_name)),
type_=raw.c[col_name].type,
),
)
return await db.execute(
insert(raw)
.from_select(
insert_cols,
select(*select_cols).where(
current.c[subject_col_name].in_(subjects_select),
current.c.to_state.in_(valid_from_states),
),
)
.returning(raw.c[subject_col_name]),
)
@staticmethod
def _validate(
states: list[str],
valid_transitions: list[tuple[str | None, str]],
) -> None:
"""Raise CodegenDatabaseValidationError if configuration is invalid."""
if not states:
msg = "CodegenDatabaseStateMachine: states must be a non-empty list"
raise CodegenDatabaseValidationError(msg)
if not valid_transitions:
msg = (
"CodegenDatabaseStateMachine: "
"valid_transitions must be non-empty"
)
raise CodegenDatabaseValidationError(msg)
state_set = set(states)
for from_s, to_s in valid_transitions:
if from_s is not None and from_s not in state_set:
msg = (
f"CodegenDatabaseStateMachine: transition from_state "
f"{from_s!r} is not in states."
)
raise CodegenDatabaseValidationError(msg)
if to_s not in state_set:
msg = (
f"CodegenDatabaseStateMachine: transition to_state "
f"{to_s!r} is not in states."
)
raise CodegenDatabaseValidationError(msg)
@produces("exclude_from_mutations")
class _CurrentStateColumnPlugin(Plugin):
"""Inject the denormalized current-state column on the subject."""
def __init__(
self,
column_name: str,
column_type: type[enum.Enum],
) -> None:
"""Store the column name and enum type."""
self.column_name = column_name
self.column_type = column_type
def run(self, ctx: FactoryContext) -> None:
"""Inject the column and mark it read-only for view mutations."""
ctx.schema_items.append(
Column(self.column_name, TextEnum(self.column_type), nullable=True)
)
ctx.setdefault("exclude_from_mutations", set()).add(self.column_name)
@requires("raw_table")
class _CurrentStateSyncPlugin(Plugin):
"""Registers a function + trigger on the transitions table that
writes each new ``to_state`` onto the subject's current-state column.
"""
def __init__(
self,
sm: CodegenDatabaseStateMachine,
column_name: str,
) -> None:
"""Store the state machine and the subject column to maintain."""
self.sm = sm
self.column_name = column_name
def run(self, ctx: FactoryContext) -> None:
"""Register the sync function and AFTER INSERT trigger."""
sm = self.sm
raw = ctx["raw_table"]
subst = {
"name": sm.name,
"schema": sm.schemaname,
"column": self.column_name,
}
fn_name = resolve_name(
ctx.metadata, "sm_sync_function", subst, _NAMING_DEFAULTS
)
trig_name = resolve_name(
ctx.metadata, "sm_sync_trigger", subst, _NAMING_DEFAULTS
)
body = load_template(
_TEMPLATES / "sync_current_state.plpgsql.mako"
).render(
subject_table=f"{raw.schema}.{raw.name}",
status_column=self.column_name,
pk_column=ctx.pk_column_name,
subject_fk_column=sm._subject_col_name,
)
register_function(
ctx.metadata,
Function(
fn_name,
body,
returns="trigger",
language="plpgsql",
schema=sm.schemaname,
security=FunctionSecurity.definer,
),
)
register_trigger(
ctx.metadata,
Trigger.after(
"insert",
on=f"{sm.schemaname}.{sm.table.name}",
execute=f"{sm.schemaname}.{fn_name}",
name=trig_name,
).for_each_row(),
)
[docs]
class StateMachineCurrentStatePlugin(PluginCollection):
"""Maintain a denormalized, read-only current-state column.
class Deal(Base):
__plugins__ = [
UUIDV7PKPlugin(),
StateMachineCurrentStatePlugin(
sm=deal_sm, column_type=DealStatus
),
]
status: Mapped[DealStatus]
Args:
sm: The :class:`CodegenDatabaseStateMachine` whose current state
is mirrored onto the subject.
column_type: The :class:`enum.Enum` the SM's states belong to.
column_name: Name of the current-state column on the subject
(default ``"status"``).
"""
def __init__(
self,
*,
sm: CodegenDatabaseStateMachine,
column_type: type[enum.Enum],
column_name: str = "status",
) -> None:
"""Store the state machine and the column to maintain."""
self.sm = sm
self.column_type = column_type
self.column_name = column_name
def __iter__(self) -> Iterator[Plugin]:
"""Yield the before-root column plugin, then the sync plugin."""
return iter(
[
_CurrentStateColumnPlugin(self.column_name, self.column_type),
_CurrentStateSyncPlugin(self.sm, self.column_name),
]
)
[docs]
def construct_state_machine_current_query(
sm: CodegenDatabaseStateMachine,
) -> Select:
"""Return a select against the current-state view.
Args:
sm: A :class:`CodegenDatabaseStateMachine` instance.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select` over
``{name}_current``.
"""
return select(sm.current_view)
[docs]
def construct_state_machine_history_query(
sm: CodegenDatabaseStateMachine,
subject_id: object = None,
) -> Select:
"""Return all transitions, optionally filtered to one subject.
Args:
sm: A :class:`CodegenDatabaseStateMachine` instance.
subject_id: When provided, adds a ``WHERE subject_col =
subject_id`` filter.
Returns:
A SQLAlchemy :class:`~sqlalchemy.Select`.
"""
q = select(sm.table).order_by(
sm.table.c["created_at"],
)
if subject_id is not None:
# Use the first non-PK, non-state, non-timestamp column as
# the subject column for filtering.
subject_col = next(
c
for c in sm.table.columns
if c.name not in {"id", "from_state", "to_state", "created_at"}
)
q = q.where(subject_col == subject_id)
return q