"""Imperative and declarative PostgreSQL view registration."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, cast
from sqlalchemy import Column, MetaData, Table
from sqlalchemy import types as sa_types
from sqlalchemy.orm import registry
from sqlalchemy_declarative_extensions import View
from sqlalchemy_declarative_extensions import (
register_view as decl_register_view,
)
from codegen_database.errors import CodegenDatabaseValidationError
from codegen_database.utils.query import compile_query
if TYPE_CHECKING:
from sqlalchemy import Select
from sqlalchemy.sql.type_api import TypeEngine
[docs]
@dataclass(frozen=True)
class ViewOptions:
materialized: bool = False
def _column_type(column: object) -> TypeEngine:
return getattr(column, "type", sa_types.NullType())
def _construct_view_proxy_table(
name: str,
schema: str,
query: Select,
) -> Table:
columns = [
Column(
column.key,
_column_type(column),
primary_key=getattr(column, "primary_key", False),
)
for column in query.selected_columns
if column.key is not None
]
if not columns:
msg = f"{schema}.{name}: view query must select at least one column."
raise CodegenDatabaseValidationError(msg)
return Table(name, MetaData(), *columns, schema=schema)
def register_view(
name: str,
schema: str,
metadata: MetaData,
query: Select,
*,
materialized: bool,
) -> tuple[View, Table]:
proxy = _construct_view_proxy_table(name, schema, query)
view = View(
name,
compile_query(query),
schema=schema,
materialized=materialized,
)
decl_register_view(metadata, view)
return view, proxy
def register_view_imperatively(
name: str,
schema: str,
metadata: MetaData,
query: Select,
*,
materialized: bool,
primary_key: tuple[str, ...] | None = None,
) -> tuple[View, type[object], type[object]]:
mapper_args = (
{"primary_key": primary_key} if primary_key is not None else {}
)
view_class = cast(
"type[object]",
type(
name,
(),
{
"__module__": __name__,
"__tablename__": name,
"__table_args__": {"schema": schema},
"__query__": query,
"__options__": ViewOptions(materialized=materialized),
"__mapper_args__": mapper_args,
"metadata": metadata,
"registry": registry(),
},
),
)
view, orm_class = _register_view_class(view_class)
return view, view_class, orm_class
def _schema_from_table_args(class_: type) -> str:
raw = getattr(class_, "__table_args__", None)
if isinstance(raw, dict):
schema = raw.get("schema")
elif isinstance(raw, tuple) and raw and isinstance(raw[-1], dict):
schema = raw[-1].get("schema")
else:
schema = None
if schema is None:
msg = f"{class_.__name__}: must specify a schema in __table_args__."
raise CodegenDatabaseValidationError(msg)
return schema
[docs]
class CodegenDatabasePlainView:
"""Register a view with a detached, queryable proxy table."""
def __init__(
self,
name: str,
schema: str,
metadata: MetaData,
query: Select,
) -> None:
self.view, self.table = register_view(
name,
schema,
metadata,
query,
materialized=False,
)
self.name = name
self.schema = schema
self.metadata = metadata
[docs]
class CodegenDatabaseMaterializedView:
"""Register a materialized view with a detached proxy table."""
def __init__(
self,
name: str,
schema: str,
metadata: MetaData,
query: Select,
) -> None:
self.view, self.table = register_view(
name,
schema,
metadata,
query,
materialized=True,
)
self.name = name
self.schema = schema
self.metadata = metadata
def _register_view_class(class_: type) -> tuple[View, type]:
name = class_.__dict__.get("__tablename__")
if name is None:
msg = f"{class_.__name__}: must define __tablename__."
raise CodegenDatabaseValidationError(msg)
query = class_.__dict__["__query__"]
schema = _schema_from_table_args(class_)
options = class_.__dict__.get("__options__", ViewOptions())
mapper_args = dict(class_.__dict__.get("__mapper_args__", {}))
explicit_primary_key = mapper_args.get("primary_key")
if explicit_primary_key is not None and not explicit_primary_key:
msg = f"{schema}.{name}: primary_key must not be empty."
raise CodegenDatabaseValidationError(msg)
view, proxy = register_view(
name,
schema,
class_.metadata,
query,
materialized=options.materialized,
)
if explicit_primary_key is not None:
mapper_args["primary_key"] = [
proxy.c[
value
if isinstance(value, str)
else getattr(value, "key", value)
]
for value in explicit_primary_key
]
elif not proxy.primary_key:
id_column = proxy.c.get("id")
mapper_args["primary_key"] = [
id_column if id_column is not None else next(iter(proxy.c))
]
class_.registry.map_imperatively(class_, proxy, **mapper_args)
class_.view = view
class_.table = proxy
return view, class_