Extension system

While Plugin architecture compose behaviour within a single factory, extensions sit one level above: they bundle plugins, metadata hooks, Alembic hooks, and CLI commands into a single installable unit.

Extensions make it possible for third-party packages to extend codegen_database, and for codegen_database’s own opt-in subsystems (pg_cron) to be cleanly separated.

Built-in extensions

codegen_database ships three opt-in extensions:

  • pg_cron extension — declarative cron job scheduling and CREATE EXTENSION management via pg_cron.

  • PostGIS extension — declares the postgis extension required by PostGIS-backed types such as COORDINATE.

  • Chart extension — the codegen_database_date_bin polyfill needed by the ledger chart function builders.

Plugins vs extensions

If a plugin’s effects are fully contained within one factory, you don’t need an extension — just list the plugin in __plugins__:

from codegen_database.factory import CodegenDatabaseLedger
from codegen_database.plugins.ledger import DoubleEntryPlugin


class Journal(Base):
    __tablename__ = "journal"
    __table_args__ = {"schema": "finance"}
    __factory__ = CodegenDatabaseLedger
    __plugins__ = [DoubleEntryPlugin()]

    account = Column(String, nullable=False)

Extensions are for cross-cutting concerns that don’t belong to any single factory:

  • ``configure_metadata()`` — registers roles, grants, or schemas at the metadata level, outside any factory.

  • ``configure_alembic()`` — installs custom Alembic renderers or rewriters that apply globally, not per table.

  • ``register_cli(app)`` — adds subcommands to the codegen_database CLI.

  • ``validate()`` — runs cross-extension dependency checks after all extensions are resolved.

Quick start

Register an extension on your CodegenDatabaseConfig and attach it to the metadata that your CodegenDatabaseBase subclass uses:

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

config = CodegenDatabaseConfig()
config.use(ChartExtension())
metadata.info["codegen_database_config"] = config

Every model on Base now inherits the extension automatically. In env.py, pass the same config to configure_metadata():

from codegen_database.alembic.register import (
    configure_metadata,
)

configure_metadata(metadata, config)

Extension hooks

CodegenDatabaseExtension provides five hooks. Override only the ones you need — every hook is a no-op by default.

plugins()

Return a list of Plugin instances that are prepended to every factory’s plugin list.

configure_metadata(metadata)

Register roles, grants, schemas, or other metadata-level objects. Called by configure_metadata().

configure_alembic()

Register custom Alembic renderers or rewriters. Called by alembic_hook().

register_cli(app)

Add subcommands to the codegen_database CLI.

validate(registered_names)

Check that required peer extensions are present. Called after all extensions are loaded.

Inter-extension dependencies

Declare dependencies using the depends_on class variable:

from dataclasses import dataclass
from typing import ClassVar

from codegen_database.extension import CodegenDatabaseExtension

@dataclass
class MyExtension(CodegenDatabaseExtension):
    name: str = "my-ext"
    depends_on: ClassVar[list[str]] = ["another-extension"]

codegen_database validates that all declared dependencies are present when extensions are resolved. A CodegenDatabaseValidationError is raised if any are missing.

Entry point discovery

Third-party packages can register extensions via the codegen_database.ext entry point group in pyproject.toml:

[project.entry-points."codegen_database.ext"]
nanoid = "codegen_database_nanoid:NanoIDExtension"

Discovered extensions are automatically loaded unless auto_discover=False is set on the config. Manually registered extensions take precedence over discovered ones with the same name.

Writing an extension

Here are two example tiers, from simple to complex.

Column-level extension (NanoID PK)

An extension that contributes a single plugin to replace the default serial PK with a NanoID:

from dataclasses import dataclass

from codegen_database.extension import CodegenDatabaseExtension
from codegen_database.plugin import Plugin


class NanoIDPKPlugin(Plugin):
    """Replace serial PK with a NanoID column."""
    # ... plugin implementation ...


@dataclass
class NanoIDExtension(CodegenDatabaseExtension):
    name: str = "nanoid"

    def plugins(self) -> list[Plugin]:
        return [NanoIDPKPlugin()]

Composite extension (audit trail)

An extension that bundles multiple plugins — a shadow table and a trigger that writes to it:

from dataclasses import dataclass

from codegen_database.extension import CodegenDatabaseExtension
from codegen_database.plugin import Plugin


@dataclass
class AuditExtension(CodegenDatabaseExtension):
    name: str = "audit"

    def plugins(self) -> list[Plugin]:
        return [
            ShadowTablePlugin(),
            ShadowTriggerPlugin(),
        ]

pg_cron extension

PGCronExtension wires declarative cron job scheduling into the codegen_database lifecycle. Register it on your config to enable:

  • Automatic CREATE EXTENSION IF NOT EXISTS pg_cron CASCADE in Alembic-generated migrations when pg_cron is absent.

  • Diffing of declared CronJob objects against the live cron.job table during alembic revision --autogenerate, emitting op.execute("SELECT cron.schedule(...)") for new or changed jobs.

  • A clear error at factory-construction time if a plugin requests pg_cron scheduling without the extension being registered (e.g. forgetting config.use(PGCronExtension())).

How it works

PGCronExtension.configure_metadata() registers a PGExtension for pg_cron in metadata. The Alembic comparator queries pg_extension and emits CREATE EXTENSION IF NOT EXISTS pg_cron CASCADE if the extension is missing.

PGCronExtension.configure_alembic() registers two Alembic comparators:

  1. Extension comparator — queries pg_extension, emits CreateExtensionOp for any declared extension not yet installed. Extensions are never dropped automatically.

  2. Cron comparator — queries cron.job, emits ScheduleCronJobOp for new or changed jobs. Jobs are never dropped automatically; use UnscheduleCronJobOp in a hand-written migration to remove one.

Both registrations are idempotent — safe to call multiple times from env.py.

Minimal example

# models.py
from sqlalchemy import Column, Integer, MetaData, Numeric

from codegen_database import CodegenDatabaseBase, construct_naming_conventions_dict
from codegen_database.config import CodegenDatabaseConfig
from codegen_database.ext.pg_cron import PGCronExtension
from codegen_database.ext.refresh.plugin import IncrementalRefreshPlugin

metadata = MetaData(
    naming_convention=construct_naming_conventions_dict(),
)

config = CodegenDatabaseConfig()
config.use(PGCronExtension())
metadata.info["codegen_database_config"] = config


class Base(CodegenDatabaseBase):
    metadata = metadata


class Orders(Base):
    __tablename__ = "orders"
    __table_args__ = {"schema": "app"}
    __plugins__ = [
        IncrementalRefreshPlugin(
            body=\"\"\"
INSERT INTO app.orders_daily_agg
    (customer_id, day, total)
SELECT customer_id,
       date_trunc('day', created_at),
       SUM(amount)
FROM app.orders_raw
WHERE created_at >= _since
GROUP BY 1, 2
ON CONFLICT (customer_id, day)
DO UPDATE SET
    total = orders_daily_agg.total + EXCLUDED.total;
\"\"\",
            fallback_sql=(
                "SELECT MIN(day::timestamptz)"
                " - INTERVAL '1 second'"
                " FROM app.orders_daily_agg"
            ),
            schedule="*/5 * * * *",
        ),
    ]

    customer_id = Column(Integer, nullable=False)
    amount = Column(Numeric, nullable=False)

Alembic will then emit, in order:

  1. CREATE EXTENSION IF NOT EXISTS pg_cron CASCADE

  2. CREATE FUNCTION app.app_orders_refresh(...)

  3. SELECT cron.schedule('app.app_orders_refresh', ...)

Pass config to configure_metadata(metadata, config) in env.py — see Setting up a new project for the full Alembic wiring.

Registering extensions without cron

Some extensions are registered for you: configure_metadata() adds DEFAULT_PG_EXTENSIONS (currently pg_trgm, which backs the % trigram operator that fuzzy text search and trigram_indexes() rely on) to every metadata wired through the standard env.py hooks. A manual registration of the same extension name takes precedence and is never duplicated.

You can also declare arbitrary PostgreSQL extensions in metadata without using PGCronExtension. The register_pg_extension() function and PGExtension dataclass are available standalone, but you must call register_pg_extension_alembic_events() yourself in env.py (already done if you call alembic_hook()):

from codegen_database.pg_extension import (
    PGExtension,
    register_pg_extension,
    register_pg_extension_alembic_events,
)

register_pg_extension(metadata, PGExtension("pg_trgm"))

# in env.py:
register_pg_extension_alembic_events()

PGCronExtension.configure_alembic() calls both register_cron_alembic_events() and register_pg_extension_alembic_events() for you — prefer the extension if you are using pg_cron.

Removing a cron job

The comparator never removes jobs automatically to avoid accidentally dropping externally-managed pg_cron jobs. To remove a job, write a migration by hand:

from codegen_database.ext.cron.compare import UnscheduleCronJobOp

def upgrade():
    op.execute(
        UnscheduleCronJobOp(
            name="app.app_orders_refresh"
        ).to_sql()[0]
    )

def downgrade():
    pass

Validation

Plugins that use pg_cron scheduling call assert_pg_extension_declared() internally. If PGCronExtension is not registered on the config, constructing a factory with IncrementalRefreshPlugin(schedule=...) raises CodegenDatabaseValidationError immediately with a message explaining what to do.

TemporalPlugin works the same way for btree_gist: it calls register_pg_extension() automatically, so btree_gist is always declared in metadata when the plugin is used. No manual registration is needed.

PostGIS extension

PostGISExtension groups the PostGIS suite. It exposes the PostGIS column types (currently STDADDR, the stdaddr composite) and declares the PostgreSQL extensions those types need so Alembic can create them.

Extension-owned tables are protected from autogenerate

Installing postgis creates relations that the extension owns but your metadata never declares — most notably the spatial_ref_sys table. Left alone, Alembic autogenerate sees a table in the database with no matching model and emits DROP TABLE spatial_ref_sys, which fails outright:

cannot drop table spatial_ref_sys because extension postgis requires it

codegen_database’s process_revision_directives hook (wired into the standard env.py) prevents this: before sorting migration ops it queries pg_depend for every relation owned by an installed extension and drops any CreateTableOp / DropTableOp that targets one. This is automatic and applies to any extension, not just PostGIS, so you never need to move these tables into a separate schema or hand-edit generated migrations. Online autogenerate only — offline runs have no database to inspect, so nothing is filtered.

Chart extension

ChartExtension installs the codegen_database_date_bin polyfill used by the ledger chart function builders. Native date_bin rejects intervals of a month or larger because months have no fixed length; the polyfill handles any interval ('15 minutes', '1 day', '1 month', '3 months', '1 year') by switching between native date_bin and calendar arithmetic.

Register it on your config whenever you use construct_ledger_chart_function(), construct_ledger_rollup_chart_function(), or construct_double_entry_chart_function().

How it works

ChartExtension.configure_metadata():

  1. Resolves the schema: the schema argument when set, otherwise the config’s utility_schema (default "codegen_database") — so the polyfill shares the one codegen_database utility schema with every other codegen_database-managed object.

  2. Ensures that schema is added to metadata.info["schemas"] so Alembic autogenerate creates it.

  3. Registers codegen_database_date_bin in it.

  4. Stores the schema name under metadata.info["codegen_database_chart_schema"] so the chart function builders can resolve the polyfill without a user-supplied override.

Without this extension, the chart function builders raise CodegenDatabaseValidationError at factory-construction time with a message pointing at config.use(ChartExtension()).

Minimal example

from sqlalchemy import MetaData

from codegen_database import (
    CodegenDatabaseBase,
    CodegenDatabaseFunctionMixin,
    construct_naming_conventions_dict,
)
from codegen_database.config import CodegenDatabaseConfig
from codegen_database.ext.chart import ChartExtension
from codegen_database.ext.ledger import construct_ledger_chart_function

metadata = MetaData(
    naming_convention=construct_naming_conventions_dict(),
)

config = CodegenDatabaseConfig()
config.use(ChartExtension())
metadata.info["codegen_database_config"] = config


class Base(CodegenDatabaseBase):
    metadata = metadata


class StockChart(CodegenDatabaseFunctionMixin, Base):
    __table_args__ = {"schema": "inventory"}
    __funcspec__ = construct_ledger_chart_function(
        StockMovements,
        name="stock_chart",
        dimensions=["warehouse", "sku"],
        period_default="1 day",
    )

Pass config to configure_metadata(metadata, config) in env.py — see Setting up a new project for the full Alembic wiring.

Alembic will emit, in order:

  1. CREATE SCHEMA IF NOT EXISTS codegen_database

  2. CREATE FUNCTION codegen_database.codegen_database_date_bin(...)

  3. CREATE FUNCTION inventory.stock_chart(...) (which calls codegen_database.codegen_database_date_bin internally)

Changing the polyfill schema

By default the polyfill lives in the config’s utility_schema (codegen_database) alongside every other codegen_database-managed object. To move just the chart polyfill elsewhere, pass an explicit schema:

config.use(ChartExtension(schema="reporting"))

To move all codegen_database utility objects together, set CodegenDatabaseConfig(utility_schema="...") instead and leave ChartExtension() at its default.

The ledger chart function builders read metadata.info["codegen_database_chart_schema"] and qualify calls to codegen_database_date_bin accordingly. You can also point the builders at an existing helper by passing date_bin_schema and date_bin_name explicitly — skip the extension entirely in that case.

See Ledger tables for the full chart function walkthrough, including the split_by kwarg and the double-entry helper.