lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/profiling/postgres.py

@@ -0,0 +1,211 @@
+"""Postgres profiler: catalog reflection (SQLAlchemy) plus a sampled, read-only SELECT per table.
+
+Declared facts come from the catalog: column types (mapped to Attribute types with a fixed
+table, design 9.1), nullability, primary keys, single-column unique constraints, foreign keys
+and comments. Values come from ``SELECT ... LIMIT sample_rows`` in a read-only transaction
+with a statement timeout. Row counts are the planner estimate unless ``exact_row_count``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from sqlalchemy import Connection, column, func, inspect, select, table, text
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+from sqlalchemy.types import (
+    ARRAY,
+    JSON,
+    Boolean,
+    Date,
+    DateTime,
+    Enum,
+    Float,
+    Integer,
+    Interval,
+    Numeric,
+    String,
+    Text,
+    TypeEngine,
+    Uuid,
+)
+
+from lorebase_core.infer import ColumnProfile, DeclaredForeignKey, TableProfile
+from lorebase_core.types import AttributeType, TypeSpec
+from lorebase_store.settings import PostgresProfilingSettings
+
+DEFAULT_SCHEMA = "public"
+_ESTIMATE = text(
+    "SELECT c.reltuples::bigint FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace "
+    "WHERE n.nspname = :schema AND c.relname = :table"
+)
+
+
+class PostgresProfileError(ValueError):
+    """The requested tables cannot be profiled."""
+
+
+@dataclass(frozen=True)
+class _Column:
+    name: str
+    spec: TypeSpec | None
+    max_length: int | None
+    nullable: bool
+    unique: bool
+    comment: str | None
+
+
+@dataclass(frozen=True)
+class _Table:
+    name: str
+    columns: tuple[_Column, ...]
+    primary_key: tuple[str, ...]
+    foreign_keys: tuple[DeclaredForeignKey, ...]
+    comment: str | None
+
+
+def type_spec(sql_type: TypeEngine[Any]) -> tuple[TypeSpec | None, int | None]:
+    """Attribute type and declared max length of a catalog type; ``(None, None)`` when the type
+    has no fixed mapping (inference then detects it from values)."""
+    if isinstance(sql_type, Boolean):
+        return TypeSpec.of(AttributeType.BOOLEAN), None
+    if isinstance(sql_type, Integer):
+        return TypeSpec.of(AttributeType.INTEGER), None
+    if isinstance(sql_type, Float):
+        return TypeSpec.of(AttributeType.FLOAT), None
+    if isinstance(sql_type, Numeric):
+        precision, scale = sql_type.precision, sql_type.scale
+        if precision is not None and scale is not None and 0 <= scale <= precision:
+            return TypeSpec.decimal(precision, scale), None
+        return None, None
+    if isinstance(sql_type, DateTime):
+        return TypeSpec.of(AttributeType.TIMESTAMP), None
+    if isinstance(sql_type, Date):
+        return TypeSpec.of(AttributeType.DATE), None
+    if isinstance(sql_type, Interval):
+        return TypeSpec.of(AttributeType.DURATION), None
+    if isinstance(sql_type, JSON | postgresql.JSONB):
+        return TypeSpec.of(AttributeType.JSON), None
+    if isinstance(sql_type, Uuid | postgresql.UUID):
+        return TypeSpec.of(AttributeType.STRING), None
+    if isinstance(sql_type, Enum):
+        values = tuple(str(v) for v in sql_type.enums)
+        return (TypeSpec.enum(*values), None) if values else (None, None)
+    if isinstance(sql_type, Text):
+        return TypeSpec.of(AttributeType.TEXT), None
+    if isinstance(sql_type, String):
+        return TypeSpec.of(AttributeType.STRING), sql_type.length
+    if isinstance(sql_type, ARRAY):
+        element, _ = type_spec(sql_type.item_type)
+        if element is not None and element.name not in (AttributeType.ARRAY, AttributeType.JSON):
+            return TypeSpec.array(element), None
+        return None, None
+    return None, None
+
+
+def _reflect(
+    connection: Connection, schema: str, names: Sequence[str] | None, max_tables: int
+) -> list[_Table]:
+    inspector = inspect(connection)
+    available = sorted(inspector.get_table_names(schema=schema))
+    if names:
+        missing = sorted(set(names) - set(available))
+        if missing:
+            raise PostgresProfileError(f"unknown tables in schema {schema!r}: {', '.join(missing)}")
+        chosen = [n for n in available if n in set(names)]
+    else:
+        chosen = available
+    if len(chosen) > max_tables:
+        raise PostgresProfileError(
+            f"{len(chosen)} tables requested, more than the {max_tables} allowed"
+        )
+    tables = []
+    for name in chosen:
+        primary = tuple(
+            inspector.get_pk_constraint(name, schema=schema).get("constrained_columns") or ()
+        )
+        unique = {
+            u["column_names"][0]
+            for u in inspector.get_unique_constraints(name, schema=schema)
+            if len(u["column_names"]) == 1
+        }
+        columns = []
+        for info in inspector.get_columns(name, schema=schema):
+            spec, max_length = type_spec(info["type"])
+            columns.append(
+                _Column(
+                    name=info["name"],
+                    spec=spec,
+                    max_length=max_length,
+                    nullable=bool(info.get("nullable", True)) and info["name"] not in primary,
+                    unique=info["name"] in unique,
+                    comment=info.get("comment"),
+                )
+            )
+        foreign_keys = tuple(
+            DeclaredForeignKey(
+                tuple(fk["constrained_columns"]),
+                fk["referred_table"],
+                tuple(fk["referred_columns"]),
+            )
+            for fk in inspector.get_foreign_keys(name, schema=schema)
+            if fk.get("referred_schema") in (None, schema)
+        )
+        comment = inspector.get_table_comment(name, schema=schema).get("text")
+        tables.append(_Table(name, tuple(columns), primary, foreign_keys, comment))
+    return tables
+
+
+async def _sample(
+    connection: AsyncConnection, schema: str, found: _Table, settings: PostgresProfilingSettings
+) -> TableProfile:
+    source = table(found.name, *(column(c.name) for c in found.columns), schema=schema)
+    query = select(*source.columns).limit(settings.sample_rows)
+    rows = (await connection.execute(query)).all()
+    count: int | None
+    if settings.exact_row_count:
+        count = (await connection.execute(select(func.count()).select_from(source))).scalar_one()
+    else:
+        estimate = (
+            await connection.execute(_ESTIMATE, {"schema": schema, "table": found.name})
+        ).scalar()
+        count = int(estimate) if estimate is not None and estimate >= 0 else None
+    columns = tuple(
+        ColumnProfile(
+            name=c.name,
+            values=[row[index] for row in rows],
+            declared_type=c.spec,
+            declared_max_length=c.max_length,
+            nullable=c.nullable,
+            unique=c.unique or None,
+            comment=c.comment,
+        )
+        for index, c in enumerate(found.columns)
+    )
+    return TableProfile(
+        name=found.name,
+        columns=columns,
+        row_count=max(count or 0, len(rows)),
+        primary_key=found.primary_key or None,
+        foreign_keys=found.foreign_keys,
+        comment=found.comment,
+    )
+
+
+async def profile_postgres(
+    engine: AsyncEngine,
+    settings: PostgresProfilingSettings,
+    *,
+    schema: str = DEFAULT_SCHEMA,
+    tables: Sequence[str] | None = None,
+) -> list[TableProfile]:
+    """Profiles of ``tables`` (default: every table) of ``schema``, read-only."""
+    async with engine.connect() as connection, connection.begin():
+        await connection.execute(text("SET TRANSACTION READ ONLY"))
+        await connection.execute(
+            text(f"SET LOCAL statement_timeout = {int(settings.statement_timeout_ms)}")
+        )
+        found = await connection.run_sync(_reflect, schema, tables, settings.max_tables)
+        return [await _sample(connection, schema, t, settings) for t in found]

lorebase_store/profiling/record_files.py

@@ -0,0 +1,127 @@
+"""Profilers of uploaded record files: JSON documents, JSON Lines and Parquet files.
+
+The bytes are already bounded by the caller (an upload). A JSON document is an array of
+objects, or holds one at ``records_path`` (dotted keys); JSON Lines hold one object per line
+(blank lines skipped); both are UTF-8, with or without a byte order mark. Parquet needs
+pyarrow (an optional extra of the installing package; ``ParquetUnavailableError`` without it)
+and is decoded ``batch_rows`` rows at a time. Values keep their JSON or Arrow types and nested
+objects stay JSON values; the table is named after the file (``delimited.table_name``).
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator, Mapping
+from typing import Any
+
+from lorebase_core.infer import TableProfile
+from lorebase_store.profiling.delimited import table_name
+from lorebase_store.profiling.records import RecordProfileError, RecordProfiler
+from lorebase_store.settings import RecordProfilingSettings
+
+
+class ParquetUnavailableError(RecordProfileError):
+    """Parquet files cannot be read: pyarrow is not installed."""
+
+
+def _text(data: bytes, name: str) -> str:
+    try:
+        return data.decode("utf-8-sig")
+    except UnicodeDecodeError:
+        raise RecordProfileError(f"{name}: not UTF-8 text") from None
+
+
+def _add(profiler: RecordProfiler, row: Mapping[str, Any], name: str) -> None:
+    try:
+        profiler.add(row)
+    except RecordProfileError as error:
+        raise RecordProfileError(f"{name}: {error}") from None
+
+
+def _profile(
+    rows: Iterator[Mapping[str, Any]], name: str, settings: RecordProfilingSettings
+) -> TableProfile:
+    profiler = RecordProfiler(settings)
+    for row in rows:
+        _add(profiler, row, name)
+    return profiler.profile(table_name(name))
+
+
+def _records_at(document: Any, path: str | None, name: str) -> list[Any]:
+    """The array of records of a JSON document: the document itself, or the array at the dotted
+    key ``path`` (``data.items``); null there is no records."""
+    node = document
+    if path is not None and not isinstance(document, list):
+        for part in path.split("."):
+            if not isinstance(node, dict) or part not in node:
+                raise RecordProfileError(f"{name}: no {path!r} in the document")
+            node = node[part]
+        if node is None:
+            return []
+    if not isinstance(node, list):
+        what = "the document" if path is None else repr(path)
+        hint = "" if path is not None else " (give records_path)"
+        raise RecordProfileError(f"{name}: {what} is not an array of records{hint}")
+    return node
+
+
+def profile_json_bytes(
+    data: bytes, name: str, settings: RecordProfilingSettings, records_path: str | None = None
+) -> TableProfile:
+    """Profile a JSON document holding an array of objects, or the array at ``records_path``
+    (dotted keys) of an object."""
+    try:
+        document = json.loads(_text(data, name))
+    except ValueError as error:
+        raise RecordProfileError(f"{name}: not valid JSON ({error})") from None
+    document = _records_at(document, records_path, name)
+
+    def rows() -> Iterator[Mapping[str, Any]]:
+        for number, record in enumerate(document, start=1):
+            if not isinstance(record, dict):
+                raise RecordProfileError(f"{name}: record {number} is not an object")
+            yield record
+
+    return _profile(rows(), name, settings)
+
+
+def profile_jsonl_bytes(data: bytes, name: str, settings: RecordProfilingSettings) -> TableProfile:
+    """Profile JSON Lines: one object per line."""
+    text = _text(data, name)
+
+    def rows() -> Iterator[Mapping[str, Any]]:
+        for number, line in enumerate(text.splitlines(), start=1):
+            if not line.strip():
+                continue
+            try:
+                record = json.loads(line)
+            except ValueError:
+                raise RecordProfileError(f"{name}: line {number} is not valid JSON") from None
+            if not isinstance(record, dict):
+                raise RecordProfileError(f"{name}: line {number} is not an object")
+            yield record
+
+    return _profile(rows(), name, settings)
+
+
+def profile_parquet_bytes(
+    data: bytes, name: str, settings: RecordProfilingSettings
+) -> TableProfile:
+    """Profile a Parquet file; ``ParquetUnavailableError`` without pyarrow."""
+    try:
+        import pyarrow  # type: ignore[import-untyped]
+        import pyarrow.parquet as parquet  # type: ignore[import-untyped]
+    except ImportError:
+        raise ParquetUnavailableError(f"{name}: Parquet files need pyarrow") from None
+    profiler = RecordProfiler(settings)
+    try:
+        file = parquet.ParquetFile(pyarrow.BufferReader(data))
+        columns = tuple(str(field.name) for field in file.schema_arrow)
+        for batch in file.iter_batches(batch_size=settings.batch_rows):
+            for row in batch.to_pylist():
+                _add(profiler, row, name)
+    except RecordProfileError:
+        raise
+    except (ValueError, OSError) as error:  # pyarrow.ArrowInvalid, truncated files
+        raise RecordProfileError(f"{name}: not a readable Parquet file ({error})") from None
+    return profiler.profile(table_name(name), columns)

lorebase_store/profiling/records.py

@@ -0,0 +1,92 @@
+"""Profiler of rows read as records: JSON documents, JSON Lines, Parquet files, JSON APIs.
+
+Rows are mappings from column name to native value (``str``, ``int``, ``float``, ``Decimal``,
+``date``, ``datetime``, ``bool``, nested ``dict`` or ``list``, ``None``). Columns are the keys
+seen in any row, in first-seen order; a row without a column holds null there. Every row is
+counted and ``sample_rows`` of them are kept, the first ones (``head``) or a uniform seeded
+sample (``reservoir``, kept in read order). Rows can be pushed one at a time
+(``RecordProfiler``), so an asynchronous reader never holds more than the sample.
+"""
+
+from __future__ import annotations
+
+import random
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Literal
+
+from lorebase_core.infer import ColumnProfile, TableProfile
+from lorebase_store.settings import RecordProfilingSettings
+
+
+class RecordProfileError(ValueError):
+    """The rows cannot be profiled as one table."""
+
+
+class RowSampler[T]:
+    """Counts rows and keeps ``size`` of them: the first ones, or a seeded uniform sample."""
+
+    def __init__(self, size: int, sampling: Literal["head", "reservoir"], seed: int) -> None:
+        self._size = size
+        self._head = sampling == "head"
+        self._random = random.Random(seed)
+        self._kept: list[tuple[int, T]] = []
+        self.total = 0
+
+    def add(self, row: T) -> None:
+        index = self.total
+        self.total += 1
+        if len(self._kept) < self._size:
+            self._kept.append((index, row))
+        elif not self._head:
+            slot = self._random.randint(0, index)
+            if slot < self._size:
+                self._kept[slot] = (index, row)
+
+    def kept(self) -> list[T]:
+        """The kept rows, in read order."""
+        return [row for _, row in sorted(self._kept, key=lambda item: item[0])]
+
+
+class RecordProfiler:
+    """Builds the ``TableProfile`` of rows pushed with ``add``."""
+
+    def __init__(self, settings: RecordProfilingSettings) -> None:
+        self._settings = settings
+        self._sampler: RowSampler[Mapping[str, object]] = RowSampler(
+            settings.sample_rows, settings.sampling, settings.seed
+        )
+        self._columns: dict[str, None] = {}
+
+    def add(self, row: Mapping[str, object]) -> None:
+        for column in row:
+            if column not in self._columns:
+                if len(self._columns) >= self._settings.max_columns:
+                    raise RecordProfileError(
+                        f"more than the {self._settings.max_columns} columns allowed"
+                    )
+                self._columns[column] = None
+        self._sampler.add(row)
+
+    def profile(self, name: str, columns: Sequence[str] = ()) -> TableProfile:
+        """The profile of table ``name``; ``columns`` (a declared header) come first."""
+        names = list(dict.fromkeys([*columns, *self._columns]))
+        if len(names) > self._settings.max_columns:
+            raise RecordProfileError(
+                f"{name}: more than the {self._settings.max_columns} columns allowed"
+            )
+        kept = self._sampler.kept()
+        return TableProfile(
+            name,
+            tuple(ColumnProfile(column, [row.get(column) for row in kept]) for column in names),
+            row_count=self._sampler.total,
+        )
+
+
+def profile_records(
+    rows: Iterable[Mapping[str, object]], name: str, settings: RecordProfilingSettings
+) -> TableProfile:
+    """Profile rows as table ``name``."""
+    profiler = RecordProfiler(settings)
+    for row in rows:
+        profiler.add(row)
+    return profiler.profile(name)