"""Declarative and imperative PostgreSQL function registration.
:class:`CodegenDatabaseFunction` works both imperatively and as a declarative
base class, mirroring how
:class:`~codegen_database.declarative.CodegenDatabaseView`
and :class:`~codegen_database.views.view.CodegenDatabasePlainView` work
for views.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, NotRequired, TypedDict
from sqlalchemy_declarative_extensions import register_function
from sqlalchemy_declarative_extensions.dialects.postgresql import (
Function,
FunctionParam,
FunctionSecurity,
FunctionVolatility,
)
from codegen_database.declarative import _parse_table_args
from codegen_database.errors import CodegenDatabaseValidationError
if TYPE_CHECKING:
from sqlalchemy import MetaData
[docs]
class CodegenDatabaseFunctionSpec(TypedDict):
"""Specification dict for a PostgreSQL function.
Returned by spec-builders such as
:func:`~codegen_database.ext.ledger.functions.ledger_event_function`,
:func:`~codegen_database.ext.ledger.chart_functions.construct_ledger_chart_function`,
and
:func:`~codegen_database.ext.chart.date_bin.construct_date_bin_function`.
Pass directly as ``__funcspec__`` on a
:class:`~codegen_database.declarative.CodegenDatabaseFunctionMixin`
subclass, or unpack into :class:`CodegenDatabaseFunction` for
imperative registration::
class InventoryAdjust(CodegenDatabaseFunctionMixin, Base):
__table_args__ = {"schema": "private"}
__funcspec__ = ledger_event_function(inventory, adjust_event)
Attributes:
name: Unqualified function name.
definition: SQL function body.
language: Function language (``"sql"``, ``"plpgsql"``, ...).
parameters: List of ``FunctionParam`` instances.
returns: Return type string (e.g. ``"SETOF schema.table"``).
security: Security mode.
volatility: Optional volatility
(``STABLE`` / ``IMMUTABLE`` / ``VOLATILE``).
Defaults to ``VOLATILE`` when omitted.
Note:
Schema is intentionally excluded. Pass it via
``__table_args__`` (declarative) or read it from the
source's ``ctx.schemaname`` (imperative).
"""
name: str
definition: str
language: str
parameters: list
returns: str
security: FunctionSecurity
volatility: NotRequired[FunctionVolatility]
[docs]
@dataclass(frozen=True)
class FunctionOptions:
"""Options for :class:`CodegenDatabaseFunction` declarative subclasses.
Set on a subclass via ``__options__``.
Attributes:
returns: Return type string.
language: Function language (``"sql"``, ``"plpgsql"``, ...).
parameters: List of ``FunctionParam`` instances from
``sqlalchemy_declarative_extensions.dialects.postgresql``.
security: ``FunctionSecurity.invoker`` or
``FunctionSecurity.definer``.
volatility: Volatility classification.
"""
returns: str = "void"
language: str = "sql"
parameters: list[FunctionParam] = field(default_factory=list)
security: FunctionSecurity = FunctionSecurity.invoker
volatility: FunctionVolatility = FunctionVolatility.VOLATILE
[docs]
class CodegenDatabaseFunction:
"""Register a PostgreSQL function on SQLAlchemy metadata.
Works both imperatively (direct instantiation) and as a declarative
base class (subclassing).
**Imperative usage** (``metadata`` read from class when omitted)::
class FunctionBase(CodegenDatabaseFunction):
metadata = Base.metadata
spec = ledger_event_function(source, event)
FunctionBase(
spec["name"],
source.ctx.schemaname,
spec["definition"],
parameters=spec["parameters"],
returns=spec["returns"],
security=spec["security"],
)
Or with explicit metadata::
CodegenDatabaseFunction(
name="inventory_adjust",
schema="private",
definition="...",
metadata=metadata,
returns="SETOF private.inventory_raw",
parameters=[...],
security=FunctionSecurity.definer,
)
**Declarative usage:**
.. code-block:: python
class FunctionBase(CodegenDatabaseFunction):
metadata = metadata
class InventoryAdjust(FunctionBase):
__funcname__ = "inventory_adjust"
__table_args__ = {"schema": "private"}
__definition__ = "..."
__options__ = FunctionOptions(
returns="SETOF private.inventory_raw",
parameters=[...],
security=FunctionSecurity.definer,
)
In both cases a
``sqlalchemy_declarative_extensions.dialects.postgresql.Function``
is registered on the provided ``MetaData`` instance so that
Alembic autogeneration picks it up.
Args:
name: Unqualified function name.
schema: PostgreSQL schema.
metadata: SQLAlchemy ``MetaData`` to register on.
definition: SQL (or PL/pgSQL) function body.
returns: Return type string (default ``"void"``).
language: Function language (default ``"sql"``).
parameters: List of ``FunctionParam`` instances from
``sqlalchemy_declarative_extensions.dialects.postgresql``
(default ``[]``).
security: Security mode — ``FunctionSecurity.invoker`` or
``FunctionSecurity.definer``
(default ``FunctionSecurity.invoker``).
volatility: Volatility classification
(default ``FunctionVolatility.VOLATILE``).
"""
# Set on project-level base subclasses by __init_subclass__.
metadata: ClassVar[MetaData]
# Set on declarative subclasses by __init_subclass__, and on
# imperative instances by __init__.
function: Function
def __init__( # noqa: PLR0913
self,
name: str,
schema: str,
definition: str,
*,
metadata: MetaData | None = None,
returns: str = "void",
language: str = "sql",
parameters: list[FunctionParam] | None = None,
security: FunctionSecurity = FunctionSecurity.invoker,
volatility: FunctionVolatility = FunctionVolatility.VOLATILE,
) -> None:
"""Create and register the function.
``metadata`` may be omitted when called on a subclass that
declares ``metadata`` as a class attribute::
class FunctionBase(CodegenDatabaseFunction):
metadata = Base.metadata
FunctionBase(**ledger_event_function(source, event))
"""
if metadata is None:
for klass in type(self).__mro__:
if "metadata" in klass.__dict__:
metadata = klass.__dict__["metadata"]
break
if metadata is None:
msg = (
f"{type(self).__name__}: metadata must be provided "
f"or set as a class attribute on the base class."
)
raise CodegenDatabaseValidationError(msg)
fn = Function(
name,
definition,
returns=returns,
language=language,
schema=schema,
parameters=parameters or [],
security=security,
volatility=volatility,
)
register_function(metadata, fn)
self.function = fn
self.name = name
self.schema = schema
def __init_subclass__(cls, **kwargs: Any) -> None: # noqa: ANN401
"""Register the function when a concrete subclass is defined."""
super().__init_subclass__(**kwargs)
# Direct subclasses of CodegenDatabaseFunction are project-level bases
# (e.g. the user's ``FunctionBase``). Stop here.
if CodegenDatabaseFunction in cls.__bases__:
return
# Skip intermediate classes that have no __funcname__.
if not hasattr(cls, "__funcname__"):
return
# Resolve metadata from MRO (set on the base class).
md: MetaData | None = None
for klass in cls.__mro__:
if "metadata" in klass.__dict__:
md = klass.__dict__["metadata"]
break
if md is None:
msg = (
f"{cls.__name__}: base class must define "
f"'metadata' before declaring function classes."
)
raise CodegenDatabaseValidationError(msg)
schema, _ = _parse_table_args(cls, schema_required=True)
definition: str = getattr(cls, "__definition__", "")
options: FunctionOptions = cls.__dict__.get(
"__options__", FunctionOptions()
)
fn = Function(
cls.__funcname__,
definition,
returns=options.returns,
language=options.language,
schema=schema,
parameters=options.parameters,
security=options.security,
volatility=options.volatility,
)
register_function(md, fn)
cls.function = fn