From f60407ac9dc5e8810af013c14a64a1aac1ca1d2a Mon Sep 17 00:00:00 2001 From: Tobias Macey Date: Thu, 23 Jul 2026 11:45:19 -0400 Subject: [PATCH 1/4] feat(dlt): add a reusable database source, with Keycloak as its first instance Migrating database sources off Airbyte one at a time only pays off if each one is a declaration rather than a new pipeline. A DatabaseSourceSpec of tables becomes a dlt source whose resources already carry the warehouse's raw-table names, so the existing Dagster translator and asset conventions apply unchanged. Two constraints shaped the credential path. Dagster instantiates every dlt source when the code location is imported, so anything resolved there is long stale by the time a scheduled run fires -- Vault credentials are therefore injected per connection through a SQLAlchemy do_connect hook, and table reflection is deferred so no connection is opened at import either. The deployed connection URL carries only a host and database name; a username or password never appears in it. Which database gets read follows the deployment rather than the spec, so local development reads the local-dev Postgres cluster and no developer needs access to QA or production to exercise a pipeline. Keycloak is the first instance: a curated identity subset, deliberately excluding the credential and session tables and dropping client secret columns, since Keycloak's remaining ~80 tables are its own bookkeeping and several hold secrets. Every table is replaced on each run because Keycloak records no row-level modification time -- emails, names and enablement all change in place -- and reconciling missed updates would cost more than a full re-read of a small realm. Co-Authored-By: Claude Opus 4.8 --- .../data_loading/defs/ingestion/assets.py | 7 + .../data_loading/defs/ingestion/schedules.py | 10 + dg_projects/data_loading/uv.lock | 11 +- src/ol_dlt/README.md | 49 +++- src/ol_dlt/ol_dlt/database.py | 265 ++++++++++++++++++ .../ol_dlt/sources/keycloak/__init__.py | 86 ++++++ .../ol_dlt/sources/keycloak/__main__.py | 14 + src/ol_dlt/ol_dlt/vault.py | 123 ++++++++ src/ol_dlt/pyproject.toml | 5 +- src/ol_dlt/tests/sources/test_keycloak.py | 46 +++ src/ol_dlt/tests/test_database.py | 161 +++++++++++ src/ol_dlt/tests/test_vault.py | 121 ++++++++ src/ol_dlt/uv.lock | 53 +++- 13 files changed, 945 insertions(+), 6 deletions(-) create mode 100644 src/ol_dlt/ol_dlt/database.py create mode 100644 src/ol_dlt/ol_dlt/sources/keycloak/__init__.py create mode 100644 src/ol_dlt/ol_dlt/sources/keycloak/__main__.py create mode 100644 src/ol_dlt/ol_dlt/vault.py create mode 100644 src/ol_dlt/tests/sources/test_keycloak.py create mode 100644 src/ol_dlt/tests/test_database.py create mode 100644 src/ol_dlt/tests/test_vault.py diff --git a/dg_projects/data_loading/data_loading/defs/ingestion/assets.py b/dg_projects/data_loading/data_loading/defs/ingestion/assets.py index 050c95b88..fe2663422 100644 --- a/dg_projects/data_loading/data_loading/defs/ingestion/assets.py +++ b/dg_projects/data_loading/data_loading/defs/ingestion/assets.py @@ -14,6 +14,7 @@ from dagster_dlt import DagsterDltResource, dlt_assets from ol_dlt.sources import ( edxorg_s3, + keycloak, mit_climate, mit_edx_programs, mitpe, @@ -73,6 +74,11 @@ def _assets( source=podcast_rss.build_source(), pipeline=podcast_rss.podcast_rss_pipeline, ) +keycloak_assets = build_ingest_assets( + name="keycloak_ingest", + source=keycloak.build_source(), + pipeline=keycloak.keycloak_pipeline, +) # --- edxorg_s3: custom upstream deps + one op per table --------------------- @@ -131,6 +137,7 @@ def _asset( mit_climate_assets, mit_edx_programs_assets, podcast_rss_assets, + keycloak_assets, *edxorg_s3_table_assets, ], ) diff --git a/dg_projects/data_loading/data_loading/defs/ingestion/schedules.py b/dg_projects/data_loading/data_loading/defs/ingestion/schedules.py index 12bc9cd66..d9f6ec570 100644 --- a/dg_projects/data_loading/data_loading/defs/ingestion/schedules.py +++ b/dg_projects/data_loading/data_loading/defs/ingestion/schedules.py @@ -48,6 +48,15 @@ execution_timezone="Etc/UTC", ) +keycloak_ingest_schedule = dg.ScheduleDefinition( + name="keycloak_ingest_daily_schedule", + # Selected by group rather than by key so adding a table to KEYCLOAK_SPEC + # does not also require editing this schedule. + target=dg.AssetSelection.groups("keycloak"), + cron_schedule="30 4 * * *", + execution_timezone="Etc/UTC", +) + defs = dg.Definitions( schedules=[ oll_ingest_schedule, @@ -55,5 +64,6 @@ mit_climate_ingest_schedule, mit_edx_programs_ingest_schedule, podcast_rss_ingest_schedule, + keycloak_ingest_schedule, ], ) diff --git a/dg_projects/data_loading/uv.lock b/dg_projects/data_loading/uv.lock index ec3f3a907..e9d164e5f 100644 --- a/dg_projects/data_loading/uv.lock +++ b/dg_projects/data_loading/uv.lock @@ -796,6 +796,9 @@ pyiceberg = [ { name = "pyiceberg-core" }, { name = "sqlalchemy" }, ] +sql-database = [ + { name = "sqlalchemy" }, +] [[package]] name = "docstring-parser" @@ -1889,8 +1892,10 @@ version = "0.1.0" source = { editable = "../../src/ol_dlt" } dependencies = [ { name = "defusedxml" }, - { name = "dlt", extra = ["filesystem", "parquet", "pyiceberg"] }, + { name = "dlt", extra = ["filesystem", "parquet", "pyiceberg", "sql-database"] }, { name = "duckdb" }, + { name = "hvac" }, + { name = "psycopg2-binary" }, { name = "pyiceberg", extra = ["glue"] }, { name = "requests" }, ] @@ -1898,8 +1903,10 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "defusedxml", specifier = ">=0.7" }, - { name = "dlt", extras = ["filesystem", "parquet", "pyiceberg"], specifier = ">=1.28,<2" }, + { name = "dlt", extras = ["filesystem", "parquet", "pyiceberg", "sql-database"], specifier = ">=1.28,<2" }, { name = "duckdb", specifier = ">=1.0,!=1.5.0,!=1.5.1" }, + { name = "hvac", specifier = ">=2.3" }, + { name = "psycopg2-binary", specifier = ">=2.9" }, { name = "pyiceberg", extras = ["glue"], specifier = ">=0.9" }, { name = "requests", specifier = ">=2.32" }, ] diff --git a/src/ol_dlt/README.md b/src/ol_dlt/README.md index f156b7858..1671c39dc 100644 --- a/src/ol_dlt/README.md +++ b/src/ol_dlt/README.md @@ -21,6 +21,8 @@ src/ol_dlt/ │ └── secrets.toml.template ├── ol_dlt/ │ ├── config.py # profile -> destination/dataset/table_format factory +│ ├── database.py # reusable relational-database source (spec -> @dlt.source) +│ ├── vault.py # Vault-issued database credentials (pure hvac) │ └── sources// # @dlt.source / @dlt.resource bodies only └── tests/ # unit + materialization tests (ephemeral DuckDB/filesystem) ``` @@ -96,6 +98,8 @@ DLT_PROFILE=dev uv run python -m ol_dlt.sources.podcast_rss DLT_PROFILE=dev uv run python -m ol_dlt.sources.mit_edx_programs # Needs AWS creds (reads the prod S3 landing zone): DLT_PROFILE=dev uv run python -m ol_dlt.sources.edxorg_s3 +# Needs the local-dev Postgres cluster port-forwarded (see Database sources): +DLT_PROFILE=dev uv run python -m ol_dlt.sources.keycloak # 5. Inspect a pipeline's state / last load after a run. DLT_PROFILE=dev uv run dlt pipeline oll info @@ -104,10 +108,53 @@ DLT_PROFILE=dev uv run dlt pipeline oll info The hermetic suite (step 1) is what CI runs and what should gate merges; the live runs (step 4) depend on the upstream services being reachable. +## Database sources + +Relational sources do **not** get hand-written resources. Declare a +`DatabaseSourceSpec` in `ol_dlt/sources//__init__.py` and +`ol_dlt.database` builds the `@dlt.source` from it — see +`ol_dlt/sources/keycloak/` for the reference instance. + +```python +KEYCLOAK_SPEC = DatabaseSourceSpec( + name="keycloak", + raw_table_prefix="raw__keycloak__app__postgres__", + database="keycloak", + vault_mount="postgres-keycloak", + tables=(DatabaseTable(name="user_entity", primary_key="id"), ...), +) +``` + +**Which database gets read is the profile's job, not the spec's.** `qa` reads +the QA database, `production` reads production, and `dev`/`ci`/`test` read the +local-dev CloudNativePG cluster from `ol-infrastructure/local-dev`, so no +developer needs access to a deployed environment to exercise a pipeline: + +```bash +kubectl port-forward -n local-infra svc/local-pg-rw 5432:5432 +DLT_PROFILE=dev uv run python -m ol_dlt.sources.keycloak +``` + +Per-source overrides come from `_DB_HOST` / `_PORT` / `_USERNAME` / +`_PASSWORD`. Deployed profiles **require** `_DB_HOST` (the Dagster Pulumi +stack passes it through from the application's stack reference) and take their +username/password from Vault at connection time — never from the environment, +and never at code-location import, where a lease would go stale long before a +scheduled run fires. + +Two things that need doing outside this repo before a new database source can +run in a deployed environment, both in `ol-infrastructure`: + +1. Grant `/creds/readonly` in + `src/ol_infrastructure/applications/dagster/dagster_server_policy.hcl`. +2. Add `_DB_HOST` to the `data_loading` deployment env in the dagster + `__main__.py`, sourced from the application stack's exported RDS host. + ## Adding a new source 1. Author the `@dlt.source` in `ol_dlt/sources//__init__.py` — no env branching, no Dagster imports. Resolve destination/dataset/table_format via - `ol_dlt.config`. + `ol_dlt.config`. For a relational source, declare a `DatabaseSourceSpec` + instead (see **Database sources** above). 2. Add a `@dlt_assets` wrapper in the `data_loading` code location (`defs/ingestion/assets.py`) using the shared `RawDataDltTranslator`. diff --git a/src/ol_dlt/ol_dlt/database.py b/src/ol_dlt/ol_dlt/database.py new file mode 100644 index 000000000..f47a3805b --- /dev/null +++ b/src/ol_dlt/ol_dlt/database.py @@ -0,0 +1,265 @@ +"""Reusable relational-database source for ol_dlt pipelines. + +One declarative spec per source database produces a ``@dlt.source`` whose +resources are named with the warehouse's raw-table convention +(``raw____app__postgres__``), so a new database source is a +data-only change: declare the tables, wrap it in the ``data_loading`` code +location, done. + +Environment alignment + The database a pipeline reads is chosen by the active ``DLT_PROFILE`` + (which ``data_loading`` derives from ``DAGSTER_ENVIRONMENT``): the QA + Dagster deployment reads the QA database, production reads production, and + ``dev``/``ci``/``test`` read the local-dev CloudNativePG cluster that + ``ol-infrastructure/local-dev`` stands up — so a developer never needs + access to a deployed environment to exercise a pipeline. + +Credentials + Deployed profiles fetch short-lived credentials from the Vault database + secrets engine (``/creds/readonly``) at *connection* time, via a + SQLAlchemy ``do_connect`` hook. Resolving them there rather than when the + source is constructed matters: Dagster instantiates every source at code- + location import, and a lease taken then would be stale (or revoked) by the + time a scheduled run actually executes. + +Reflection is deferred (``defer_table_reflect=True``) for the same reason — no +database connection is opened until a resource is extracted. +""" + +import logging +from collections.abc import Generator, Sequence +from dataclasses import dataclass, field +from typing import Any + +import dlt +from dlt.sources.sql_database import sql_table +from sqlalchemy import event +from sqlalchemy.engine import URL, Engine + +from ol_dlt import config, vault + +logger = logging.getLogger(__name__) + +# Rows per batch pulled from the source. The pyarrow backend materializes a +# batch in memory, so this bounds peak memory per table. +DEFAULT_CHUNK_SIZE = 50_000 + +# Matches the shared CloudNativePG cluster in ol-infrastructure/local-dev +# (local-dev/infra/modules/database.py). Reach it with: +# kubectl port-forward -n local-infra svc/local-pg-rw 5432:5432 +LOCAL_DEV_HOST = "localhost" +LOCAL_DEV_USERNAME = "app" +LOCAL_DEV_PASSWORD = "localdev" # noqa: S105 # pragma: allowlist secret + + +@dataclass(frozen=True) +class DatabaseTable: + """One source table to ingest. + + Args: + name: Table name in the source schema. + primary_key: Column(s) uniquely identifying a row. Required when + ``cursor_column`` is set, since incremental loads merge on it. + cursor_column: Monotonic column driving incremental loads. When unset + the table is fully re-read and replaced each run — the right + default for small tables with no reliable modification timestamp. + excluded_columns: Columns never to select. Use this for secrets and + credential material, which must not reach the warehouse at all. + """ + + name: str + primary_key: str | Sequence[str] | None = None + cursor_column: str | None = None + excluded_columns: Sequence[str] = () + + def __post_init__(self) -> None: + """Reject an incremental table that has no key to merge on.""" + if self.cursor_column and not self.primary_key: + msg = ( + f"{self.name}: primary_key is required when cursor_column is set " + "(incremental loads merge on the primary key)" + ) + raise ValueError(msg) + + +@dataclass(frozen=True) +class DatabaseSourceSpec: + """Declarative description of one source database. + + Args: + name: Short source name. Drives the dlt pipeline name, the per-source + destination prefix, and the ``_DB_*`` environment variables. + raw_table_prefix: Prefix applied to each table name to form the raw + warehouse table (and Dagster asset) name. + database: Database name to connect to. + vault_mount: Vault database secrets-engine mount serving this database, + e.g. ``postgres-keycloak``. Used by deployed profiles only. + tables: The tables to ingest. + db_schema: Schema the tables live in, or ``None`` for engines that have + no named schema. + vault_role: Vault role to request. Ingestion is read-only, always. + port: Database port. + drivername: SQLAlchemy driver. + """ + + name: str + raw_table_prefix: str + database: str + vault_mount: str + tables: Sequence[DatabaseTable] = field(default_factory=tuple) + db_schema: str | None = "public" + vault_role: str = "readonly" + port: int = 5432 + drivername: str = "postgresql+psycopg2" + + @property + def env_prefix(self) -> str: + """Environment-variable prefix for per-deployment overrides.""" + return f"{self.name.upper()}_DB" + + def raw_table_name(self, table: DatabaseTable) -> str: + """Return the raw warehouse table name for ``table``.""" + return f"{self.raw_table_prefix}{table.name}" + + +def _uses_vault(profile: str) -> bool: + """Return True when the profile reads a deployed (Vault-backed) database.""" + return profile in config.ICEBERG_PROFILES + + +def _connection_url(spec: DatabaseSourceSpec, profile: str) -> URL: + """Build the connection URL for ``spec`` under ``profile``. + + Deployed profiles omit the username/password — those are injected per + connection from Vault — and require TLS, which the RDS parameter group + enforces anyway (``rds.force_ssl``). + """ + host = config.resolve_secret(None, f"{spec.env_prefix}_HOST") + port = config.resolve_secret(None, f"{spec.env_prefix}_PORT") + if _uses_vault(profile): + if not host: + msg = ( + f"{spec.env_prefix}_HOST is not set. The {profile} deployment must " + f"pass the {spec.name} database host through from its Pulumi stack " + "reference." + ) + raise ValueError(msg) + return URL.create( + spec.drivername, + host=host, + port=int(port) if port else spec.port, + database=spec.database, + query={"sslmode": "require"}, + ) + return URL.create( + spec.drivername, + username=config.resolve_secret(None, f"{spec.env_prefix}_USERNAME") + or LOCAL_DEV_USERNAME, + password=config.resolve_secret(None, f"{spec.env_prefix}_PASSWORD") + or LOCAL_DEV_PASSWORD, + host=host or LOCAL_DEV_HOST, + port=int(port) if port else spec.port, + database=spec.database, + ) + + +def _vault_credential_injector(spec: DatabaseSourceSpec) -> Any: # noqa: ANN401 + """Return an engine adapter that supplies Vault credentials per connection. + + SQLAlchemy's ``do_connect`` event fires when a pooled connection is + actually opened, so each connection gets credentials that are live *then* + rather than whenever the pipeline was defined. + """ + + def adapt(engine: Engine) -> Engine: + @event.listens_for(engine, "do_connect") + def _inject( + _dialect: Any, # noqa: ANN401 + _conn_rec: Any, # noqa: ANN401 + _cargs: Any, # noqa: ANN401 + cparams: dict[str, Any], + ) -> None: + username, password = vault.read_database_credentials( + spec.vault_mount, spec.vault_role + ) + cparams["user"] = username + cparams["password"] = password + + return engine + + return adapt + + +def build_table_resource( + spec: DatabaseSourceSpec, + table: DatabaseTable, + *, + profile: str | None = None, +) -> Any: # noqa: ANN401 + """Return the dlt resource for one table of ``spec``.""" + resolved_profile = profile or config.active_profile() + raw_name = spec.raw_table_name(table) + incremental = ( + dlt.sources.incremental(table.cursor_column) if table.cursor_column else None + ) + resource = sql_table( + credentials=_connection_url(spec, resolved_profile).render_as_string( + hide_password=False + ), + table=table.name, + schema=spec.db_schema, + incremental=incremental, + chunk_size=DEFAULT_CHUNK_SIZE, + # pyarrow reflects source types faithfully instead of re-inferring them + # from sampled rows, which keeps the Iceberg schema stable run to run. + backend="pyarrow", + reflection_level="full", + # No connection until extraction — see the module docstring. + defer_table_reflect=True, + excluded_columns=list(table.excluded_columns) or None, + write_disposition="merge" if table.cursor_column else "replace", + primary_key=table.primary_key, + engine_adapter_callback=( + _vault_credential_injector(spec) if _uses_vault(resolved_profile) else None + ), + ) + return resource.with_name(raw_name).apply_hints( + table_name=raw_name, + table_format=config.active_table_format(), + ) + + +def build_database_source( + spec: DatabaseSourceSpec, + *, + tables: Sequence[str] | None = None, + profile: str | None = None, +) -> Any: # noqa: ANN401 + """Return the ``@dlt.source`` for ``spec``. + + Args: + spec: The database source description. + tables: Optional subset of ``spec.tables`` names to include, for + targeted local runs. + profile: Override the active profile (mainly for tests). + """ + selected = ( + [table for table in spec.tables if table.name in set(tables)] + if tables is not None + else list(spec.tables) + ) + if not selected: + msg = f"{spec.name}: no tables selected from {tables!r}" + raise ValueError(msg) + + @dlt.source(name=f"{spec.name}_ingest") + def _source() -> Generator[Any]: + for table in selected: + yield build_table_resource(spec, table, profile=profile) + + return _source() + + +def pipeline_for(spec: DatabaseSourceSpec) -> dlt.Pipeline: + """Return the dlt pipeline for ``spec``'s destination.""" + return config.pipeline_for(spec.name) diff --git a/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py b/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py new file mode 100644 index 000000000..f5017aaa6 --- /dev/null +++ b/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py @@ -0,0 +1,86 @@ +"""Keycloak application-database ingestion via dlt. + +First instance of the reusable ``ol_dlt.database`` source. Keycloak is the SSO +system of record for MIT Open Learning identities, so these tables are what let +the warehouse resolve a person across the applications that federate to it. + +Data flow: + Keycloak RDS Postgres -> raw__keycloak__app__postgres__
+ +The database read is aligned with the deployment: QA Dagster reads the QA +Keycloak database, production reads production, and local development reads the +``keycloak`` database in the local-dev CloudNativePG cluster (see +``ol_dlt.database``). + +Scope is a curated identity subset rather than the whole Keycloak schema. The +~80 remaining tables are Keycloak's own internal machinery (session, event, +migration and policy bookkeeping) with no downstream modelling value, and +several of them hold credential material. Notably absent, deliberately: + + credential password hashes and OTP secrets + component_config identity-provider and LDAP bind secrets + *_session, *_event high-churn operational state, not identity facts + +Every table is loaded with ``write_disposition="replace"``: Keycloak has no +reliable row-level modification timestamp to drive an incremental cursor +(``user_entity.created_timestamp`` only records creation, while emails, names +and enablement change in place), and the realm is small enough that a full +re-read per run is cheaper than reconciling missed updates. + +Run standalone against local-dev (port-forward the CNPG cluster first): + kubectl port-forward -n local-infra svc/local-pg-rw 5432:5432 + DLT_PROFILE=dev python -m ol_dlt.sources.keycloak +""" + +from typing import Any + +from ol_dlt.database import ( + DatabaseSourceSpec, + DatabaseTable, + build_database_source, + pipeline_for, +) + +KEYCLOAK_SPEC = DatabaseSourceSpec( + name="keycloak", + raw_table_prefix="raw__keycloak__app__postgres__", + database="keycloak", + vault_mount="postgres-keycloak", + tables=( + # --- identities ----------------------------------------------------- + DatabaseTable(name="user_entity", primary_key="id"), + DatabaseTable(name="user_attribute", primary_key="id"), + DatabaseTable( + name="federated_identity", primary_key=("identity_provider", "user_id") + ), + DatabaseTable( + name="user_required_action", primary_key=("required_action", "user_id") + ), + # --- group membership ----------------------------------------------- + DatabaseTable(name="keycloak_group", primary_key="id"), + DatabaseTable(name="group_attribute", primary_key="id"), + DatabaseTable( + name="user_group_membership", primary_key=("group_id", "user_id") + ), + # --- role assignment ------------------------------------------------- + DatabaseTable(name="keycloak_role", primary_key="id"), + DatabaseTable(name="user_role_mapping", primary_key=("role_id", "user_id")), + # --- realm / client context ----------------------------------------- + DatabaseTable(name="realm", primary_key="id"), + DatabaseTable( + name="client", + primary_key="id", + # Client secrets and registration tokens are credentials; they have + # no analytical use and must never land in the warehouse. + excluded_columns=("secret", "registration_token"), + ), + DatabaseTable(name="identity_provider", primary_key="internal_id"), + ), +) + +keycloak_pipeline = pipeline_for(KEYCLOAK_SPEC) + + +def build_source(tables: list[str] | None = None) -> Any: # noqa: ANN401 + """Instantiate the Keycloak source (uniform entrypoint for Dagster).""" + return build_database_source(KEYCLOAK_SPEC, tables=tables) diff --git a/src/ol_dlt/ol_dlt/sources/keycloak/__main__.py b/src/ol_dlt/ol_dlt/sources/keycloak/__main__.py new file mode 100644 index 000000000..f8fa1636f --- /dev/null +++ b/src/ol_dlt/ol_dlt/sources/keycloak/__main__.py @@ -0,0 +1,14 @@ +"""Standalone smoke run: ``DLT_PROFILE=dev python -m ol_dlt.sources.keycloak``. + +Requires the local-dev Postgres cluster to be reachable, e.g. +``kubectl port-forward -n local-infra svc/local-pg-rw 5432:5432``. +""" + +import logging + +from ol_dlt.sources.keycloak import build_source, keycloak_pipeline + +logging.basicConfig(level=logging.INFO) +logging.getLogger(__name__).info( + "Pipeline completed: %s", keycloak_pipeline.run(build_source()) +) diff --git a/src/ol_dlt/ol_dlt/vault.py b/src/ol_dlt/ol_dlt/vault.py new file mode 100644 index 000000000..d210736c9 --- /dev/null +++ b/src/ol_dlt/ol_dlt/vault.py @@ -0,0 +1,123 @@ +"""Vault credential resolution for ol_dlt sources. + +Pure ``hvac`` — this module must not import Dagster (see the ruff banned-api +rule in ``pyproject.toml``). It mirrors the authentication contract that +``ol_orchestrate.lib.utils.authenticate_vault`` uses in the Dagster code +locations, so a dlt pipeline resolves credentials the same way whether it runs +inside a Dagster pod or standalone: + + - In Kubernetes (the service-account token file exists): Vault Kubernetes + auth with the role/mount from ``DAGSTER_VAULT_ROLE`` / + ``DAGSTER_VAULT_MOUNT``. + - Anywhere else: an existing token from ``VAULT_TOKEN`` (e.g. after + ``vault login`` or ``bin/starrocks-auth``). + +Credentials are fetched at pipeline *run* time (never at import), and cached +for a fraction of the Vault lease so every table in a single run shares one +lease instead of opening a new one per table. +""" + +import logging +import os +import time +from pathlib import Path +from typing import Any + +import hvac + +logger = logging.getLogger(__name__) + +K8S_SERVICE_ACCOUNT_TOKEN = Path( + "/var/run/secrets/kubernetes.io/serviceaccount/token" # noqa: S108 +) + +# Re-fetch once this fraction of the lease has elapsed. Long-lived code-location +# pods outlive any single lease, so a process-lifetime cache would eventually +# hand out revoked credentials. +_LEASE_USE_FRACTION = 0.5 + +# (mount, role) -> (expires_at_monotonic, username, password) +_CREDENTIAL_CACHE: dict[tuple[str, str], tuple[float, str, str]] = {} + + +def vault_address() -> str: + """Return the Vault address for the current environment. + + Mirrors ``ol_orchestrate.lib.constants.VAULT_ADDRESS``: ``VAULT_ADDR`` wins, + otherwise it is derived from ``DAGSTER_ENVIRONMENT`` (``dev`` points at QA + so local development never touches production Vault). + """ + explicit = os.getenv("VAULT_ADDR") + if explicit: + return explicit + environment = os.getenv("DAGSTER_ENVIRONMENT", "dev") + if environment == "dev": + environment = "qa" + return f"https://vault-{environment}.odl.mit.edu" + + +def _authenticated_client() -> hvac.Client: + """Return an authenticated Vault client (Kubernetes auth, else token).""" + client = hvac.Client(url=vault_address()) + if K8S_SERVICE_ACCOUNT_TOKEN.exists(): + client.auth.kubernetes.login( + role=os.getenv("DAGSTER_VAULT_ROLE", "dagster"), + jwt=K8S_SERVICE_ACCOUNT_TOKEN.read_text(), + use_token=True, + mount_point=os.getenv("DAGSTER_VAULT_MOUNT", "k8s-data"), + ) + return client + token = os.getenv("VAULT_TOKEN") + if not token: + msg = ( + "No Vault Kubernetes service-account token and no VAULT_TOKEN in the " + "environment. Authenticate first (`vault login`) or run against the " + "dev profile, which uses the local-dev database instead of Vault." + ) + raise RuntimeError(msg) + client.token = token + return client + + +def read_database_credentials(mount: str, role: str = "readonly") -> tuple[str, str]: + """Return ``(username, password)`` from Vault's database secrets engine. + + Args: + mount: Database secrets-engine mount point, e.g. ``postgres-keycloak``. + role: Role to issue credentials for. ``readonly`` is the only role an + ingestion pipeline should ever ask for. + """ + cache_key = (mount, role) + cached = _CREDENTIAL_CACHE.get(cache_key) + if cached and time.monotonic() < cached[0]: + return cached[1], cached[2] + + vault_path = f"{mount}/creds/{role}" + client = _authenticated_client() + try: + response: dict[str, Any] | None = client.read(vault_path) + except hvac.exceptions.Forbidden as exc: + msg = ( + f"Vault denied {vault_path!r} — the token's policy does not grant it. " + "Source database mounts must be added to the dagster policy in " + "ol-infrastructure (dagster_server_policy.hcl)." + ) + raise RuntimeError(msg) from exc + if response is None: + msg = f"Vault path not found: {vault_path!r} — check the mount and role name" + raise RuntimeError(msg) + + lease_duration = int(response.get("lease_duration") or 0) + username = response["data"]["username"] + password = response["data"]["password"] + logger.info( + "Issued Vault database credentials from %s (lease %ss)", + vault_path, + lease_duration, + ) + _CREDENTIAL_CACHE[cache_key] = ( + time.monotonic() + lease_duration * _LEASE_USE_FRACTION, + username, + password, + ) + return username, password diff --git a/src/ol_dlt/pyproject.toml b/src/ol_dlt/pyproject.toml index 619d4967c..896578172 100644 --- a/src/ol_dlt/pyproject.toml +++ b/src/ol_dlt/pyproject.toml @@ -6,11 +6,14 @@ authors = [{ name = "MIT Open Learning Engineering", email = "ol-data@mit.edu" } requires-python = "~=3.13,<3.15" license = "BSD-3-Clause" dependencies = [ - "dlt[filesystem,parquet,pyiceberg]>=1.28,<2", + "dlt[filesystem,parquet,pyiceberg,sql_database]>=1.28,<2", "pyiceberg[glue]>=0.9", "requests>=2.32", "duckdb>=1.0,!=1.5.0,!=1.5.1", "defusedxml>=0.7", + # Relational sources (ol_dlt.database) and their Vault-issued credentials. + "psycopg2-binary>=2.9", + "hvac>=2.3", ] [dependency-groups] diff --git a/src/ol_dlt/tests/sources/test_keycloak.py b/src/ol_dlt/tests/sources/test_keycloak.py new file mode 100644 index 000000000..55ae0b0e0 --- /dev/null +++ b/src/ol_dlt/tests/sources/test_keycloak.py @@ -0,0 +1,46 @@ +"""Tests for the Keycloak database source.""" + +from ol_dlt.sources import keycloak + +# Tables that hold credential material or high-churn operational state. These +# must never appear in the spec; the assertion exists so adding one is a +# deliberate, reviewed act rather than an accident. +FORBIDDEN_TABLES = frozenset( + { + "credential", + "component_config", + "user_session", + "offline_user_session", + "offline_client_session", + "admin_event_entity", + "event_entity", + } +) + + +def test_spec_excludes_credential_tables() -> None: + selected = {table.name for table in keycloak.KEYCLOAK_SPEC.tables} + assert not selected & FORBIDDEN_TABLES + + +def test_client_secret_columns_are_excluded() -> None: + client = next( + table for table in keycloak.KEYCLOAK_SPEC.tables if table.name == "client" + ) + assert set(client.excluded_columns) == {"secret", "registration_token"} + + +def test_every_table_declares_a_primary_key() -> None: + assert all(table.primary_key for table in keycloak.KEYCLOAK_SPEC.tables) + + +def test_resources_follow_the_raw_naming_convention() -> None: + source = keycloak.build_source() + assert "raw__keycloak__app__postgres__user_entity" in source.resources + assert all( + name.startswith("raw__keycloak__app__postgres__") for name in source.resources + ) + + +def test_pipeline_targets_the_keycloak_prefix() -> None: + assert keycloak.keycloak_pipeline.pipeline_name == "keycloak" diff --git a/src/ol_dlt/tests/test_database.py b/src/ol_dlt/tests/test_database.py new file mode 100644 index 000000000..e00c16ede --- /dev/null +++ b/src/ol_dlt/tests/test_database.py @@ -0,0 +1,161 @@ +"""Unit + materialization tests for the reusable database source.""" + +import sqlite3 +from pathlib import Path +from typing import NoReturn + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.engine import URL + +from ol_dlt import config, database +from ol_dlt.database import DatabaseSourceSpec, DatabaseTable + +_SPEC = DatabaseSourceSpec( + name="example", + raw_table_prefix="raw__example__app__postgres__", + database="example", + vault_mount="postgres-example", + tables=( + DatabaseTable(name="widget", primary_key="id"), + DatabaseTable(name="gadget", primary_key="id", excluded_columns=("secret",)), + ), +) + + +def test_cursor_column_requires_primary_key() -> None: + with pytest.raises(ValueError, match="primary_key is required"): + DatabaseTable(name="thing", cursor_column="updated_at") + + +def test_local_profile_url_uses_local_dev_defaults() -> None: + url = database._connection_url(_SPEC, "dev") # noqa: SLF001 + assert url.host == database.LOCAL_DEV_HOST + assert url.username == database.LOCAL_DEV_USERNAME + assert url.password == database.LOCAL_DEV_PASSWORD + assert url.database == "example" + assert "sslmode" not in url.query + + +def test_local_profile_url_honors_env_overrides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("EXAMPLE_DB_HOST", "db.internal") + monkeypatch.setenv("EXAMPLE_DB_PORT", "6543") + monkeypatch.setenv("EXAMPLE_DB_USERNAME", "reader") + url = database._connection_url(_SPEC, "dev") # noqa: SLF001 + assert (url.host, url.port, url.username) == ("db.internal", 6543, "reader") + + +def test_deployed_profile_url_omits_credentials_and_requires_tls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("EXAMPLE_DB_HOST", "example.rds.amazonaws.com") + url = database._connection_url(_SPEC, "production") # noqa: SLF001 + # Credentials are injected per connection from Vault, never in the URL. + assert url.username is None + assert url.password is None + assert url.query["sslmode"] == "require" + + +def test_deployed_profile_requires_host(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("EXAMPLE_DB_HOST", raising=False) + with pytest.raises(ValueError, match="EXAMPLE_DB_HOST is not set"): + database._connection_url(_SPEC, "production") # noqa: SLF001 + + +def test_deployed_profile_never_reads_vault_at_definition_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Building the source must not open a Vault lease. + + Dagster instantiates every source when the code location is imported; a + lease taken there would be long expired by the time a scheduled run fires. + """ + monkeypatch.setenv("EXAMPLE_DB_HOST", "example.rds.amazonaws.com") + + def _fail(*_args: object, **_kwargs: object) -> NoReturn: + msg = "Vault was read while the source was being defined" + raise AssertionError(msg) + + monkeypatch.setattr(database.vault, "read_database_credentials", _fail) + source = database.build_database_source(_SPEC, profile="production") + assert set(source.resources) == { + "raw__example__app__postgres__widget", + "raw__example__app__postgres__gadget", + } + + +def test_vault_credentials_are_injected_at_connection_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + database.vault, + "read_database_credentials", + lambda mount, role: (f"v-{mount}", f"p-{role}"), + ) + engine = create_engine("postgresql+psycopg2://example.rds.amazonaws.com/example") + database._vault_credential_injector(_SPEC)(engine) # noqa: SLF001 + + connect_params: dict[str, str] = {} + engine.dialect.dispatch.do_connect(engine.dialect, None, [], connect_params) + assert connect_params == { + "user": "v-postgres-example", + "password": "p-readonly", # pragma: allowlist secret + } + + +def test_build_database_source_selects_a_subset() -> None: + source = database.build_database_source(_SPEC, tables=["widget"], profile="dev") + assert set(source.resources) == {"raw__example__app__postgres__widget"} + + +def test_build_database_source_rejects_an_empty_selection() -> None: + with pytest.raises(ValueError, match="no tables selected"): + database.build_database_source(_SPEC, tables=["nope"], profile="dev") + + +@pytest.mark.integration +def test_materializes_a_database_to_the_test_profile( + test_profile: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """End-to-end: reflect, extract and write a real relational source. + + SQLite stands in for Postgres so the test stays hermetic; everything under + test (deferred reflection, resource renaming, column exclusion, the raw + table layout) is driver-independent. + """ + db_path = tmp_path / "example.db" + with sqlite3.connect(db_path) as connection: + connection.execute("CREATE TABLE widget (id INTEGER PRIMARY KEY, name TEXT)") + connection.execute( + "CREATE TABLE gadget (id INTEGER PRIMARY KEY, name TEXT, secret TEXT)" + ) + connection.executemany( + "INSERT INTO widget VALUES (?, ?)", [(1, "first"), (2, "second")] + ) + connection.execute("INSERT INTO gadget VALUES (1, 'gizmo', 'hunter2')") + + monkeypatch.setattr( + database, + "_connection_url", + lambda *_a, **_k: URL.create("sqlite", database=str(db_path)), + ) + spec = DatabaseSourceSpec( + name="example", + raw_table_prefix="raw__example__app__postgres__", + database="example", + vault_mount="postgres-example", + db_schema=None, # SQLite has no named schema + tables=_SPEC.tables, + ) + pipeline = config.pipeline_for("example") + info = pipeline.run(database.build_database_source(spec, profile="test")) + assert not info.has_failed_jobs + + tables = pipeline.default_schema.tables + assert "raw__example__app__postgres__widget" in tables + gadget_columns = tables["raw__example__app__postgres__gadget"]["columns"] + assert "name" in gadget_columns + # excluded_columns must keep credential material out of the warehouse. + assert "secret" not in gadget_columns diff --git a/src/ol_dlt/tests/test_vault.py b/src/ol_dlt/tests/test_vault.py new file mode 100644 index 000000000..ccda85fad --- /dev/null +++ b/src/ol_dlt/tests/test_vault.py @@ -0,0 +1,121 @@ +"""Tests for Vault credential resolution.""" + +from typing import Any + +import hvac +import pytest + +from ol_dlt import vault + + +@pytest.fixture(autouse=True) +def _clear_credential_cache() -> None: + vault._CREDENTIAL_CACHE.clear() # noqa: SLF001 + + +class FakeVaultClient: + """Records reads and returns a canned database-credentials response.""" + + def __init__(self, response: Any = None, error: Exception | None = None) -> None: # noqa: ANN401 + self.response = response + self.error = error + self.reads: list[str] = [] + + def read(self, path: str) -> Any: # noqa: ANN401 + self.reads.append(path) + if self.error: + raise self.error + return self.response + + +def _credentials_response(lease_duration: int = 3600) -> dict[str, Any]: + return { + "lease_duration": lease_duration, + "data": { + "username": "v-user-1", + "password": "v-pass-1", # pragma: allowlist secret + }, + } + + +def test_vault_address_prefers_the_explicit_env_var( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("VAULT_ADDR", "https://vault.example.test") + assert vault.vault_address() == "https://vault.example.test" + + +def test_vault_address_for_dev_points_at_qa(monkeypatch: pytest.MonkeyPatch) -> None: + """Local development must never reach production Vault by default.""" + monkeypatch.delenv("VAULT_ADDR", raising=False) + monkeypatch.setenv("DAGSTER_ENVIRONMENT", "dev") + assert vault.vault_address() == "https://vault-qa.odl.mit.edu" + + +def test_vault_address_follows_the_dagster_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("VAULT_ADDR", raising=False) + monkeypatch.setenv("DAGSTER_ENVIRONMENT", "production") + assert vault.vault_address() == "https://vault-production.odl.mit.edu" + + +def test_reads_the_readonly_credentials_path(monkeypatch: pytest.MonkeyPatch) -> None: + client = FakeVaultClient(_credentials_response()) + monkeypatch.setattr(vault, "_authenticated_client", lambda: client) + assert vault.read_database_credentials("postgres-keycloak") == ( + "v-user-1", + "v-pass-1", + ) + assert client.reads == ["postgres-keycloak/creds/readonly"] + + +def test_credentials_are_cached_within_the_lease( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = FakeVaultClient(_credentials_response(lease_duration=3600)) + monkeypatch.setattr(vault, "_authenticated_client", lambda: client) + vault.read_database_credentials("postgres-keycloak") + vault.read_database_credentials("postgres-keycloak") + # One lease serves every table in a run rather than one lease per table. + assert len(client.reads) == 1 + + +def test_credentials_are_refetched_once_the_lease_is_half_spent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A code-location pod outlives any single lease, so the cache must expire.""" + client = FakeVaultClient(_credentials_response(lease_duration=100)) + monkeypatch.setattr(vault, "_authenticated_client", lambda: client) + clock = [1000.0] + monkeypatch.setattr(vault.time, "monotonic", lambda: clock[0]) + + vault.read_database_credentials("postgres-keycloak") + clock[0] += 60 # past 50% of the 100s lease + vault.read_database_credentials("postgres-keycloak") + assert len(client.reads) == 2 # noqa: PLR2004 + + +def test_a_denied_path_names_the_policy_to_fix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = FakeVaultClient(error=hvac.exceptions.Forbidden("denied")) + monkeypatch.setattr(vault, "_authenticated_client", lambda: client) + with pytest.raises(RuntimeError, match="dagster_server_policy.hcl"): + vault.read_database_credentials("postgres-keycloak") + + +def test_a_missing_path_is_reported_clearly(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(vault, "_authenticated_client", lambda: FakeVaultClient(None)) + with pytest.raises(RuntimeError, match="Vault path not found"): + vault.read_database_credentials("postgres-nonexistent") + + +def test_no_kubernetes_token_and_no_vault_token_fails_loudly( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Any, # noqa: ANN401 +) -> None: + monkeypatch.setattr(vault, "K8S_SERVICE_ACCOUNT_TOKEN", tmp_path / "absent") + monkeypatch.delenv("VAULT_TOKEN", raising=False) + with pytest.raises(RuntimeError, match="No Vault Kubernetes service-account token"): + vault._authenticated_client() # noqa: SLF001 diff --git a/src/ol_dlt/uv.lock b/src/ol_dlt/uv.lock index ba1203166..ec49bba76 100644 --- a/src/ol_dlt/uv.lock +++ b/src/ol_dlt/uv.lock @@ -383,6 +383,9 @@ pyiceberg = [ { name = "pyiceberg-core" }, { name = "sqlalchemy" }, ] +sql-database = [ + { name = "sqlalchemy" }, +] [[package]] name = "duckdb" @@ -561,6 +564,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b0/aa/0b7365d30fed43e7a3449aba1fe20a0a7174d9cf13e282af4e69ac825441/humanize-4.16.0-py3-none-any.whl", hash = "sha256:353eb2f34c09d098b2880eee8bef21832eae6d174f48c5762fff7e5fcb74d01d", size = 137209, upload-time = "2026-06-30T16:17:28.36Z" }, ] +[[package]] +name = "hvac" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e2/57/b46c397fb3842cfb02a44609aa834c887f38dd75f290c2fc5a34da4b2fee/hvac-2.4.0.tar.gz", hash = "sha256:e0056ad9064e7923e874e6769015b032580b639e29246f5ab1044f7959c1c7e0", size = 332543, upload-time = "2025-10-30T12:57:47.512Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/33/71e45a6bd6875f44a26f99da31c63b6840123e88bedf2c0b1ce429b8be12/hvac-2.4.0-py3-none-any.whl", hash = "sha256:008db5efd8c2f77bd37d2368ea5f713edceae1c65f11fd608393179478649e0f", size = 155921, upload-time = "2025-10-30T12:57:46.253Z" }, +] + [[package]] name = "idna" version = "3.18" @@ -771,8 +786,10 @@ version = "0.1.0" source = { editable = "." } dependencies = [ { name = "defusedxml" }, - { name = "dlt", extra = ["filesystem", "parquet", "pyiceberg"] }, + { name = "dlt", extra = ["filesystem", "parquet", "pyiceberg", "sql-database"] }, { name = "duckdb" }, + { name = "hvac" }, + { name = "psycopg2-binary" }, { name = "pyiceberg", extra = ["glue"] }, { name = "requests" }, ] @@ -787,8 +804,10 @@ dev = [ [package.metadata] requires-dist = [ { name = "defusedxml", specifier = ">=0.7" }, - { name = "dlt", extras = ["filesystem", "parquet", "pyiceberg"], specifier = ">=1.28,<2" }, + { name = "dlt", extras = ["filesystem", "parquet", "pyiceberg", "sql-database"], specifier = ">=1.28,<2" }, { name = "duckdb", specifier = ">=1.0,!=1.5.0,!=1.5.1" }, + { name = "hvac", specifier = ">=2.3" }, + { name = "psycopg2-binary", specifier = ">=2.9" }, { name = "pyiceberg", extras = ["glue"], specifier = ">=0.9" }, { name = "requests", specifier = ">=2.32" }, ] @@ -975,6 +994,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/ed/1cdcab6ba3d6ab7feca11fc14f0eeea80755bb53ef4e892079f31b10a25f/propcache-0.5.2-py3-none-any.whl", hash = "sha256:be1ddfcbb376e3de5d2e2db1d58d6d67463e6b4f9f040c000de8e300295465fe", size = 14036, upload-time = "2026-05-08T21:02:10.673Z" }, ] +[[package]] +name = "psycopg2-binary" +version = "2.9.12" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2a/60/a3624f79acea344c16fbef3a94d28b89a8042ddfb8f3e4ca83f538671409/psycopg2_binary-2.9.12.tar.gz", hash = "sha256:5ac9444edc768c02a6b6a591f070b8aae28ff3a99be57560ac996001580f294c", size = 379686, upload-time = "2026-04-21T09:40:34.304Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/bb/4608c96f970f6e0c56572e87027ef4404f709382a3503e9934526d7ba051/psycopg2_binary-2.9.12-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7c729a73c7b1b84de3582f73cdd27d905121dc2c531f3d9a3c32a3011033b965", size = 3712419, upload-time = "2026-04-20T23:34:58.754Z" }, + { url = "https://files.pythonhosted.org/packages/5e/af/48f76af9d50d61cf390f8cd657b503168b089e2e9298e48465d029fcc713/psycopg2_binary-2.9.12-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4413d0caef93c5cf50b96863df4c2efe8c269bf2267df353225595e7e15e8df7", size = 3822990, upload-time = "2026-04-20T23:35:00.821Z" }, + { url = "https://files.pythonhosted.org/packages/7a/df/aba0f99397cd811d32e06fc0cc781f1f3ce98bc0e729cb423925085d781a/psycopg2_binary-2.9.12-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:4dfcf8e45ebb0c663be34a3442f65e17311f3367089cd4e5e3a3e8e62c978777", size = 4578696, upload-time = "2026-04-20T23:35:03.409Z" }, + { url = "https://files.pythonhosted.org/packages/95/9c/eaa74021ac4e4d5c2f83d82fc6615a63f4fe6c94dc4e94c3990427053f67/psycopg2_binary-2.9.12-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c41321a14dd74aceb6a9a643b9253a334521babfa763fa873e33d89cfa122fb5", size = 4274982, upload-time = "2026-04-20T23:35:05.583Z" }, + { url = "https://files.pythonhosted.org/packages/35/ed/c25deff98bd26187ba48b3b250a3ffc3037c46c5b89362534a15d200e0db/psycopg2_binary-2.9.12-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:83946ba43979ebfdc99a3cd0ee775c89f221df026984ba19d46133d8d75d3cd9", size = 5894867, upload-time = "2026-04-20T23:35:07.902Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/8d0e21ca77373c6c9589e5c4528f6e8f0c08c62cafc76fb0bddb7a2cee22/psycopg2_binary-2.9.12-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:411e85815652d13560fbe731878daa5d92378c4995a22302071890ec3397d019", size = 4110578, upload-time = "2026-04-20T23:35:10.149Z" }, + { url = "https://files.pythonhosted.org/packages/00/fc/f481e2435bd8f742d0123309174aae4165160ad3ef17c1b99c3622c241d2/psycopg2_binary-2.9.12-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1c8ad4c08e00f7679559eaed7aff1edfffc60c086b976f93972f686384a95e2c", size = 3655816, upload-time = "2026-04-20T23:35:12.56Z" }, + { url = "https://files.pythonhosted.org/packages/53/79/b9f46466bdbe9f239c96cde8be33c1aace4842f06013b47b730dc9759187/psycopg2_binary-2.9.12-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:00814e40fa23c2b37ef0a1e3c749d89982c73a9cb5046137f0752a22d432e82f", size = 3301307, upload-time = "2026-04-20T23:35:15.029Z" }, + { url = "https://files.pythonhosted.org/packages/3f/19/7dc003b32fe35024df89b658104f7c8538a8b2dcbde7a4e746ce929742e7/psycopg2_binary-2.9.12-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:98062447aebc20ed20add1f547a364fd0ef8933640d5372ff1873f8deb9b61be", size = 3048968, upload-time = "2026-04-20T23:35:16.757Z" }, + { url = "https://files.pythonhosted.org/packages/91/58/2dbd7db5c604d45f4950d988506aae672a14126ec22998ced5021cbb76bb/psycopg2_binary-2.9.12-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:66a7685d7e548f10fb4ce32fb01a7b7f4aa702134de92a292c7bd9e0d3dbd290", size = 3351369, upload-time = "2026-04-20T23:35:18.933Z" }, + { url = "https://files.pythonhosted.org/packages/42/ee/dee8dcaad07f735824de3d6563bc67119fa6c28257b17977a8d624f02fab/psycopg2_binary-2.9.12-cp313-cp313-win_amd64.whl", hash = "sha256:b6937f5fe4e180aeee87de907a2fa982ded6f7f15d7218f78a083e4e1d68f2a0", size = 2757347, upload-time = "2026-04-20T23:35:21.283Z" }, + { url = "https://files.pythonhosted.org/packages/13/1b/708c0dca874acfad6d65314271859899a79007686f3a1f74e82a2ed4b645/psycopg2_binary-2.9.12-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:6f3b3de8a74ef8db215f22edffb19e32dc6fa41340456de7ec99efdc8a7b3ec2", size = 3712428, upload-time = "2026-04-20T23:35:23.453Z" }, + { url = "https://files.pythonhosted.org/packages/d6/39/ddbea9d4b4de6aca9431b6ed253f530f8a02d3b8f9bcfd0dbfe2b3de6fe4/psycopg2_binary-2.9.12-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1006fb62f0f0bc5ce256a832356c6262e91be43f5e4eb15b5eaf38079464caf2", size = 3823184, upload-time = "2026-04-20T23:35:25.92Z" }, + { url = "https://files.pythonhosted.org/packages/bf/a0/bc2fef74b106fa345567122a0659e6d94512ed7dc0131ec44c9e5aba3725/psycopg2_binary-2.9.12-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:840066105706cd2eb29b9a1c2329620056582a4bf3e8169dec5c447042d0869f", size = 4579157, upload-time = "2026-04-20T23:35:28.542Z" }, + { url = "https://files.pythonhosted.org/packages/57/d7/d4e3b2005d3de607ca4fbb0e8742e248056e52184a6b94ebda3c1c2c329b/psycopg2_binary-2.9.12-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:863f5d12241ebe1c76a72a04c2113b6dc905f90b9cef0e9be0efd994affd9354", size = 4274970, upload-time = "2026-04-20T23:35:30.418Z" }, + { url = "https://files.pythonhosted.org/packages/2e/42/c9853f8db3967fe08bcde11f53d53b85d351750cae726ce001cb68afa9c1/psycopg2_binary-2.9.12-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a99eaab34a9010f1a086b126de467466620a750634d114d20455f3a824aae033", size = 5895175, upload-time = "2026-04-20T23:35:33.584Z" }, + { url = "https://files.pythonhosted.org/packages/eb/fd/b82b5601a97630308bef079f545ffec481bbbc795c2ba5ec416a01d03f60/psycopg2_binary-2.9.12-cp314-cp314-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:ffdd7dc5463ccd61845ac37b7012d0f35a1548df9febe14f8dd549be4a0bc81e", size = 4110658, upload-time = "2026-04-20T23:35:35.638Z" }, + { url = "https://files.pythonhosted.org/packages/62/8c/32ca69b0389ef25dd22937bf9e8fbe2ce27aea20b05ded48c4ce4cb42475/psycopg2_binary-2.9.12-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:54a0dfecab1b48731f934e06139dfe11e24219fb6d0ceb32177cf0375f14c7b5", size = 3656251, upload-time = "2026-04-20T23:35:37.854Z" }, + { url = "https://files.pythonhosted.org/packages/c4/29/96992a2b59e3b9d730fcf9612d0a387305025dc867a9fc490a9e496e074e/psycopg2_binary-2.9.12-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:96937c9c5d891f772430f418a7a8b4691a90c3e6b93cf72b5bd7cad8cbca32a5", size = 3301810, upload-time = "2026-04-20T23:35:39.927Z" }, + { url = "https://files.pythonhosted.org/packages/56/ad/44b06659949b243ae10112cd3b20a197f9bf3e81d5651379b9eb889bfaad/psycopg2_binary-2.9.12-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:77b348775efd4cdab410ec6609d81ccecd1139c90265fa583a7255c8064bc03d", size = 3048977, upload-time = "2026-04-20T23:35:41.806Z" }, + { url = "https://files.pythonhosted.org/packages/1d/f2/10a1bcebadb6aa55e280e1f58975c36a7b560ea525184c7aa4064c466633/psycopg2_binary-2.9.12-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:527e6342b3e44c2f0544f6b8e927d60de7f163f5723b8f1dfa7d2a84298738cd", size = 3351466, upload-time = "2026-04-20T23:35:43.993Z" }, + { url = "https://files.pythonhosted.org/packages/20/be/b732c8418ffa5bcfda002890f5dc4c869fc17db66ff11f53b17cfe44afc0/psycopg2_binary-2.9.12-cp314-cp314-win_amd64.whl", hash = "sha256:f12ae41fcafadb39b2785e64a40f9db05d6de2ac114077457e0e7c597f3af980", size = 2848762, upload-time = "2026-04-20T23:35:46.421Z" }, +] + [[package]] name = "pyarrow" version = "25.0.0" From 91827f366bfe593bb3c7490a4380795b39c623c8 Mon Sep 17 00:00:00 2001 From: Tobias Macey Date: Thu, 23 Jul 2026 12:17:29 -0400 Subject: [PATCH 2/4] fix(dlt): resolve table_format from the explicit profile, not the env var Everything else in build_table_resource keys off the resolved profile, so leaving the table-format hint to fall back to DLT_PROFILE meant a targeted run or test could be configured for a deployed destination while hinting plain parquet -- or the reverse. active_table_format now takes an optional profile, matching dataset_name and bucket_root which already do. Also hoist the selection set out of the comprehension so it is built once rather than per candidate table. Both raised by Copilot on PR #2472. Co-Authored-By: Claude Opus 4.8 --- src/ol_dlt/ol_dlt/config.py | 11 ++++++++--- src/ol_dlt/ol_dlt/database.py | 12 ++++++------ src/ol_dlt/tests/test_database.py | 31 +++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 9 deletions(-) diff --git a/src/ol_dlt/ol_dlt/config.py b/src/ol_dlt/ol_dlt/config.py index 5eefecdf8..2386461b4 100644 --- a/src/ol_dlt/ol_dlt/config.py +++ b/src/ol_dlt/ol_dlt/config.py @@ -65,14 +65,19 @@ def active_profile() -> str: return os.getenv("DLT_PROFILE", DEFAULT_PROFILE) -def active_table_format() -> TableFormat: - """Return the table format for the active profile. +def active_table_format(profile: str | None = None) -> TableFormat: + """Return the table format for the (active) profile. ``"iceberg"`` for qa/production, ``"native"`` (plain filesystem parquet) for dev/ci/test. Use this in ``@dlt.resource(..., table_format=...)`` so a source body never branches on the environment itself. + + Takes an optional explicit ``profile`` for the same reason ``dataset_name`` + and ``bucket_root`` do: a caller that has already resolved a profile (a test, + or a targeted run) must not have the table format silently resolved from a + different one via ``DLT_PROFILE``. """ - return "iceberg" if active_profile() in ICEBERG_PROFILES else "native" + return "iceberg" if (profile or active_profile()) in ICEBERG_PROFILES else "native" def dataset_name(profile: str | None = None) -> str: diff --git a/src/ol_dlt/ol_dlt/database.py b/src/ol_dlt/ol_dlt/database.py index f47a3805b..ab1783f8b 100644 --- a/src/ol_dlt/ol_dlt/database.py +++ b/src/ol_dlt/ol_dlt/database.py @@ -225,7 +225,7 @@ def build_table_resource( ) return resource.with_name(raw_name).apply_hints( table_name=raw_name, - table_format=config.active_table_format(), + table_format=config.active_table_format(resolved_profile), ) @@ -243,11 +243,11 @@ def build_database_source( targeted local runs. profile: Override the active profile (mainly for tests). """ - selected = ( - [table for table in spec.tables if table.name in set(tables)] - if tables is not None - else list(spec.tables) - ) + if tables is None: + selected = list(spec.tables) + else: + wanted = set(tables) + selected = [table for table in spec.tables if table.name in wanted] if not selected: msg = f"{spec.name}: no tables selected from {tables!r}" raise ValueError(msg) diff --git a/src/ol_dlt/tests/test_database.py b/src/ol_dlt/tests/test_database.py index e00c16ede..a3404b737 100644 --- a/src/ol_dlt/tests/test_database.py +++ b/src/ol_dlt/tests/test_database.py @@ -105,6 +105,37 @@ def test_vault_credentials_are_injected_at_connection_time( } +def test_table_format_follows_the_explicit_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An explicit profile must drive the table format, not ``DLT_PROFILE``. + + Everything else in ``build_table_resource`` keys off the resolved profile, + so letting the table-format hint fall back to the ambient env var would + write plain parquet hints onto a run that is otherwise configured for a + deployed (Iceberg) destination, and vice versa. + """ + monkeypatch.setenv("DLT_PROFILE", "test") + monkeypatch.setenv("EXAMPLE_DB_HOST", "example.rds.amazonaws.com") + monkeypatch.setattr( + database.vault, "read_database_credentials", lambda *_a: ("u", "p") + ) + deployed = database.build_database_source(_SPEC, profile="production") + local = database.build_database_source(_SPEC, profile="dev") + assert ( + deployed.resources["raw__example__app__postgres__widget"] + .compute_table_schema() + .get("table_format") + == "iceberg" + ) + assert ( + local.resources["raw__example__app__postgres__widget"] + .compute_table_schema() + .get("table_format") + == "native" + ) + + def test_build_database_source_selects_a_subset() -> None: source = database.build_database_source(_SPEC, tables=["widget"], profile="dev") assert set(source.resources) == {"raw__example__app__postgres__widget"} From 416b58e18c31450114ab216429e2a72ae11685d6 Mon Sep 17 00:00:00 2001 From: Tobias Macey Date: Thu, 23 Jul 2026 13:10:22 -0400 Subject: [PATCH 3/4] feat(dlt): ingest Keycloak organization tables Organizations are how Keycloak represents the institutional affiliations that downstream models need to attribute a learner to, so leaving them out would have meant resolving people without resolving who they belong to. Membership needs no new join table: Keycloak models it as membership of the organization's backing group (org.group_id -> keycloak_group.id), tagged by user_group_membership.membership_type, and both of those were already in scope. Only the organization itself and its email domains were missing. Schema verified against Keycloak 26.7.0's Liquibase changelogs rather than assumed -- org and org_domain are created in jpa-changelog-25.0.0, and org_domain's primary key is the composite (id, name), not a single id. Co-Authored-By: Claude Opus 4.8 --- .../ol_dlt/sources/keycloak/__init__.py | 19 +++++++++++++++---- src/ol_dlt/tests/sources/test_keycloak.py | 18 ++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py b/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py index f5017aaa6..1f8f410cd 100644 --- a/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py +++ b/src/ol_dlt/ol_dlt/sources/keycloak/__init__.py @@ -12,10 +12,12 @@ ``keycloak`` database in the local-dev CloudNativePG cluster (see ``ol_dlt.database``). -Scope is a curated identity subset rather than the whole Keycloak schema. The -~80 remaining tables are Keycloak's own internal machinery (session, event, -migration and policy bookkeeping) with no downstream modelling value, and -several of them hold credential material. Notably absent, deliberately: +Scope is a curated identity subset rather than the whole Keycloak schema: +users, group membership, role assignment, organizations, and the realm/client +context needed to interpret them. The ~80 remaining tables are Keycloak's own +internal machinery (session, event, migration and policy bookkeeping) with no +downstream modelling value, and several of them hold credential material. +Notably absent, deliberately: credential password hashes and OTP secrets component_config identity-provider and LDAP bind secrets @@ -65,6 +67,15 @@ # --- role assignment ------------------------------------------------- DatabaseTable(name="keycloak_role", primary_key="id"), DatabaseTable(name="user_role_mapping", primary_key=("role_id", "user_id")), + # --- organizations --------------------------------------------------- + # Keycloak models organization membership as membership of the org's + # backing group (org.group_id -> keycloak_group.id), distinguished by + # user_group_membership.membership_type (MANAGED/UNMANAGED). Both of + # those tables are already ingested above, so no join table is needed + # here; org and org_domain supply the organization itself and the email + # domains that route users to it. + DatabaseTable(name="org", primary_key="id"), + DatabaseTable(name="org_domain", primary_key=("id", "name")), # --- realm / client context ----------------------------------------- DatabaseTable(name="realm", primary_key="id"), DatabaseTable( diff --git a/src/ol_dlt/tests/sources/test_keycloak.py b/src/ol_dlt/tests/sources/test_keycloak.py index 55ae0b0e0..f4c1f13fd 100644 --- a/src/ol_dlt/tests/sources/test_keycloak.py +++ b/src/ol_dlt/tests/sources/test_keycloak.py @@ -30,6 +30,24 @@ def test_client_secret_columns_are_excluded() -> None: assert set(client.excluded_columns) == {"secret", "registration_token"} +def test_organization_membership_is_resolvable() -> None: + """Orgs need their backing group tables, not a dedicated join table. + + Keycloak models organization membership as membership of the org's backing + group (``org.group_id`` -> ``keycloak_group.id``), tagged by + ``user_group_membership.membership_type``. Dropping either group table + would leave `org` present but its members unresolvable, so pin all four + together. + """ + selected = {table.name for table in keycloak.KEYCLOAK_SPEC.tables} + assert { + "org", + "org_domain", + "keycloak_group", + "user_group_membership", + } <= selected + + def test_every_table_declares_a_primary_key() -> None: assert all(table.primary_key for table in keycloak.KEYCLOAK_SPEC.tables) From 1c7295da8f7c0325d8d4e8403647d6998354db8d Mon Sep 17 00:00:00 2001 From: Tobias Macey Date: Thu, 23 Jul 2026 13:18:42 -0400 Subject: [PATCH 4/4] fix(dlt): guard the Vault credential cache and name a non-database mount Two findings from Sentry's review of PR #2472. The credential cache's purpose is that one lease serves a whole run, but the check-then-fetch was unsynchronized. Nothing races for it today -- extraction runs with workers = 1 -- except that setting lives in .dlt/config.toml for an entirely unrelated reason (a dlt injectable-context bug), so raising it would have silently turned one lease per run into one per worker. A lock makes the cache correct on its own terms rather than by coincidence. The response parsing assumed a database secrets engine's shape. Pointing vault_mount at a KV-v2 path -- which nests its payload under data.data -- is a plausible enough misconfiguration that it deserves the same named error the denied and not-found cases already get, instead of a bare KeyError. The concurrency test needed a delay in the fake client to be worth anything: without one the fake read returns before a second thread is ever scheduled inside the window, and the test passed with the lock removed. It now fails without the lock and passes with it. Co-Authored-By: Claude Opus 4.8 --- src/ol_dlt/ol_dlt/vault.py | 29 ++++++++++++++-- src/ol_dlt/tests/test_vault.py | 63 ++++++++++++++++++++++++++++++++-- 2 files changed, 87 insertions(+), 5 deletions(-) diff --git a/src/ol_dlt/ol_dlt/vault.py b/src/ol_dlt/ol_dlt/vault.py index d210736c9..a0196e2ac 100644 --- a/src/ol_dlt/ol_dlt/vault.py +++ b/src/ol_dlt/ol_dlt/vault.py @@ -19,6 +19,7 @@ import logging import os +import threading import time from pathlib import Path from typing import Any @@ -39,6 +40,13 @@ # (mount, role) -> (expires_at_monotonic, username, password) _CREDENTIAL_CACHE: dict[tuple[str, str], tuple[float, str, str]] = {} +# Guards the check-then-fetch above. Extraction is single-threaded today +# (``workers = 1`` in .dlt/config.toml), so nothing races for this lock — but +# that setting lives in another file for an unrelated reason (a dlt injectable- +# context bug), and raising it would otherwise silently turn "one lease per +# run" into "one lease per worker". +_CREDENTIAL_LOCK = threading.Lock() + def vault_address() -> str: """Return the Vault address for the current environment. @@ -88,6 +96,14 @@ def read_database_credentials(mount: str, role: str = "readonly") -> tuple[str, ingestion pipeline should ever ask for. """ cache_key = (mount, role) + with _CREDENTIAL_LOCK: + return _read_database_credentials_locked(cache_key, mount, role) + + +def _read_database_credentials_locked( + cache_key: tuple[str, str], mount: str, role: str +) -> tuple[str, str]: + """Return cached credentials, or issue and cache a fresh lease.""" cached = _CREDENTIAL_CACHE.get(cache_key) if cached and time.monotonic() < cached[0]: return cached[1], cached[2] @@ -108,8 +124,17 @@ def read_database_credentials(mount: str, role: str = "readonly") -> tuple[str, raise RuntimeError(msg) lease_duration = int(response.get("lease_duration") or 0) - username = response["data"]["username"] - password = response["data"]["password"] + issued = response.get("data") or {} + username, password = issued.get("username"), issued.get("password") + if not (username and password): + # The realistic cause is a mount that exists but isn't a database + # secrets engine (a KV-v2 mount, say, nests its payload under + # data.data), which otherwise surfaces as a bare KeyError. + msg = ( + f"Vault response from {vault_path!r} carries no username/password — " + f"is {mount!r} a database secrets-engine mount?" + ) + raise RuntimeError(msg) logger.info( "Issued Vault database credentials from %s (lease %ss)", vault_path, diff --git a/src/ol_dlt/tests/test_vault.py b/src/ol_dlt/tests/test_vault.py index ccda85fad..b32729c7e 100644 --- a/src/ol_dlt/tests/test_vault.py +++ b/src/ol_dlt/tests/test_vault.py @@ -1,5 +1,8 @@ """Tests for Vault credential resolution.""" +import threading +import time +from concurrent.futures import ThreadPoolExecutor from typing import Any import hvac @@ -14,15 +17,30 @@ def _clear_credential_cache() -> None: class FakeVaultClient: - """Records reads and returns a canned database-credentials response.""" - - def __init__(self, response: Any = None, error: Exception | None = None) -> None: # noqa: ANN401 + """Records reads and returns a canned database-credentials response. + + ``read_delay`` widens the window between a cache miss and the cache write, + which is what a real network round-trip to Vault does. Without it a fake + read returns so fast that a racing thread is never actually scheduled + inside the window, and a concurrency test passes whether or not the code + under test is correct. + """ + + def __init__( + self, + response: Any = None, # noqa: ANN401 + error: Exception | None = None, + read_delay: float = 0.0, + ) -> None: self.response = response self.error = error + self.read_delay = read_delay self.reads: list[str] = [] def read(self, path: str) -> Any: # noqa: ANN401 self.reads.append(path) + if self.read_delay: + time.sleep(self.read_delay) if self.error: raise self.error return self.response @@ -96,6 +114,45 @@ def test_credentials_are_refetched_once_the_lease_is_half_spent( assert len(client.reads) == 2 # noqa: PLR2004 +def test_concurrent_callers_share_one_lease( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Parallel extraction must not open a lease per worker. + + Extraction is single-threaded today, but the cache's whole purpose is + "one lease per run", so it should not quietly depend on that. + """ + client = FakeVaultClient( + _credentials_response(lease_duration=3600), read_delay=0.05 + ) + monkeypatch.setattr(vault, "_authenticated_client", lambda: client) + start = threading.Barrier(8) + + def _fetch() -> tuple[str, str]: + start.wait() # maximize the odds of a check-then-act overlap + return vault.read_database_credentials("postgres-keycloak") + + with ThreadPoolExecutor(max_workers=8) as pool: + results = [ + future.result() for future in [pool.submit(_fetch) for _ in range(8)] + ] + + assert len(client.reads) == 1 + assert set(results) == {("v-user-1", "v-pass-1")} + + +def test_a_response_without_credentials_names_the_likely_cause( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A non-database mount must not surface as a bare KeyError.""" + kv_v2_shaped = {"lease_duration": 0, "data": {"data": {"foo": "bar"}}} + monkeypatch.setattr( + vault, "_authenticated_client", lambda: FakeVaultClient(kv_v2_shaped) + ) + with pytest.raises(RuntimeError, match="database secrets-engine mount"): + vault.read_database_credentials("secret-data") + + def test_a_denied_path_names_the_policy_to_fix( monkeypatch: pytest.MonkeyPatch, ) -> None: