contextbase-shared-plugins 0.2.3__py3-none-any.whl

shared_plugins/pg_column_comments.py

@@ -0,0 +1,59 @@
+"""Monkey-patch PostgresClient to emit COMMENT ON COLUMN for dlt column descriptions.
+
+dlt's ``TColumnSchema`` supports a ``description`` field, but the Postgres destination
+does not act on it.  Databricks and Snowflake do (inline ``COMMENT`` in DDL).  Postgres
+uses a separate ``COMMENT ON COLUMN`` statement, so we hook into
+PostgresClient's ``_build_schema_update_sql`` to append COMMENT statements for the
+same tables dlt is updating — regardless of whether those tables have new columns.
+
+ctxb docs reads ``COMMENT ON COLUMN`` metadata and surfaces it in generated schema docs
+that are injected into the agent context.  This closes the loop:
+
+    dlt schema (description) → COMMENT ON COLUMN → ctxb docs → agent prompt
+"""
+
+from __future__ import annotations
+
+from typing import Iterable, List, Tuple
+
+from dlt.common.data_writers.escape import (
+    escape_postgres_identifier,
+    escape_postgres_literal,
+)
+from dlt.common.schema.typing import (
+    TSchemaTables,
+    TTableSchemaColumns,
+)
+from dlt.destinations.impl.postgres.postgres import PostgresClient
+
+_original_build_schema_update_sql = PostgresClient._build_schema_update_sql
+
+
+def _build_schema_update_sql_with_comments(
+    self: PostgresClient,
+    storage_tables: Iterable[Tuple[str, TTableSchemaColumns]],
+) -> Tuple[List[str], TSchemaTables]:
+    storage_tables = list(storage_tables)
+    sql_updates, schema_update = _original_build_schema_update_sql(self, storage_tables)
+
+    dlt_table_names = set(self.schema.dlt_table_names())
+    for table_name, _storage_columns in storage_tables:
+        if table_name in dlt_table_names:
+            continue
+
+        full_table = self.prepare_load_table(table_name)
+        qualified_table = self.sql_client.make_qualified_table_name(table_name)
+        for col in full_table["columns"].values():
+            description = col.get("description")
+            comment_literal = (
+                escape_postgres_literal(description) if description else "NULL"
+            )
+            col_name = escape_postgres_identifier(col["name"])
+            sql_updates.append(
+                f"COMMENT ON COLUMN {qualified_table}.{col_name} IS {comment_literal}"
+            )
+
+    return sql_updates, schema_update
+
+
+PostgresClient._build_schema_update_sql = _build_schema_update_sql_with_comments  # type: ignore[method-assign]

shared_plugins/pyairbyte.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from datetime import datetime, timezone
+from math import isfinite
+from typing import Any, ClassVar, cast
+
+import airbyte as ab
+import dagster as dg
+import sqlalchemy
+from dagster import AssetExecutionContext
+from airbyte._processors.sql.postgres import PostgresSqlProcessor
+from airbyte.caches.postgres import PostgresCache
+from airbyte.progress import ProgressTracker
+from airbyte.records import StreamRecordHandler
+from airbyte_protocol.models import AirbyteRecordMessage
+from pydantic import Field
+from sqlalchemy.engine import make_url
+from sqlalchemy.types import TypeEngine
+
+from shared_types.dagster_binding_plan import DagsterAllPlanBinding
+
+from .env import load_shared_python_settings
+from .exceptions import PluginConfigurationError
+from .naming import dagster_airbyte_sync_asset_key
+from .values import parse_utc_datetime_from_str
+
+_CTX_BINDING_ID_COLUMN = "_ctx_binding_id"
+_CTX_SOURCE_UPDATED_AT_COLUMN = "_ctx_source_updated_at"
+_CTX_BINDING_ID_SCHEMA: dict[str, object] = {"type": "string"}
+# WARNING: This *looks* like it produces a Postgres timestamptz column, but
+# pyairbyte/airbyte currently drops the timezone flag in its DDL path.
+#
+# In airbyte/shared/sql_processor.py, `_create_table_for_loading()`,
+# `_ensure_final_table_exists()`, and `_add_column_to_table()` render SQLAlchemy
+# types via `f"{sql_type}"` / `f"{column_type}"` instead of dialect-aware
+# compilation. For `sqlalchemy.TIMESTAMP(timezone=True)`, `str(...)` is just
+# `"TIMESTAMP"`, so Postgres creates `timestamp without time zone` instead of
+# `TIMESTAMP WITH TIME ZONE`.
+#
+# This means `_ctx_source_updated_at` can be stored incorrectly: psycopg/Postgres
+# will coerce aware datetimes through the session timezone and then drop tz info.
+# The value may therefore vary by session/database timezone instead of preserving a
+# stable absolute instant.
+#
+# We are documenting this intentionally rather than fixing it right now. A future
+# fix should keep the upstream catalog/row-shaping behavior, but override the DDL
+# rendering path so these types compile with the Postgres dialect (for example,
+# `TIMESTAMP(timezone=True)` -> `TIMESTAMP WITH TIME ZONE`) before creating or
+# altering tables.
+_CTX_SOURCE_UPDATED_AT_SCHEMA: dict[str, object] = {
+    "type": ["string", "null"],
+    "format": "date-time",
+}
+
+
+def _coerce_source_updated_at(value: object) -> datetime | None:
+    if value is None:
+        return None
+
+    if isinstance(value, bool):
+        raise RuntimeError("_ctx_source_updated_at does not accept bool values.")
+
+    if isinstance(value, datetime):
+        if value.tzinfo is None or value.utcoffset() is None:
+            return value.replace(tzinfo=timezone.utc)
+        return value.astimezone(timezone.utc)
+
+    if isinstance(value, int):
+        return _coerce_unix_timestamp(value)
+
+    if isinstance(value, float):
+        if not isfinite(value):
+            raise RuntimeError(
+                "_ctx_source_updated_at does not accept non-finite unix timestamps."
+            )
+        return _coerce_unix_timestamp(value)
+
+    if isinstance(value, str):
+        try:
+            return parse_utc_datetime_from_str(value)
+        except ValueError as exc:
+            raise RuntimeError(
+                f"_ctx_source_updated_at must be an ISO-8601 timestamp, got {value!r}."
+            ) from exc
+
+    raise RuntimeError(
+        "_ctx_source_updated_at must be None, a datetime, an ISO-8601 timestamp string, "
+        f"or a unix timestamp, got {type(value).__name__}."
+    )
+
+
+def _coerce_unix_timestamp(value: int | float) -> datetime | None:
+    if value == 0:
+        return None
+    if value < 0:
+        raise RuntimeError(
+            "_ctx_source_updated_at does not accept negative unix timestamps."
+        )
+
+    try:
+        return datetime.fromtimestamp(value, tz=timezone.utc)
+    except (OverflowError, OSError, ValueError) as exc:
+        raise RuntimeError(
+            "_ctx_source_updated_at must be a valid unix timestamp."
+        ) from exc
+
+
+def _build_source_updated_at_fields(
+    stream_names: Sequence[str],
+    source_updated_at_fields: Mapping[str, str | None] | None,
+) -> dict[str, str | None]:
+    return {
+        stream_name: (
+            None
+            if source_updated_at_fields is None
+            else source_updated_at_fields.get(stream_name)
+        )
+        for stream_name in stream_names
+    }
+
+
+def _extract_source_updated_at_value(
+    *,
+    stream_name: str,
+    record_data: Mapping[str, Any],
+    source_updated_at_field: str | None,
+) -> object:
+    if source_updated_at_field is None:
+        return None
+
+    if source_updated_at_field not in record_data:
+        raise RuntimeError(
+            f"PyAirbyte stream '{stream_name}' declared SOURCE_UPDATED_AT_FIELDS[{stream_name!r}] = "
+            f"{source_updated_at_field!r}, but that field was not present in the record payload."
+        )
+
+    return record_data[source_updated_at_field]
+
+
+def _with_ctx_binding_primary_key(
+    *,
+    stream_name: str,
+    source_primary_key: Sequence[Sequence[str]] | None,
+) -> list[list[str]]:
+    if not source_primary_key:
+        raise RuntimeError(
+            f"PyAirbyte stream '{stream_name}' does not define a primary key. "
+            "ContextBase requires source-defined keys so _ctx_binding_id can be prepended safely."
+        )
+
+    normalized_primary_key = [[_CTX_BINDING_ID_COLUMN]]
+    for field_path in source_primary_key:
+        if len(field_path) != 1:
+            raise RuntimeError(
+                f"PyAirbyte stream '{stream_name}' uses nested primary keys {field_path!r}. "
+                "ContextBase only supports top-level primary-key columns."
+            )
+
+        field_name = field_path[0]
+        if field_name == _CTX_BINDING_ID_COLUMN:
+            continue
+        normalized_primary_key.append([field_name])
+
+    return normalized_primary_key
+
+
+def _apply_ctx_catalog_contract(
+    *,
+    source: Any,
+    selected_stream_names: Sequence[str],
+) -> None:
+    selected_stream_name_set = set(selected_stream_names)
+    applied_stream_names: set[str] = set()
+
+    for stream in source.discovered_catalog.streams:
+        if stream.name not in selected_stream_name_set:
+            continue
+
+        json_schema = stream.json_schema
+        properties = json_schema.setdefault("properties", {})
+        if not isinstance(properties, dict):
+            raise RuntimeError(
+                f"PyAirbyte stream '{stream.name}' JSON schema must define a top-level "
+                "properties mapping."
+            )
+
+        properties[_CTX_BINDING_ID_COLUMN] = dict(_CTX_BINDING_ID_SCHEMA)
+        properties[_CTX_SOURCE_UPDATED_AT_COLUMN] = dict(_CTX_SOURCE_UPDATED_AT_SCHEMA)
+        stream.source_defined_primary_key = _with_ctx_binding_primary_key(
+            stream_name=stream.name,
+            source_primary_key=stream.source_defined_primary_key,
+        )
+        applied_stream_names.add(stream.name)
+
+    missing_stream_names = selected_stream_name_set - applied_stream_names
+    if missing_stream_names:
+        missing_stream_names_str = ", ".join(sorted(missing_stream_names))
+        raise RuntimeError(
+            "PyAirbyte selected streams were not found in the discovered catalog: "
+            f"{missing_stream_names_str}."
+        )
+
+
+def _normalize_record_data_for_sql_columns(
+    *,
+    record_data: dict[str, Any],
+    column_definitions: dict[str, TypeEngine[Any]],
+) -> None:
+    for key, value in record_data.items():
+        if not isinstance(value, (dict, list)):
+            continue
+
+        column_type = column_definitions.get(key)
+        if isinstance(column_type, sqlalchemy.types.JSON):
+            continue
+        if isinstance(column_type, sqlalchemy.types.String):
+            record_data[key] = json.dumps(value)
+            continue
+
+        raise RuntimeError(
+            f"PyAirbyte record field {key!r} has structured "
+            f"{type(value).__name__} data for unsupported SQL column type "
+            f"{type(column_type).__name__ if column_type is not None else 'None'}."
+        )
+
+
+class CtxPostgresProcessor(PostgresSqlProcessor):
+    def process_record_message(
+        self,
+        record_msg: AirbyteRecordMessage,
+        stream_record_handler: StreamRecordHandler,
+        progress_tracker: ProgressTracker,
+    ) -> None:
+        cache = cast(CtxPostgresCache, self.sql_config)
+        record_data = record_msg.data
+        record_data[_CTX_BINDING_ID_COLUMN] = cache.ctx_binding_id
+
+        source_updated_at_field = cache.ctx_source_updated_at_fields.get(
+            record_msg.stream
+        )
+        raw_source_updated_at = _extract_source_updated_at_value(
+            stream_name=record_msg.stream,
+            record_data=record_data,
+            source_updated_at_field=source_updated_at_field,
+        )
+        record_data[_CTX_SOURCE_UPDATED_AT_COLUMN] = _coerce_source_updated_at(
+            raw_source_updated_at
+        )
+
+        column_definitions = self._get_sql_column_definitions(record_msg.stream)
+        _normalize_record_data_for_sql_columns(
+            record_data=record_data,
+            column_definitions=column_definitions,
+        )
+
+        super().process_record_message(
+            record_msg,
+            stream_record_handler,
+            progress_tracker,
+        )
+
+
+class CtxPostgresCache(PostgresCache):
+    _sql_processor_class: ClassVar[type] = CtxPostgresProcessor
+
+    ctx_binding_id: str = ""
+    ctx_source_updated_at_fields: dict[str, str | None] = Field(default_factory=dict)
+
+
+def _build_cache(
+    *,
+    schema_name: str,
+    binding_id: str,
+    source_updated_at_fields: dict[str, str | None],
+) -> CtxPostgresCache:
+    database_url = load_shared_python_settings().ctx_database_url
+    parsed_url = make_url(database_url)
+
+    return CtxPostgresCache(
+        host=parsed_url.host or "localhost",
+        port=parsed_url.port or 5432,
+        database=parsed_url.database,
+        username=parsed_url.username or "postgres",
+        password=parsed_url.password or "",
+        schema_name=schema_name,
+        ctx_binding_id=binding_id,
+        ctx_source_updated_at_fields=source_updated_at_fields,
+    )
+
+
+def build_pyairbyte_source(
+    *,
+    docker_image: str,
+    connector_config: dict[str, Any],
+) -> Any:
+    image_name = docker_image.split(":")[0].split("/")[-1]
+    return ab.get_source(
+        image_name,
+        docker_image=docker_image,
+        config=connector_config,
+    )
+
+
+def require_pyairbyte_selected_stream_names(
+    binding: DagsterAllPlanBinding,
+) -> tuple[str, ...]:
+    if binding.models is not None and binding.models.filter is not None:
+        raise PluginConfigurationError(
+            f"{binding.plugin_id} models.filter is not supported for PyAirbyte connectors."
+        )
+
+    if binding.models is None or binding.models.active is None:
+        raise PluginConfigurationError(
+            f"{binding.plugin_id} models.active must include at least one stream."
+        )
+
+    selected_stream_names = tuple(binding.models.active)
+    if not selected_stream_names:
+        raise PluginConfigurationError(
+            f"{binding.plugin_id} models.active must include at least one stream."
+        )
+
+    return selected_stream_names
+
+
+def run_pyairbyte_sync(
+    *,
+    context: AssetExecutionContext,
+    plugin_id: str,
+    binding_id: str,
+    selected_stream_names: Sequence[str],
+    source_updated_at_fields: Mapping[str, str | None] | None = None,
+    force_full_refresh: bool = False,
+    source: Any,
+) -> dg.MaterializeResult:
+    ordered_selected_stream_names = tuple(selected_stream_names)
+    context.log.info(
+        f"Starting {plugin_id} sync for binding_id={binding_id} "
+        f"selected_streams={list(ordered_selected_stream_names)}"
+    )
+
+    cache = _build_cache(
+        schema_name=plugin_id,
+        binding_id=binding_id,
+        source_updated_at_fields=_build_source_updated_at_fields(
+            ordered_selected_stream_names,
+            source_updated_at_fields,
+        ),
+    )
+
+    _apply_ctx_catalog_contract(
+        source=source,
+        selected_stream_names=ordered_selected_stream_names,
+    )
+    # Airbyte's internal sync model lives below this call. Our runtime does not
+    # choose per-stream incremental vs full-refresh policy; it only chooses:
+    #   1. which raw stream IDs to pass to Source.read(streams=[...])
+    #   2. whether to request force_full_refresh=True
+    #
+    # Verified against the installed PyAirbyte package in this repo:
+    # - `airbyte/sources/base.py:420`
+    #   `Source.get_configured_catalog(...)` builds each configured stream and
+    #   assigns `sync_mode`.
+    # - `airbyte/sources/base.py:449`
+    #   When `force_full_refresh=True`, PyAirbyte prefers `full_refresh` if the
+    #   stream supports it; otherwise it uses `incremental`.
+    # - `airbyte/sources/base.py:873`
+    #   Full refresh also disables state loading by setting
+    #   `state_provider = None`.
+    # - `airbyte/caches/_state_backend.py:31`
+    #   Incremental state is stored in `_airbyte_state`.
+    #
+    # So the repo seam stays intentionally narrow: `binding.models.active`
+    # selects streams, and `force_full_refresh` is the only runtime override.
+    source.read(
+        cache=cache,
+        streams=list(ordered_selected_stream_names),
+        force_full_refresh=force_full_refresh,
+    )
+
+    context.log.info(
+        f"Completed {plugin_id} sync for binding_id={binding_id} "
+        f"selected_streams={list(ordered_selected_stream_names)}"
+    )
+
+    return dg.MaterializeResult(
+        asset_key=dagster_airbyte_sync_asset_key(plugin_id),
+        metadata={
+            "selected_streams": dg.MetadataValue.json(
+                list(ordered_selected_stream_names)
+            ),
+            "selected_stream_count": dg.MetadataValue.int(
+                len(ordered_selected_stream_names)
+            ),
+            "force_full_refresh": dg.MetadataValue.bool(force_full_refresh),
+        },
+    )

shared_plugins/resources.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from copy import deepcopy
+from functools import wraps
+from typing import Any
+
+import dlt
+from dagster_dlt import DagsterDltResource, DagsterDltTranslator
+
+from .models import CtxModel
+
+_CTX_METADATA_COLUMNS = (
+    "_ctx_binding_id",
+    "_ctx_source_updated_at",
+)
+
+
+def _normalize_primary_key(primary_key: object) -> tuple[str, ...]:
+    if isinstance(primary_key, str):
+        return (primary_key,)
+    if isinstance(primary_key, Sequence):
+        normalized = tuple(str(value) for value in primary_key)
+        if normalized:
+            return normalized
+    raise ValueError(
+        "ctx_dlt_* wrappers require a non-empty primary_key that starts with '_ctx_binding_id'."
+    )
+
+
+def _validate_primary_key_prefix(primary_key: object) -> None:
+    normalized = _normalize_primary_key(primary_key)
+    if normalized[0] != "_ctx_binding_id":
+        raise ValueError(
+            "ctx_dlt_* wrappers require primary_key to start with '_ctx_binding_id'. "
+            f"Received: {normalized!r}"
+        )
+
+
+def _apply_max_table_nesting(kwargs: dict[str, Any]) -> None:
+    max_table_nesting = kwargs.pop("max_table_nesting", None)
+    if max_table_nesting not in (None, 0):
+        raise ValueError(
+            "ctx_dlt_* wrappers enforce max_table_nesting=0. Omit the argument or set it to 0."
+        )
+    kwargs["max_table_nesting"] = 0
+
+
+def _with_metadata_columns(columns: object) -> dict[str, Any]:
+    if columns is None:
+        merged_columns: dict[str, Any] = {}
+    elif isinstance(columns, Mapping):
+        merged_columns = deepcopy(dict(columns))
+    else:
+        raise TypeError("ctx_dlt_* wrappers expect 'columns' to be a mapping when set.")
+
+    merged_columns.setdefault("_ctx_source_updated_at", {"data_type": "timestamp"})
+    return merged_columns
+
+
+def _validate_ctx_row(
+    row: object,
+    *,
+    wrapper_name: str,
+    resource_name: str,
+) -> CtxModel:
+    if not isinstance(row, CtxModel):
+        raise TypeError(
+            f"{wrapper_name}('{resource_name}') emitted non-CtxModel row "
+            f"of type {type(row).__name__}."
+        )
+
+    serialized = row.model_dump(by_alias=True)
+    missing_columns = [name for name in _CTX_METADATA_COLUMNS if name not in serialized]
+    if missing_columns:
+        raise TypeError(
+            f"{wrapper_name}('{resource_name}') emitted CtxModel row missing required "
+            f"metadata columns: {', '.join(missing_columns)}."
+        )
+
+    binding_id = serialized.get("_ctx_binding_id")
+    if not isinstance(binding_id, str) or not binding_id.strip():
+        raise TypeError(
+            f"{wrapper_name}('{resource_name}') emitted row with empty '_ctx_binding_id'."
+        )
+
+    return row
+
+
+def _iter_validated_rows(
+    emitted: object,
+    *,
+    wrapper_name: str,
+    resource_name: str,
+) -> Iterable[CtxModel]:
+    if emitted is None:
+        return ()
+
+    if isinstance(emitted, CtxModel):
+        return (
+            _validate_ctx_row(
+                emitted, wrapper_name=wrapper_name, resource_name=resource_name
+            ),
+        )
+
+    if isinstance(emitted, (str, bytes, bytearray, Mapping)):
+        raise TypeError(
+            f"{wrapper_name}('{resource_name}') emitted invalid row container "
+            f"of type {type(emitted).__name__}."
+        )
+
+    if not isinstance(emitted, Iterable):
+        raise TypeError(
+            f"{wrapper_name}('{resource_name}') emitted non-iterable value "
+            f"of type {type(emitted).__name__}."
+        )
+
+    def _generator() -> Iterable[CtxModel]:
+        for row in emitted:
+            yield _validate_ctx_row(
+                row,
+                wrapper_name=wrapper_name,
+                resource_name=resource_name,
+            )
+
+    return _generator()
+
+
+def _wrap_emitter(
+    func: Callable[..., object],
+    *,
+    wrapper_name: str,
+) -> Callable[..., Iterable[CtxModel]]:
+    @wraps(func)
+    def wrapped(*args: Any, **kwargs: Any) -> Iterable[CtxModel]:
+        emitted = func(*args, **kwargs)
+        yield from _iter_validated_rows(
+            emitted,
+            wrapper_name=wrapper_name,
+            resource_name=func.__name__,
+        )
+
+    return wrapped
+
+
+def ctx_dlt_resource(**kwargs: Any) -> Callable[[Callable[..., object]], Any]:
+    decorator_kwargs = dict(kwargs)
+    _apply_max_table_nesting(decorator_kwargs)
+    _validate_primary_key_prefix(decorator_kwargs.get("primary_key"))
+    decorator_kwargs["columns"] = _with_metadata_columns(
+        decorator_kwargs.get("columns")
+    )
+
+    def decorator(func: Callable[..., object]) -> Any:
+        return dlt.resource(**decorator_kwargs)(
+            _wrap_emitter(func, wrapper_name="ctx_dlt_resource")
+        )
+
+    return decorator
+
+
+def ctx_dlt_transformer(**kwargs: Any) -> Callable[[Callable[..., object]], Any]:
+    decorator_kwargs = dict(kwargs)
+    _apply_max_table_nesting(decorator_kwargs)
+    _validate_primary_key_prefix(decorator_kwargs.get("primary_key"))
+    decorator_kwargs["columns"] = _with_metadata_columns(
+        decorator_kwargs.get("columns")
+    )
+
+    def decorator(func: Callable[..., object]) -> Any:
+        return dlt.transformer(**decorator_kwargs)(
+            _wrap_emitter(func, wrapper_name="ctx_dlt_transformer")
+        )
+
+    return decorator
+
+
+DLT_TRANSLATOR = DagsterDltTranslator()
+DLT_RESOURCE = DagsterDltResource()