pyoq-sql 1.0.2__py3-none-any.whl

pyoq/schema/postgres/reflection.py

@@ -0,0 +1,466 @@
+"""PostgreSQL metadata reflection into the immutable schema model."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeVar
+
+from pyoq.config import DatabaseDialect
+from pyoq.errors import SchemaInspectionError
+from pyoq.schema.models import (
+    Capability,
+    CapabilitySet,
+    Catalog,
+    CheckConstraint,
+    Column,
+    Domain,
+    EnumType,
+    GeneratedValue,
+    Identifier,
+    Index,
+    IndexTerm,
+    Key,
+    KeyKind,
+    ObjectReference,
+    ParameterMode,
+    Relation,
+    Routine,
+    RoutineKind,
+    RoutineParameter,
+    Schema,
+    SchemaSnapshot,
+    SqlExpression,
+    SqlType,
+    Table,
+    TypeKind,
+    View,
+)
+from pyoq.schema.postgres import sql
+from pyoq.schema.postgres.records import (
+    ColumnRecord,
+    ConstraintRecord,
+    DomainCheckRecord,
+    DomainRecord,
+    EnumRecord,
+    IndexRecord,
+    ObjectRecord,
+    RoutineRecord,
+    read_column,
+    read_constraint,
+    read_domain,
+    read_domain_check,
+    read_enum,
+    read_index,
+    read_object,
+    read_routine,
+)
+from pyoq.schema.postgres.types import postgres_type
+
+if TYPE_CHECKING:
+    from psycopg import Connection
+
+_ROUTINE_KINDS = {"f": RoutineKind.FUNCTION, "p": RoutineKind.PROCEDURE}
+_PARAMETER_MODES = {
+    "i": ParameterMode.IN,
+    "o": ParameterMode.OUT,
+    "b": ParameterMode.INOUT,
+}
+
+_DOMAIN_TYPE = "d"
+_ENUM_TYPE = "e"
+
+Record = TypeVar("Record")
+GroupKey = TypeVar("GroupKey")
+
+_CAPABILITIES = CapabilitySet(
+    frozenset(
+        {
+            Capability.ARRAYS,
+            Capability.CHECK_CONSTRAINTS,
+            Capability.ENUMS,
+            Capability.EXPRESSION_INDEXES,
+            Capability.GENERATED_COLUMNS,
+            Capability.IDENTITY_COLUMNS,
+            Capability.MATERIALIZED_VIEWS,
+            Capability.PARTIAL_INDEXES,
+            Capability.SCHEMAS,
+            Capability.VIEWS,
+        }
+    )
+)
+
+
+@dataclass(frozen=True, slots=True)
+class PostgresSchemaInspector:
+    def inspect(
+        self,
+        connection: Connection[tuple[object, ...]],
+    ) -> SchemaSnapshot:
+        catalog = _read(connection)
+        schemas = tuple(_schema(name, catalog) for name in sorted(catalog.schema_names))
+        return SchemaSnapshot(
+            DatabaseDialect.POSTGRES,
+            (Catalog(None, schemas),),
+            _CAPABILITIES,
+        )
+
+
+@dataclass(frozen=True, slots=True, eq=False)
+class _Catalog:
+    """Grouped catalog rows. An index, not a value, so it does not compare."""
+
+    objects: tuple[ObjectRecord, ...]
+    enums: dict[str, list[EnumRecord]]
+    domains: dict[str, list[DomainRecord]]
+    domain_checks: dict[tuple[str, str], list[DomainCheckRecord]]
+    routines: dict[str, list[RoutineRecord]]
+    columns: dict[tuple[str, str], list[ColumnRecord]]
+    constraints: dict[tuple[str, str], list[ConstraintRecord]]
+    indexes: dict[tuple[str, str], list[IndexRecord]]
+
+    @property
+    def schema_names(self) -> set[str]:
+        """Every schema that holds anything, not only tables and views.
+
+        A schema of nothing but domains, enums, or routines is still a schema,
+        and leaving it out would lose everything it defines.
+        """
+        return (
+            {record.schema for record in self.objects}
+            | set(self.enums)
+            | set(self.domains)
+            | set(self.routines)
+        )
+
+
+def _read(connection: Connection[tuple[object, ...]]) -> _Catalog:
+    try:
+        objects = tuple(
+            read_object(row) for row in connection.execute(sql.SCHEMA_OBJECTS)
+        )
+        columns = tuple(read_column(row) for row in connection.execute(sql.COLUMNS))
+        constraints = tuple(
+            read_constraint(row) for row in connection.execute(sql.CONSTRAINTS)
+        )
+        indexes = tuple(read_index(row) for row in connection.execute(sql.INDEXES))
+        enums = tuple(read_enum(row) for row in connection.execute(sql.ENUMS))
+        domains = tuple(read_domain(row) for row in connection.execute(sql.DOMAINS))
+        checks = tuple(
+            read_domain_check(row) for row in connection.execute(sql.DOMAIN_CHECKS)
+        )
+        routines = tuple(read_routine(row) for row in connection.execute(sql.ROUTINES))
+    except Exception as error:
+        message = f"unable to inspect PostgreSQL schema: {error}"
+        raise SchemaInspectionError(message) from error
+    return _Catalog(
+        objects,
+        _grouped(enums, lambda record: record.schema),
+        _grouped(domains, lambda record: record.schema),
+        _grouped(checks, lambda record: (record.schema, record.domain)),
+        _grouped(routines, lambda record: record.schema),
+        _grouped(columns, lambda record: (record.schema, record.table)),
+        _grouped(constraints, lambda record: (record.schema, record.table)),
+        _grouped(indexes, lambda record: (record.schema, record.table)),
+    )
+
+
+def _schema(name: str, catalog: _Catalog) -> Schema:
+    objects = tuple(record for record in catalog.objects if record.schema == name)
+    tables = tuple(_table(record, catalog) for record in objects if record.is_table)
+    views = tuple(_view(record, catalog) for record in objects if not record.is_table)
+    return Schema(
+        Identifier(name),
+        tables,
+        views,
+        _enums(catalog, name),
+        _domains(catalog, name),
+        _routines(catalog, name),
+    )
+
+
+def _domains(catalog: _Catalog, schema: str) -> tuple[Domain, ...]:
+    """Every domain the schema defines, with the rules that belong to it."""
+    return tuple(
+        Domain(
+            Identifier(record.name),
+            postgres_type(record.base_type),
+            not record.not_null,
+            None if record.default is None else SqlExpression(record.default),
+            _domain_checks(catalog, schema, record.name),
+            record.comment,
+        )
+        for record in catalog.domains.get(schema, [])
+    )
+
+
+def _domain_checks(
+    catalog: _Catalog,
+    schema: str,
+    domain: str,
+) -> tuple[CheckConstraint, ...]:
+    return tuple(
+        CheckConstraint(SqlExpression(record.definition), Identifier(record.name))
+        for record in catalog.domain_checks.get((schema, domain), [])
+    )
+
+
+def _routines(catalog: _Catalog, schema: str) -> tuple[Routine, ...]:
+    """Every routine the schema keeps, one entry per routine rather than row."""
+    grouped: dict[int, list[RoutineRecord]] = defaultdict(list)
+    for record in catalog.routines.get(schema, []):
+        grouped[record.identifier].append(record)
+    named: dict[str, int] = defaultdict(int)
+    routines: list[Routine] = []
+    for records in grouped.values():
+        first = records[0]
+        named[first.name] += 1
+        routines.append(_routine(records, named[first.name]))
+    return tuple(routines)
+
+
+def _routine(records: list[RoutineRecord], seen: int) -> Routine:
+    first = records[0]
+    kind = _ROUTINE_KINDS[first.kind]
+    defaulted = _defaulted_parameter_positions(records, first.default_count)
+    return Routine(
+        Identifier(first.name),
+        kind,
+        tuple(
+            _parameter(record, record.position in defaulted)
+            for record in records
+            if record.position is not None
+        ),
+        (
+            None
+            if first.returns is None
+            else _routine_type(
+                first.returns,
+                first.return_type_schema,
+                first.return_type_name,
+                first.return_type_kind,
+                first.return_base_type,
+                first.return_is_array,
+            )
+        ),
+        None if seen == 1 else str(seen),
+    )
+
+
+def _defaulted_parameter_positions(
+    records: list[RoutineRecord], default_count: int
+) -> frozenset[int]:
+    inputs = [
+        record.position
+        for record in records
+        if record.position is not None and record.mode in {"i", "b"}
+    ]
+    return frozenset(inputs[-default_count:]) if default_count else frozenset()
+
+
+def _parameter(record: RoutineRecord, has_default: bool) -> RoutineParameter:
+    return RoutineParameter(
+        None if record.parameter_name is None else Identifier(record.parameter_name),
+        _routine_type(
+            record.parameter_type or "text",
+            record.parameter_type_schema,
+            record.parameter_type_name,
+            record.parameter_type_kind,
+            record.parameter_base_type,
+            record.parameter_is_array,
+        ),
+        _PARAMETER_MODES.get(record.mode, ParameterMode.IN),
+        has_default,
+        nullable=not record.parameter_forbids_null,
+    )
+
+
+def _routine_type(
+    declared: str,
+    schema: str | None,
+    name: str | None,
+    kind: str | None,
+    base: str | None,
+    is_array: bool,
+) -> SqlType:
+    """A routine type, retaining named enum and domain identities."""
+    scalar = _routine_scalar_type(declared, schema, name, kind, base)
+    if not is_array:
+        return scalar
+    return SqlType(TypeKind.ARRAY, declared, element_type=scalar)
+
+
+def _routine_scalar_type(
+    declared: str,
+    schema: str | None,
+    name: str | None,
+    kind: str | None,
+    base: str | None,
+) -> SqlType:
+    if kind not in {_ENUM_TYPE, _DOMAIN_TYPE} or schema is None or name is None:
+        return postgres_type(declared)
+    reference = ObjectReference(Identifier(name), Identifier(schema))
+    if kind == _ENUM_TYPE:
+        return SqlType(TypeKind.ENUM, declared, named_type=reference)
+    return SqlType(
+        TypeKind.DOMAIN,
+        declared,
+        element_type=postgres_type(base or "text"),
+        named_type=reference,
+    )
+
+
+def _enums(catalog: _Catalog, schema: str) -> tuple[EnumType, ...]:
+    """Every enum the schema defines, each with its labels in the type's order."""
+    labels: dict[str, list[str]] = {}
+    comments: dict[str, str | None] = {}
+    for record in catalog.enums.get(schema, []):
+        labels.setdefault(record.name, []).append(record.label)
+        comments.setdefault(record.name, record.comment)
+    return tuple(
+        EnumType(Identifier(name), tuple(values), comments[name])
+        for name, values in sorted(labels.items())
+    )
+
+
+def _table(record: ObjectRecord, catalog: _Catalog) -> Table:
+    location = (record.schema, record.name)
+    constraints = catalog.constraints.get(location, [])
+    return Table(
+        Identifier(record.name),
+        _columns(catalog, location),
+        tuple(
+            _key(item, kind)
+            for item in constraints
+            if (kind := item.key_kind) is not None
+        ),
+        tuple(
+            _relation(item, target)
+            for item in constraints
+            if item.kind == "f" and (target := item.target_reference) is not None
+        ),
+        tuple(_index_model(item) for item in catalog.indexes.get(location, [])),
+        tuple(_check(item) for item in constraints if item.kind == "c"),
+        record.comment,
+    )
+
+
+def _view(record: ObjectRecord, catalog: _Catalog) -> View:
+    definition = (
+        None if record.definition is None else SqlExpression(record.definition.strip())
+    )
+    return View(
+        Identifier(record.name),
+        _columns(catalog, (record.schema, record.name)),
+        definition,
+        record.is_materialized,
+        record.comment,
+    )
+
+
+def _columns(catalog: _Catalog, location: tuple[str, str]) -> tuple[Column, ...]:
+    return tuple(_column(record) for record in catalog.columns.get(location, []))
+
+
+def _column(record: ColumnRecord) -> Column:
+    kind = record.generated_kind
+    expression = None if record.default is None else SqlExpression(record.default)
+    generated = None if kind is None else GeneratedValue(kind, expression)
+    return Column(
+        Identifier(record.name),
+        _data_type(record),
+        record.nullable,
+        None if generated is not None else expression,
+        generated,
+        record.comment,
+    )
+
+
+def _data_type(record: ColumnRecord) -> SqlType:
+    """What a column holds, as the type it was declared with rather than as text.
+
+    A domain is its base type with rules attached, so the base is kept as what
+    says which Python type reads it back. An enum is its own type with its own
+    values, so it is named rather than reduced to the string it is stored as.
+    Both keep their name: a schema that forgot it would no longer describe the
+    database it was read from.
+    """
+    if record.type_kind == _ENUM_TYPE:
+        return SqlType(
+            TypeKind.ENUM, record.data_type, named_type=_type_reference(record)
+        )
+    if record.type_kind != _DOMAIN_TYPE or record.domain_base is None:
+        return postgres_type(record.data_type)
+    return SqlType(
+        TypeKind.DOMAIN,
+        record.data_type,
+        element_type=postgres_type(record.domain_base),
+        named_type=_type_reference(record),
+    )
+
+
+def _type_reference(record: ColumnRecord) -> ObjectReference:
+    return ObjectReference(Identifier(record.type_name), Identifier(record.type_schema))
+
+
+def _key(record: ConstraintRecord, kind: KeyKind) -> Key:
+    return Key(
+        kind,
+        tuple(Identifier(name) for name in record.columns),
+        Identifier(record.name),
+    )
+
+
+def _relation(record: ConstraintRecord, target: ObjectReference) -> Relation:
+    return Relation(
+        tuple(Identifier(name) for name in record.columns),
+        target,
+        tuple(Identifier(name) for name in record.target_columns),
+        Identifier(record.name),
+        record.update_action,
+        record.delete_action,
+    )
+
+
+def _check(record: ConstraintRecord) -> CheckConstraint:
+    return CheckConstraint(
+        SqlExpression(record.definition),
+        Identifier(record.name),
+    )
+
+
+def _index_model(record: IndexRecord) -> Index:
+    terms = tuple(
+        _index_term(record, position)
+        for position in range(len(record.attribute_numbers))
+    )
+    predicate = None if record.predicate is None else SqlExpression(record.predicate)
+    return Index(
+        Identifier(record.name),
+        terms,
+        record.unique,
+        predicate,
+        record.method,
+    )
+
+
+def _index_term(record: IndexRecord, position: int) -> IndexTerm:
+    descending = record.descending(position)
+    name = record.column_names[position]
+    if record.attribute_numbers[position] > 0 and name is not None:
+        return IndexTerm(Identifier(name), None, descending)
+    return IndexTerm(None, SqlExpression(record.definitions[position]), descending)
+
+
+def _grouped(
+    records: tuple[Record, ...],
+    key: Callable[[Record], GroupKey],
+) -> dict[GroupKey, list[Record]]:
+    grouped: dict[GroupKey, list[Record]] = defaultdict(list)
+    for record in records:
+        grouped[key(record)].append(record)
+    return grouped
+
+
+__all__ = ("PostgresSchemaInspector",)

pyoq/schema/postgres/source.py

@@ -0,0 +1,30 @@
+"""Configuration-backed PostgreSQL schema source."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from pyoq.config import Configuration, DatabaseDialect
+from pyoq.errors import OperationUnavailableError
+from pyoq.schema.models import SchemaSnapshot
+from pyoq.schema.postgres.connection import PostgresConnectionProvider
+from pyoq.schema.postgres.reflection import PostgresSchemaInspector
+
+
+@dataclass(frozen=True, slots=True)
+class PostgresSchemaSource:
+    connections: PostgresConnectionProvider = field(
+        default_factory=PostgresConnectionProvider
+    )
+    inspector: PostgresSchemaInspector = field(default_factory=PostgresSchemaInspector)
+
+    def load(self, configuration: Configuration) -> SchemaSnapshot:
+        dialect = configuration.active_profile.dialect
+        if dialect is not DatabaseDialect.POSTGRES:
+            message = f"PostgreSQL inspection is not available for {dialect.value}"
+            raise OperationUnavailableError(message)
+        with self.connections.open(configuration) as connection:
+            return self.inspector.inspect(connection)
+
+
+__all__ = ("PostgresSchemaSource",)