Source code for codegen_database.ext.text_range.extension

from __future__ import annotations

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.text_range.function import (
    construct_record_key_function,
)
from codegen_database.pg_type import PGType, register_pg_type

if TYPE_CHECKING:
    from sqlalchemy import MetaData

__all__ = ["TextRangeExtension"]


_RANGE_NAME = "textrange"
_MULTIRANGE_NAME = "textmultirange"


def _constructor_function_globs(schema: str | None) -> list[str]:
    # ``sqlalchemy_declarative_extensions`` normalizes the ``public`` schema
    # away when it computes a function's ``qualified_name`` (see
    # ``sqlalchemy_declarative_extensions.sql.qualify_name``), so matching
    # globs must do the same. Prefixing ``public.`` here would never match the
    # bare ``textrange``/``textmultirange`` constructor functions that
    # PostgreSQL auto-creates alongside the range type, and autogenerate would
    # emit ``DROP FUNCTION`` for them (which fails: the range type depends on
    # those constructors).
    if schema and schema != "public":
        return [f"{schema}.{_RANGE_NAME}*", f"{schema}.{_MULTIRANGE_NAME}*"]

    return [f"{_RANGE_NAME}*", f"{_MULTIRANGE_NAME}*"]


def build_textrange_pg_type(schema: str | None = None) -> PGType:
    if schema:
        create_sql = (
            f"CREATE TYPE {schema}.textrange AS RANGE "
            f"(subtype = text, "
            f"multirange_type_name = {schema}.textmultirange)"
        )

    else:
        create_sql = (
            "CREATE TYPE textrange AS RANGE "
            "(subtype = text, multirange_type_name = textmultirange)"
        )

    return PGType(
        name="textrange",
        create_sql=create_sql,
        schema=schema,
        companions=("textmultirange",),
    )




@dataclass(frozen=True, eq=True)
class TextRangeExtension(CodegenDatabaseExtension):
    name: str = "text_range"
    schema: str | None = None



    def configure_metadata(self, metadata: MetaData) -> None:
        register_pg_type(metadata, build_textrange_pg_type(self.schema))
        register_function(
            metadata,
            construct_record_key_function().construct_function(
                self.schema or resolve_config(metadata).utility_schema
            ),
        )
        metadata.info["functions"].ignore.extend(
            _constructor_function_globs(self.schema)
        )