Source code for codegen_database.ext.chart

"""Chart extension for codegen_database ledgers.

:class:`ChartExtension` wires the :func:`codegen_database_date_bin
<codegen_database.ext.chart.date_bin.construct_date_bin_function>` polyfill
into the codegen_database lifecycle.  Register it on a
:class:`~codegen_database.config.CodegenDatabaseConfig` to make the ledger chart
function builders usable::

    from codegen_database.config import CodegenDatabaseConfig
    from codegen_database.ext.chart import ChartExtension

    config = CodegenDatabaseConfig()
    config.use(ChartExtension())

After registration, ``codegen_database_date_bin`` lives in the codegen_database
utility schema (``CodegenDatabaseConfig.utility_schema``,
``"codegen_database"`` by default).  The ledger chart function builders
find it via
``metadata.info["codegen_database_chart_schema"]`` and call into it with the
correct schema qualification.

Without this extension, the chart function builders raise
:class:`~codegen_database.errors.CodegenDatabaseValidationError` with a message
pointing the user at ``config.use(ChartExtension())``.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from sqlalchemy_declarative_extensions import (
    Schemas,
    register_function,
)
from sqlalchemy_declarative_extensions.dialects.postgresql import (
    Function,
    FunctionSecurity,
    FunctionVolatility,
)

from codegen_database.config import resolve_utility_schema
from codegen_database.ext.chart._contract import (
    CHART_EXTENSION_NAME,
    CHART_SCHEMA_KEY,
    assert_chart_extension_declared,
)
from codegen_database.ext.chart.date_bin import (
    construct_date_bin_function,
)
from codegen_database.extension import CodegenDatabaseExtension

if TYPE_CHECKING:
    from sqlalchemy import MetaData

__all__ = [
    "CHART_EXTENSION_NAME",
    "CHART_SCHEMA_KEY",
    "ChartExtension",
    "assert_chart_extension_declared",
]




@dataclass
class ChartExtension(CodegenDatabaseExtension):
    """Register the ``codegen_database_date_bin`` polyfill into metadata.

    When registered on a
    :class:`~codegen_database.config.CodegenDatabaseConfig`:

    - ``configure_metadata`` creates :func:`codegen_database_date_bin
      <codegen_database.ext.chart.date_bin.construct_date_bin_function>`
      in the resolved schema and stores it under
      ``metadata.info["codegen_database_chart_schema"]`` so the ledger chart
      function builders can resolve the helper without a
      user-supplied override.
    - that schema is also added to ``metadata.info["schemas"]`` so
      Alembic autogenerate creates it.

    Ledger chart function builders
    (:func:`~codegen_database.ext.ledger.chart_functions.construct_ledger_chart_function`
    and friends) call :func:`assert_chart_extension_declared`
    internally and raise
    :class:`~codegen_database.errors.CodegenDatabaseValidationError`
    when this extension is absent.

    Args:
        name: Extension name.  Defaults to ``"codegen_database-chart"``.
        schema: Schema for the polyfill function.  ``None`` (the
            default) uses the config's ``utility_schema``, so
            ``codegen_database_date_bin`` shares the one
            codegen_database utility schema with every other
            codegen_database-managed object.  Set a
            string to override.

    """

    name: str = CHART_EXTENSION_NAME
    schema: str | None = None



    def configure_metadata(self, metadata: MetaData) -> None:
        """Register the polyfill, schema, and marker on *metadata*."""
        schema = (
            self.schema
            if self.schema is not None
            else resolve_utility_schema(metadata)
        )
        existing_schemas: Schemas | None = metadata.info.get("schemas")

        if existing_schemas is None:
            metadata.info["schemas"] = Schemas(ignore_unspecified=True).are(
                schema
            )

        else:
            # ``Schemas.are(...)`` replaces the tuple rather than merging.
            # Accumulate so earlier-registered schemas (from tables/views)
            # survive this call.
            names = {s.name for s in existing_schemas.schemas}
            names.add(schema)
            metadata.info["schemas"] = existing_schemas.are(*sorted(names))

        metadata.info[CHART_SCHEMA_KEY] = schema

        spec = construct_date_bin_function()
        existing_fns = metadata.info.get("functions")

        if existing_fns is not None:
            for fn in existing_fns:
                if fn.name == spec["name"] and fn.schema == schema:
                    return

        register_function(
            metadata,
            Function(
                spec["name"],
                spec["definition"],
                returns=spec.get("returns", "void"),
                language=spec.get("language", "sql"),
                schema=schema,
                parameters=spec.get("parameters", []),
                security=spec.get("security", FunctionSecurity.invoker),
                volatility=spec.get("volatility", FunctionVolatility.VOLATILE),
            ),
        )