Source code for codegen_database.ext.postgis

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

from codegen_database.extension import CodegenDatabaseExtension
from codegen_database.pg_extension import (
    PGExtension,
    register_pg_extension,
)
from codegen_database.types.postgis import STDADDR, StdAddr

if TYPE_CHECKING:
    from sqlalchemy import MetaData

__all__ = [
    "STDADDR",
    "PGExtension",
    "PostGISExtension",
    "StdAddr",
    "register_pg_extension",
]

_POSTGIS_EXTENSION = PGExtension("postgis", cascade=True)
_ADDRESS_STANDARDIZER_EXTENSION = PGExtension("address_standardizer")




@dataclass(frozen=True, eq=True)
class PostGISExtension(CodegenDatabaseExtension):
    name: str = "postgis"
    postgis: bool = False
    address_standardizer: bool = True
    geography: bool = False



    def configure_metadata(self, metadata: MetaData) -> None:
        if self.postgis:
            register_pg_extension(metadata, _POSTGIS_EXTENSION)

        if self.address_standardizer:
            register_pg_extension(metadata, _ADDRESS_STANDARDIZER_EXTENSION)

        if self.geography and not self.postgis:
            register_pg_extension(metadata, _POSTGIS_EXTENSION)