Source code for codegen_database.ext.chart.extension

from dataclasses import dataclass
from typing import TYPE_CHECKING

from sqlalchemy_declarative_extensions import register_function

from codegen_database import CodegenDatabaseExtension
from codegen_database.config import resolve_config
from codegen_database.ext.chart import (
    CHART_SCHEMA_KEY,
)
from codegen_database.ext.chart.date_bin import construct_date_bin_function

if TYPE_CHECKING:
    from sqlalchemy import MetaData




@dataclass(frozen=True, eq=True)
class ChartExtension(CodegenDatabaseExtension):
    name: str = "codegen_database-chart"
    schema: str | None = None



    def configure_metadata(self, metadata: MetaData) -> None:
        schema = self.schema or resolve_config(metadata).utility_schema
        metadata.info[CHART_SCHEMA_KEY] = schema

        register_function(
            metadata,
            construct_date_bin_function().construct_function(schema),
        )