From ab8543797ea2e8b4c288fedc7d30795d2e265c81 Mon Sep 17 00:00:00 2001 From: vtulus <15098903-vtulus@users.noreply.gitlab.com> Date: Tue, 6 Oct 2026 17:03:47 +0200 Subject: [PATCH 1/5] feat(postgres): import key metadata for select-only roles --- CHANGELOG.md | 1 + datacontract/imports/postgres_importer.py | 214 +++++++-- docs/docs/imports/postgres.md | 6 + tests/fixtures/postgres/data/import.sql | 118 ++++- tests/test_import_postgres.py | 521 +++++++++++++++++++++- 5 files changed, 826 insertions(+), 34 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 679e3b61f..289d227dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `datacontract import postgres` imports primary and foreign keys for SELECT-only roles with catalog access. - `datacontract test`: check that foreign keys declared under `relationships` have no values missing from the referenced schema (#1691) ## [1.2.3] - 2026-10-05 diff --git a/datacontract/imports/postgres_importer.py b/datacontract/imports/postgres_importer.py index 7dca8d838..3918c8140 100644 --- a/datacontract/imports/postgres_importer.py +++ b/datacontract/imports/postgres_importer.py @@ -2,7 +2,9 @@ Reads table and column metadata from ``information_schema``, the same catalog ``datacontract test`` reads back to verify physical types, so an imported -contract passes on the first run without hand-editing. Comments come from +contract passes on the first run without hand-editing. Primary keys use +PostgreSQL catalog metadata when available, with the portable +information-schema result as a fallback. Comments come from ``pg_description`` via ``obj_description`` / ``col_description``, which ``information_schema`` does not expose. @@ -13,9 +15,14 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional -from open_data_contract_standard.model import OpenDataContractStandard, SchemaObject, SchemaProperty +from open_data_contract_standard.model import ( + OpenDataContractStandard, + Relationship, + SchemaObject, + SchemaProperty, +) from datacontract.config import Config from datacontract.engines.ibis.native_type import reconstruct_native_type @@ -66,17 +73,78 @@ ORDER BY table_name, ordinal_position """ -_PRIMARY_KEYS_QUERY = """ +_INFORMATION_SCHEMA_PRIMARY_KEYS_QUERY = """ SELECT kcu.table_name, kcu.column_name, kcu.ordinal_position FROM information_schema.table_constraints tc JOIN information_schema.key_column_usage kcu ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema + AND tc.table_name = kcu.table_name WHERE tc.constraint_type = 'PRIMARY KEY' AND tc.table_schema = %s ORDER BY kcu.table_name, kcu.ordinal_position """ +# pg_constraint is the preferred source of primary-key metadata. The portable +# information-schema query remains a fallback when reading the catalog fails. +# conkey retains the declared order of columns in composite primary keys. +_CATALOG_PRIMARY_KEYS_QUERY = """ + SELECT relation.relname AS table_name, + attribute.attname AS column_name, + con.oid AS constraint_oid, + cardinality(con.conkey) AS expected_column_count, + key_column.ordinality AS ordinal_position + FROM pg_catalog.pg_constraint AS con + JOIN pg_catalog.pg_class AS relation + ON relation.oid = con.conrelid + JOIN pg_catalog.pg_namespace AS namespace + ON namespace.oid = relation.relnamespace + CROSS JOIN LATERAL unnest(con.conkey) + WITH ORDINALITY AS key_column(attnum, ordinality) + JOIN pg_catalog.pg_attribute AS attribute + ON attribute.attrelid = relation.oid + AND attribute.attnum = key_column.attnum + AND NOT attribute.attisdropped + WHERE con.contype = 'p' + AND namespace.nspname = %s + ORDER BY relation.relname, key_column.ordinality +""" + +# Pairing conkey and confkey by ordinality preserves the source-to-target +# column mapping for both single-column and composite foreign keys. +_FOREIGN_KEYS_QUERY = """ + SELECT source_relation.relname AS table_name, + source_attribute.attname AS column_name, + target_relation.relname AS foreign_table_name, + target_namespace.nspname AS foreign_table_schema, + target_attribute.attname AS foreign_column_name, + con.oid AS constraint_oid, + cardinality(con.conkey) AS expected_column_count, + key_column.ordinality AS ordinal_position + FROM pg_catalog.pg_constraint AS con + JOIN pg_catalog.pg_class AS source_relation + ON source_relation.oid = con.conrelid + JOIN pg_catalog.pg_namespace AS source_namespace + ON source_namespace.oid = source_relation.relnamespace + JOIN pg_catalog.pg_class AS target_relation + ON target_relation.oid = con.confrelid + JOIN pg_catalog.pg_namespace AS target_namespace + ON target_namespace.oid = target_relation.relnamespace + CROSS JOIN LATERAL unnest(con.conkey, con.confkey) + WITH ORDINALITY AS key_column(source_attnum, target_attnum, ordinality) + JOIN pg_catalog.pg_attribute AS source_attribute + ON source_attribute.attrelid = source_relation.oid + AND source_attribute.attnum = key_column.source_attnum + AND NOT source_attribute.attisdropped + JOIN pg_catalog.pg_attribute AS target_attribute + ON target_attribute.attrelid = target_relation.oid + AND target_attribute.attnum = key_column.target_attnum + AND NOT target_attribute.attisdropped + WHERE con.contype = 'f' + AND source_namespace.nspname = %s + ORDER BY source_relation.relname, con.oid, key_column.ordinality +""" + class PostgresImporter(Importer): def import_source(self, source: str, import_args: dict, config: "Config | None" = None) -> OpenDataContractStandard: @@ -119,9 +187,16 @@ def import_postgres_from_connector( try: table_rows = _fetch(connection, _TABLES_QUERY, (schema,)) column_rows = _fetch(connection, _COLUMNS_QUERY, (schema,)) - # A user without access to information_schema constraints still gets a - # usable contract, just without primary keys. - primary_key_rows = _fetch(connection, _PRIMARY_KEYS_QUERY, (schema,), optional=True) + information_schema_primary_key_rows = _fetch( + connection, _INFORMATION_SCHEMA_PRIMARY_KEYS_QUERY, (schema,), optional=True + ) + catalog_primary_key_rows = _fetch(connection, _CATALOG_PRIMARY_KEYS_QUERY, (schema,), optional=True) + primary_key_rows = ( + catalog_primary_key_rows + if catalog_primary_key_rows is not None + else information_schema_primary_key_rows or [] + ) + foreign_key_rows = _fetch(connection, _FOREIGN_KEYS_QUERY, (schema,), optional=True) finally: connection.close() @@ -146,8 +221,33 @@ def import_postgres_from_connector( schema=schema, ) ] + selected_table_names = {table["table_name"] for table in selected} + selected_column_names = { + table_name: {row["column_name"] for row in column_rows if row["table_name"] == table_name} + for table_name in selected_table_names + } + if catalog_primary_key_rows is not None: + primary_key_rows = _complete_catalog_constraint_rows( + primary_key_rows, + lambda row: row["column_name"] in selected_column_names.get(row["table_name"], set()), + ) + selected_foreign_key_rows = [ + row + for row in foreign_key_rows or [] + if row["table_name"] in selected_table_names + and row["foreign_table_schema"] == schema + and row["foreign_table_name"] in selected_table_names + ] + if foreign_key_rows is not None: + selected_foreign_key_rows = _complete_catalog_constraint_rows( + selected_foreign_key_rows, + lambda row: ( + row["column_name"] in selected_column_names.get(row["table_name"], set()) + and row["foreign_column_name"] in selected_column_names.get(row["foreign_table_name"], set()) + ), + ) odcs.schema_ = [ - _create_schema(table, column_rows, primary_key_rows) + _create_schema(table, column_rows, primary_key_rows, selected_foreign_key_rows) for table in sorted(selected, key=lambda row: row["table_name"].lower()) ] report_unmapped_types(odcs) @@ -178,26 +278,26 @@ def postgres_connection(host: str, port: int, database: str, config: Optional[Co ) -def _fetch(connection, query: str, params: tuple, optional: bool = False) -> List[Dict[str, Any]]: +def _fetch(connection, query: str, params: tuple, optional: bool = False) -> Optional[List[Dict[str, Any]]]: """Run a catalog query and return its rows as dicts keyed by column name.""" - try: - with connection.cursor() as cursor: + with connection.cursor() as cursor: + try: cursor.execute(query, params) columns = [description[0] for description in cursor.description] return [dict(zip(columns, row)) for row in cursor.fetchall()] - except Exception as e: - if optional: - # Roll back so the failed statement doesn't poison the transaction. - connection.rollback() - return [] - raise DataContractException( - type="schema", - result="failed", - name="postgres catalog query failed", - reason=f"Could not read the Postgres catalog: {e}", - engine="datacontract-cli", - original_exception=e, - ) + except Exception as e: + if optional and getattr(e, "sqlstate", None) == "42501": + # Roll back so the failed statement doesn't poison the transaction. + connection.rollback() + return None + raise DataContractException( + type="schema", + result="failed", + name="postgres catalog query failed", + reason=f"Could not read the Postgres catalog: {e}", + engine="datacontract-cli", + original_exception=e, + ) def _select_tables(table_rows: List[Dict[str, Any]], tables: Optional[List[str]]) -> List[Dict[str, Any]]: @@ -207,26 +307,81 @@ def _select_tables(table_rows: List[Dict[str, Any]], tables: Optional[List[str]] return [row for row in table_rows if row["table_name"].lower() in wanted] +def _complete_catalog_constraint_rows( + rows: List[Dict[str, Any]], row_is_visible: Callable[[Dict[str, Any]], bool] +) -> List[Dict[str, Any]]: + """Return only complete, ordinally valid catalog constraints with visible columns.""" + rows_by_constraint = {} + for row in rows: + rows_by_constraint.setdefault(row["constraint_oid"], []).append(row) + + complete_rows = [] + for constraint_rows in rows_by_constraint.values(): + expected_column_count = constraint_rows[0]["expected_column_count"] + ordinals = {row["ordinal_position"] for row in constraint_rows} + if ( + len(constraint_rows) == expected_column_count + and ordinals == set(range(1, expected_column_count + 1)) + and all(row["expected_column_count"] == expected_column_count for row in constraint_rows) + and all(row_is_visible(row) for row in constraint_rows) + ): + complete_rows.extend(sorted(constraint_rows, key=lambda row: row["ordinal_position"])) + return complete_rows + + def _create_schema( table: Dict[str, Any], column_rows: List[Dict[str, Any]], primary_key_rows: List[Dict[str, Any]], + foreign_key_rows: List[Dict[str, Any]], ) -> SchemaObject: table_name = table["table_name"] primary_keys = { - row["column_name"]: index + 1 - for index, row in enumerate(row for row in primary_key_rows if row["table_name"] == table_name) + row["column_name"]: row["ordinal_position"] for row in primary_key_rows if row["table_name"] == table_name } - properties = [_create_property(row, primary_keys) for row in column_rows if row["table_name"] == table_name] - return create_schema_object( + relationships = {} + schema_relationships = [] + foreign_keys_by_constraint = {} + for row in foreign_key_rows: + if row["table_name"] == table_name: + foreign_keys_by_constraint.setdefault(row["constraint_oid"], []).append(row) + for constraint_rows in foreign_keys_by_constraint.values(): + if len(constraint_rows) == 1: + row = constraint_rows[0] + relationships.setdefault(row["column_name"], []).append( + Relationship(type="foreignKey", to=f"{row['foreign_table_name']}.{row['foreign_column_name']}") + ) + else: + schema_relationships.append( + Relationship( + type="foreignKey", + **{ + "from": [f"{table_name}.{row['column_name']}" for row in constraint_rows], + "to": [f"{row['foreign_table_name']}.{row['foreign_column_name']}" for row in constraint_rows], + }, + ) + ) + properties = [ + _create_property(row, primary_keys, relationships.get(row["column_name"])) + for row in column_rows + if row["table_name"] == table_name + ] + schema = create_schema_object( name=table_name, physical_type=table.get("table_type") or "table", description=_clean(table.get("remarks")), properties=properties or None, ) + if schema_relationships: + schema.relationships = schema_relationships + return schema -def _create_property(row: Dict[str, Any], primary_keys: Dict[str, int]) -> SchemaProperty: +def _create_property( + row: Dict[str, Any], + primary_keys: Dict[str, int], + relationships: Optional[List[Relationship]] = None, +) -> SchemaProperty: name = row["column_name"] max_length = row.get("character_maximum_length") precision = row.get("numeric_precision") @@ -252,6 +407,7 @@ def _create_property(row: Dict[str, Any], primary_keys: Dict[str, int]) -> Schem format=format, dimensions=dimensions, element_type=element_type, + relationships=relationships, ) diff --git a/docs/docs/imports/postgres.md b/docs/docs/imports/postgres.md index cdc6e13a3..1ad00d609 100644 --- a/docs/docs/imports/postgres.md +++ b/docs/docs/imports/postgres.md @@ -22,6 +22,12 @@ The generated contract includes a ready-to-test `servers` block, so you can run Credentials are provided as environment variables and are the same ones `datacontract test` uses: `DATACONTRACT_POSTGRES_USERNAME` and `DATACONTRACT_POSTGRES_PASSWORD` — see the [Postgres Reference](../reference/postgres.md). +## Key metadata and privileges + +When permitted, the importer reads PostgreSQL catalog metadata to enrich the contract with declared primary keys and foreign-key relationships. If catalog access is unavailable because the role lacks permission, the table and column import still succeeds: primary keys fall back to `information_schema` when available, and otherwise no key metadata is included. + +The supported read-only case is a role with full-table `SELECT`, schema `USAGE`, and ordinary read access to the PostgreSQL catalogs. Column-only grants can produce an incomplete contract, particularly when primary keys use the `information_schema` fallback. Import only the tables that belong together: a foreign-key relationship is included only when both endpoint tables and their properties are present in the selected schema. A declared `NOT VALID` foreign key is imported as a relationship, but its historical validation state is not preserved or reported. + Working from a DDL file instead of a live database? Use [`datacontract import sql --dialect postgres`](./sql.md). All options: **[`datacontract import postgres`](../commands/import/postgres.md)**. diff --git a/tests/fixtures/postgres/data/import.sql b/tests/fixtures/postgres/data/import.sql index 6f5d0995b..49d8e80f0 100644 --- a/tests/fixtures/postgres/data/import.sql +++ b/tests/fixtures/postgres/data/import.sql @@ -1,5 +1,10 @@ +CREATE TABLE public.customers ( + customer_id INT PRIMARY KEY +); + CREATE TABLE public.orders ( order_id VARCHAR(36) PRIMARY KEY, + customer_id INT NOT NULL REFERENCES public.customers(customer_id), order_total NUMERIC(10, 2), line_count INT NOT NULL, ordered_at TIMESTAMPTZ, @@ -9,9 +14,116 @@ CREATE TABLE public.orders ( COMMENT ON TABLE public.orders IS 'All orders'; COMMENT ON COLUMN public.orders.order_id IS 'The order id'; -INSERT INTO public.orders (order_id, order_total, line_count, ordered_at, payload) VALUES - ('CX-263-DU', 50.00, 2, '2023-06-16 13:12:56', '{"channel": "web"}'), - ('IK-894-MN', 47.50, 1, '2023-10-08 22:40:57', '{"channel": "app"}'); +INSERT INTO public.customers (customer_id) VALUES + (1); + +INSERT INTO public.orders (order_id, customer_id, order_total, line_count, ordered_at, payload) VALUES + ('CX-263-DU', 1, 50.00, 2, '2023-06-16 13:12:56', '{"channel": "web"}'), + ('IK-894-MN', 1, 47.50, 1, '2023-10-08 22:40:57', '{"channel": "app"}'); CREATE VIEW public.open_orders AS SELECT order_id FROM public.orders WHERE line_count > 1; + +CREATE SCHEMA select_only_keys; + +CREATE TABLE select_only_keys.simple_target ( + target_id INT PRIMARY KEY +); + +CREATE TABLE select_only_keys.simple_source ( + source_id INT PRIMARY KEY, + target_id INT NOT NULL REFERENCES select_only_keys.simple_target(target_id) +); + +CREATE TABLE select_only_keys.composite_target ( + first_id INT NOT NULL, + second_id INT NOT NULL, + PRIMARY KEY (first_id, second_id) +); + +CREATE TABLE select_only_keys.composite_source ( + source_second_id INT NOT NULL, + source_first_id INT NOT NULL, + PRIMARY KEY (source_second_id, source_first_id), + FOREIGN KEY (source_first_id, source_second_id) + REFERENCES select_only_keys.composite_target(first_id, second_id) +); + +CREATE ROLE select_only_key_reader LOGIN PASSWORD 'select-only-key-reader-password'; +GRANT USAGE ON SCHEMA select_only_keys TO select_only_key_reader; +GRANT SELECT ON ALL TABLES IN SCHEMA select_only_keys TO select_only_key_reader; + +CREATE SCHEMA key_visibility; +CREATE SCHEMA key_visibility_other; + +CREATE TABLE key_visibility.selected_target ( + target_id INT PRIMARY KEY +); + +CREATE TABLE key_visibility.selected_source ( + source_id INT PRIMARY KEY, + target_id INT NOT NULL REFERENCES key_visibility.selected_target(target_id) +); + +CREATE TABLE key_visibility.composite_target ( + first_id INT NOT NULL, + second_id INT NOT NULL, + PRIMARY KEY (first_id, second_id) +); + +CREATE TABLE key_visibility.composite_source ( + source_id INT PRIMARY KEY, + source_first_id INT NOT NULL, + source_second_id INT NOT NULL, + FOREIGN KEY (source_first_id, source_second_id) + REFERENCES key_visibility.composite_target(first_id, second_id) +); + +CREATE TABLE key_visibility_other.cross_schema_target ( + target_id INT PRIMARY KEY +); + +CREATE TABLE key_visibility.cross_schema_source ( + source_id INT PRIMARY KEY, + target_id INT NOT NULL REFERENCES key_visibility_other.cross_schema_target(target_id) +); + +CREATE TABLE key_visibility.same_named ( + local_second_id INT NOT NULL, + local_first_id INT NOT NULL, + PRIMARY KEY (local_second_id, local_first_id) +); + +CREATE TABLE key_visibility_other.same_named ( + other_id INT PRIMARY KEY +); + +CREATE TABLE key_visibility.partial_target ( + first_id INT NOT NULL, + second_id INT NOT NULL, + PRIMARY KEY (first_id, second_id) +); + +CREATE TABLE key_visibility.partial_source ( + source_first_id INT NOT NULL, + source_second_id INT NOT NULL, + PRIMARY KEY (source_first_id, source_second_id), + FOREIGN KEY (source_first_id, source_second_id) + REFERENCES key_visibility.partial_target(first_id, second_id) +); + +CREATE ROLE partial_key_reader LOGIN PASSWORD 'partial-key-reader-password'; +GRANT USAGE ON SCHEMA key_visibility TO partial_key_reader; +GRANT SELECT (source_first_id) ON key_visibility.partial_source TO partial_key_reader; +GRANT SELECT (first_id) ON key_visibility.partial_target TO partial_key_reader; + +CREATE TABLE key_visibility.column_reference_primary_key ( + first_id INT NOT NULL, + second_id INT NOT NULL, + PRIMARY KEY (second_id, first_id) +); + +CREATE ROLE column_reference_key_reader LOGIN PASSWORD 'column-reference-key-reader-password'; +GRANT USAGE ON SCHEMA key_visibility TO column_reference_key_reader; +GRANT SELECT ON key_visibility.column_reference_primary_key TO column_reference_key_reader; +GRANT REFERENCES (first_id) ON key_visibility.column_reference_primary_key TO column_reference_key_reader; diff --git a/tests/test_import_postgres.py b/tests/test_import_postgres.py index 70bf67c6f..8f02167cc 100644 --- a/tests/test_import_postgres.py +++ b/tests/test_import_postgres.py @@ -11,6 +11,7 @@ from datacontract.cli import app from datacontract.data_contract import DataContract +from datacontract.imports.postgres_importer import import_postgres_from_connector from datacontract.model.exceptions import DataContractException from datacontract.model.run import ResultEnum @@ -40,6 +41,320 @@ def _import(**kwargs): return DataContract.import_from_source("postgres", postgres.get_container_host_ip(), **kwargs) +class _InsufficientPrivilegeError(Exception): + sqlstate = "42501" + + +class _FetchFailure: + def __init__(self, exception): + self.exception = exception + + +class _CatalogConnection: + def __init__(self, results): + self.results = results + self.rollbacks = 0 + self.closed = False + + def cursor(self): + return _CatalogCursor(self) + + def rollback(self): + self.rollbacks += 1 + + def close(self): + self.closed = True + + +class _CatalogCursor: + def __init__(self, connection): + self.connection = connection + self.description = None + self.rows = None + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def execute(self, query, params): + result = self.connection.results[query] + if isinstance(result, Exception): + raise result + self.rows = result + if isinstance(result, _FetchFailure): + self.description = [("table_name",)] + return + self.description = [(column,) for column in result[0]] if result else [("table_name",)] + + def fetchall(self): + if isinstance(self.rows, _FetchFailure): + raise self.rows.exception + return [tuple(row[column] for column, *_ in self.description) for row in self.rows] + + +def _catalog_results(catalog_primary_keys, foreign_keys): + from datacontract.imports import postgres_importer + + return { + postgres_importer._TABLES_QUERY: [ + {"table_name": "orders", "table_type": "BASE TABLE", "remarks": None}, + {"table_name": "customers", "table_type": "BASE TABLE", "remarks": None}, + ], + postgres_importer._COLUMNS_QUERY: [ + { + "table_name": "orders", + "column_name": "id", + "data_type": "integer", + "character_maximum_length": None, + "numeric_precision": 32, + "numeric_scale": 0, + "is_nullable": "NO", + "remarks": None, + }, + { + "table_name": "orders", + "column_name": "customer_id", + "data_type": "integer", + "character_maximum_length": None, + "numeric_precision": 32, + "numeric_scale": 0, + "is_nullable": "YES", + "remarks": None, + }, + { + "table_name": "customers", + "column_name": "id", + "data_type": "integer", + "character_maximum_length": None, + "numeric_precision": 32, + "numeric_scale": 0, + "is_nullable": "NO", + "remarks": None, + }, + ], + postgres_importer._INFORMATION_SCHEMA_PRIMARY_KEYS_QUERY: [ + {"table_name": "orders", "column_name": "id", "ordinal_position": 1} + ], + postgres_importer._CATALOG_PRIMARY_KEYS_QUERY: catalog_primary_keys, + postgres_importer._FOREIGN_KEYS_QUERY: foreign_keys, + } + + +def _import_with_catalog_results(monkeypatch, catalog_primary_keys, foreign_keys): + connection = _CatalogConnection(_catalog_results(catalog_primary_keys, foreign_keys)) + monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) + return import_postgres_from_connector(host="localhost", database="postgres"), connection + + +def test_import_postgres_catalog_primary_keys_override_information_schema_when_empty(monkeypatch): + result, connection = _import_with_catalog_results(monkeypatch, [], []) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + assert orders.properties[0].primaryKey is None + assert connection.rollbacks == 0 + + +def test_import_postgres_falls_back_to_information_schema_primary_keys_after_catalog_failure(monkeypatch): + foreign_key = { + "table_name": "orders", + "column_name": "customer_id", + "foreign_table_name": "customers", + "foreign_table_schema": "public", + "foreign_column_name": "id", + "constraint_oid": 1, + "expected_column_count": 1, + "ordinal_position": 1, + } + result, connection = _import_with_catalog_results( + monkeypatch, _InsufficientPrivilegeError("catalog PK unavailable"), [foreign_key] + ) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + assert orders.properties[0].primaryKey is True + assert orders.properties[1].relationships[0].to == "customers.id" + assert connection.rollbacks == 1 + + +def test_import_postgres_omits_foreign_keys_after_catalog_failure(monkeypatch): + catalog_primary_key = { + "table_name": "orders", + "column_name": "id", + "constraint_oid": 1, + "expected_column_count": 1, + "ordinal_position": 1, + } + result, connection = _import_with_catalog_results( + monkeypatch, [catalog_primary_key], _InsufficientPrivilegeError("catalog FK unavailable") + ) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + assert orders.properties[0].primaryKey is True + assert orders.properties[1].relationships is None + assert connection.rollbacks == 1 + + +def test_import_postgres_keeps_composite_catalog_keys_atomic_and_ordered(monkeypatch): + from datacontract.imports import postgres_importer + + results = _catalog_results( + [ + { + "table_name": "orders", + "column_name": "customer_id", + "constraint_oid": 1, + "expected_column_count": 2, + "ordinal_position": 2, + }, + { + "table_name": "orders", + "column_name": "id", + "constraint_oid": 1, + "expected_column_count": 2, + "ordinal_position": 1, + }, + ], + [ + { + "table_name": "orders", + "column_name": "customer_id", + "foreign_table_name": "customers", + "foreign_table_schema": "public", + "foreign_column_name": "id", + "constraint_oid": 2, + "expected_column_count": 2, + "ordinal_position": 2, + }, + { + "table_name": "orders", + "column_name": "id", + "foreign_table_name": "customers", + "foreign_table_schema": "public", + "foreign_column_name": "id", + "constraint_oid": 2, + "expected_column_count": 2, + "ordinal_position": 1, + }, + ], + ) + results[postgres_importer._COLUMNS_QUERY].append( + { + "table_name": "customers", + "column_name": "customer_id", + "data_type": "integer", + "character_maximum_length": None, + "numeric_precision": 32, + "numeric_scale": 0, + "is_nullable": "NO", + "remarks": None, + } + ) + results[postgres_importer._FOREIGN_KEYS_QUERY][0]["foreign_column_name"] = "customer_id" + results[postgres_importer._FOREIGN_KEYS_QUERY][1]["foreign_column_name"] = "id" + connection = _CatalogConnection(results) + monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) + + result = import_postgres_from_connector(host="localhost", database="postgres") + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + assert [property.primaryKeyPosition for property in orders.properties if property.primaryKey] == [1, 2] + assert orders.properties[0].relationships is None + assert orders.properties[1].relationships is None + assert len(orders.relationships) == 1 + assert orders.relationships[0].from_ == ["orders.id", "orders.customer_id"] + assert orders.relationships[0].to == ["customers.id", "customers.customer_id"] + assert DataContract(data_contract_str=result.to_yaml()).lint().result == ResultEnum.passed + + +def test_import_postgres_omits_incomplete_catalog_constraints(monkeypatch): + incomplete_primary_key = { + "table_name": "orders", + "column_name": "id", + "constraint_oid": 1, + "expected_column_count": 2, + "ordinal_position": 1, + } + incomplete_foreign_key = { + "table_name": "orders", + "column_name": "customer_id", + "foreign_table_name": "customers", + "foreign_table_schema": "public", + "foreign_column_name": "id", + "constraint_oid": 2, + "expected_column_count": 2, + "ordinal_position": 1, + } + + result, _ = _import_with_catalog_results(monkeypatch, [incomplete_primary_key], [incomplete_foreign_key]) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + assert all(property.primaryKey is None for property in orders.properties) + assert orders.properties[1].relationships is None + assert orders.relationships is None + + +def test_import_postgres_keeps_distinct_foreign_keys_on_the_same_source_property(monkeypatch): + foreign_keys = [ + { + "table_name": "orders", + "column_name": "customer_id", + "foreign_table_name": "customers", + "foreign_table_schema": "public", + "foreign_column_name": "id", + "constraint_oid": constraint_oid, + "expected_column_count": 1, + "ordinal_position": 1, + } + for constraint_oid in (1, 2) + ] + + result, _ = _import_with_catalog_results(monkeypatch, [], foreign_keys) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + customer_id = next(property for property in orders.properties if property.name == "customer_id") + assert [relationship.to for relationship in customer_id.relationships] == ["customers.id", "customers.id"] + + +def test_import_postgres_raises_when_a_required_catalog_query_fails(monkeypatch): + from datacontract.imports import postgres_importer + + connection = _CatalogConnection(_catalog_results([], [])) + connection.results[postgres_importer._TABLES_QUERY] = RuntimeError("tables unavailable") + monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) + + with pytest.raises(DataContractException, match="Could not read the Postgres catalog"): + import_postgres_from_connector(host="localhost", database="postgres") + + assert connection.rollbacks == 0 + + +def test_import_postgres_raises_for_an_optional_catalog_failure_that_is_not_an_access_error(monkeypatch): + connection = _CatalogConnection(_catalog_results(RuntimeError("catalog syntax error"), [])) + monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) + + with pytest.raises(DataContractException, match="Could not read the Postgres catalog") as exc_info: + import_postgres_from_connector(host="localhost", database="postgres") + + assert isinstance(exc_info.value.original_exception, RuntimeError) + assert connection.rollbacks == 0 + + +def test_import_postgres_raises_when_a_required_catalog_row_cannot_be_processed(monkeypatch): + failure = RuntimeError("table rows unavailable") + connection = _CatalogConnection(_catalog_results([], [])) + from datacontract.imports import postgres_importer + + connection.results[postgres_importer._TABLES_QUERY] = _FetchFailure(failure) + monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) + + with pytest.raises(DataContractException, match="Could not read the Postgres catalog") as exc_info: + import_postgres_from_connector(host="localhost", database="postgres") + + assert exc_info.value.original_exception is failure + assert connection.rollbacks == 0 + + def test_import_postgres(): result = _import(schema="public", postgres_table=["orders"]) @@ -73,6 +388,10 @@ def test_import_postgres(): required: true primaryKey: true primaryKeyPosition: 1 + - name: customer_id + logicalType: integer + physicalType: integer + required: true - name: order_total logicalType: number physicalType: numeric(10,2) @@ -97,6 +416,204 @@ def test_import_postgres(): assert yaml.safe_load(result.to_yaml()) == yaml.safe_load(expected) +def test_import_postgres_imports_foreign_key_relationships(): + result = _import(schema="public", postgres_table=["customers", "orders"]) + + orders = next(schema for schema in result.schema_ if schema.name == "orders") + customer_id = next(property for property in orders.properties if property.name == "customer_id") + relationship = customer_id.relationships[0] + + assert relationship.type == "foreignKey" + assert relationship.to == "customers.customer_id" + orders_yaml = next(schema for schema in yaml.safe_load(result.to_yaml())["schema"] if schema["name"] == "orders") + customer_id_yaml = next(property for property in orders_yaml["properties"] if property["name"] == "customer_id") + assert "from" not in customer_id_yaml["relationships"][0] + + +def test_import_postgres_recovers_keys_for_a_select_only_role(monkeypatch): + import psycopg + + reader = "select_only_key_reader" + password = "select-only-key-reader-password" + schema = "select_only_keys" + host = postgres.get_container_host_ip() + port = postgres.get_exposed_port(5432) + + with psycopg.connect( + dbname=postgres.dbname, user=reader, password=password, host=host, port=port, autocommit=True + ) as connection: + with connection.cursor() as cursor: + cursor.execute( + """ + WITH application_tables AS ( + SELECT relation.oid + FROM pg_catalog.pg_class AS relation + JOIN pg_catalog.pg_namespace AS namespace ON namespace.oid = relation.relnamespace + WHERE namespace.nspname = %s + AND relation.relkind = 'r' + ) + SELECT bool_and(has_schema_privilege(current_user, %s, 'USAGE')), + bool_and(has_table_privilege(current_user, oid, 'SELECT')), + NOT bool_or(has_table_privilege(current_user, oid, 'INSERT')), + NOT bool_or(has_table_privilege(current_user, oid, 'UPDATE')), + NOT bool_or(has_table_privilege(current_user, oid, 'DELETE')), + NOT bool_or(has_table_privilege(current_user, oid, 'REFERENCES')), + bool_and(has_table_privilege(current_user, 'pg_catalog.pg_constraint', 'SELECT')), + bool_and(has_table_privilege(current_user, 'pg_catalog.pg_class', 'SELECT')), + bool_and(has_table_privilege(current_user, 'pg_catalog.pg_namespace', 'SELECT')), + bool_and(has_table_privilege(current_user, 'pg_catalog.pg_attribute', 'SELECT')) + FROM application_tables + """, + (schema, schema), + ) + privileges = cursor.fetchone() + + assert privileges == (True, True, True, True, True, True, True, True, True, True) + monkeypatch.setenv("DATACONTRACT_POSTGRES_USERNAME", reader) + monkeypatch.setenv("DATACONTRACT_POSTGRES_PASSWORD", password) + + result = _import(schema=schema) + + simple_source = next(schema_object for schema_object in result.schema_ if schema_object.name == "simple_source") + simple_target = next(schema_object for schema_object in result.schema_ if schema_object.name == "simple_target") + composite_source = next( + schema_object for schema_object in result.schema_ if schema_object.name == "composite_source" + ) + composite_target = next( + schema_object for schema_object in result.schema_ if schema_object.name == "composite_target" + ) + + assert [property.name for property in simple_source.properties if property.primaryKey] == ["source_id"] + assert [property.name for property in simple_target.properties if property.primaryKey] == ["target_id"] + assert [property.name for property in composite_source.properties if property.primaryKey] == [ + "source_second_id", + "source_first_id", + ] + assert [property.primaryKeyPosition for property in composite_source.properties if property.primaryKey] == [1, 2] + assert [property.name for property in composite_target.properties if property.primaryKey] == [ + "first_id", + "second_id", + ] + assert [property.primaryKeyPosition for property in composite_target.properties if property.primaryKey] == [1, 2] + simple_fk = next(property for property in simple_source.properties if property.name == "target_id").relationships[0] + assert simple_fk.to == "simple_target.target_id" + assert simple_fk.from_ is None + assert all(property.relationships is None for property in composite_source.properties) + assert composite_source.relationships[0].from_ == [ + "composite_source.source_first_id", + "composite_source.source_second_id", + ] + assert composite_source.relationships[0].to == [ + "composite_target.first_id", + "composite_target.second_id", + ] + assert DataContract(data_contract_str=result.to_yaml()).lint().result == ResultEnum.passed + + +def test_import_postgres_omits_relationships_outside_selected_tables_and_other_schemas(): + source_only = _import(schema="key_visibility", postgres_table=["selected_source", "composite_source"]) + selected_source = next(schema for schema in source_only.schema_ if schema.name == "selected_source") + composite_source = next(schema for schema in source_only.schema_ if schema.name == "composite_source") + + assert ( + next(property for property in selected_source.properties if property.name == "target_id").relationships is None + ) + assert composite_source.relationships is None + + both_endpoints = _import( + schema="key_visibility", + postgres_table=["selected_source", "selected_target", "composite_source", "composite_target"], + ) + selected_source = next(schema for schema in both_endpoints.schema_ if schema.name == "selected_source") + composite_source = next(schema for schema in both_endpoints.schema_ if schema.name == "composite_source") + + assert next(property for property in selected_source.properties if property.name == "target_id").relationships[ + 0 + ].to == ("selected_target.target_id") + assert composite_source.relationships[0].from_ == [ + "composite_source.source_first_id", + "composite_source.source_second_id", + ] + assert composite_source.relationships[0].to == [ + "composite_target.first_id", + "composite_target.second_id", + ] + + cross_schema = _import(schema="key_visibility", postgres_table=["cross_schema_source"]) + cross_schema_source = cross_schema.schema_[0] + assert ( + next(property for property in cross_schema_source.properties if property.name == "target_id").relationships + is None + ) + + +def test_import_postgres_isolates_same_named_objects_in_other_schemas(): + result = _import(schema="key_visibility", postgres_table=["same_named"]) + + same_named = result.schema_[0] + assert [property.name for property in same_named.properties if property.primaryKey] == [ + "local_second_id", + "local_first_id", + ] + assert [property.primaryKeyPosition for property in same_named.properties if property.primaryKey] == [1, 2] + assert "other_id" not in [property.name for property in same_named.properties] + + +def test_import_postgres_omits_catalog_constraints_with_partially_visible_columns(monkeypatch): + monkeypatch.setenv("DATACONTRACT_POSTGRES_USERNAME", "partial_key_reader") + monkeypatch.setenv("DATACONTRACT_POSTGRES_PASSWORD", "partial-key-reader-password") + + result = _import(schema="key_visibility", postgres_table=["partial_source", "partial_target"]) + partial_source = next(schema for schema in result.schema_ if schema.name == "partial_source") + partial_target = next(schema for schema in result.schema_ if schema.name == "partial_target") + + assert [property.name for property in partial_source.properties] == ["source_first_id"] + assert [property.name for property in partial_target.properties] == ["first_id"] + assert all(property.primaryKey is None for property in partial_source.properties + partial_target.properties) + assert partial_source.relationships is None + + +def test_import_postgres_catalog_primary_keys_override_partial_column_references(monkeypatch): + import psycopg + + reader = "column_reference_key_reader" + password = "column-reference-key-reader-password" + host = postgres.get_container_host_ip() + port = postgres.get_exposed_port(5432) + + with psycopg.connect( + dbname=postgres.dbname, user=reader, password=password, host=host, port=port, autocommit=True + ) as connection: + with connection.cursor() as cursor: + cursor.execute( + """ + SELECT has_column_privilege( + current_user, + 'key_visibility.column_reference_primary_key', + 'first_id', + 'REFERENCES' + ), + has_column_privilege( + current_user, + 'key_visibility.column_reference_primary_key', + 'second_id', + 'REFERENCES' + ) + """ + ) + assert cursor.fetchone() == (True, False) + + monkeypatch.setenv("DATACONTRACT_POSTGRES_USERNAME", reader) + monkeypatch.setenv("DATACONTRACT_POSTGRES_PASSWORD", password) + result = _import(schema="key_visibility", postgres_table=["column_reference_primary_key"]) + + primary_key = [property for property in result.schema_[0].properties if property.primaryKey] + assert {property.name: property.primaryKeyPosition for property in primary_key} == { + "second_id": 1, + "first_id": 2, + } + + def test_import_postgres_produces_a_valid_contract(): result = _import(schema="public") @@ -118,8 +635,8 @@ def test_imported_contract_passes_test_without_editing(): def test_import_postgres_imports_all_tables_by_default(): result = _import() - assert [schema.name for schema in result.schema_] == ["open_orders", "orders"] - assert result.schema_[0].physicalType == "view" + assert [schema.name for schema in result.schema_] == ["customers", "open_orders", "orders"] + assert result.schema_[1].physicalType == "view" def test_import_postgres_defaults_to_the_public_schema(): From 813d74f422021dc034fda5d26df22aa88c4758b5 Mon Sep 17 00:00:00 2001 From: vtulus <15098903-vtulus@users.noreply.gitlab.com> Date: Fri, 9 Oct 2026 12:52:34 +0200 Subject: [PATCH 2/5] Address Postgres key metadata review feedback --- datacontract/imports/postgres_importer.py | 61 +-- docs/docs/imports/postgres.md | 6 +- tests/fixtures/postgres/data/import.sql | 33 ++ tests/test_import_postgres.py | 446 ++++------------------ 4 files changed, 138 insertions(+), 408 deletions(-) diff --git a/datacontract/imports/postgres_importer.py b/datacontract/imports/postgres_importer.py index 3918c8140..d2fb2bf80 100644 --- a/datacontract/imports/postgres_importer.py +++ b/datacontract/imports/postgres_importer.py @@ -15,6 +15,7 @@ from __future__ import annotations +import logging from typing import Any, Callable, Dict, List, Optional from open_data_contract_standard.model import ( @@ -40,6 +41,8 @@ DEFAULT_PORT = 5432 DEFAULT_SCHEMA = "public" +logger = logging.getLogger(__name__) + # to_regclass() yields NULL instead of raising for objects the user cannot see, # so a missing comment never fails the whole import. _TABLES_QUERY = """ @@ -92,7 +95,6 @@ SELECT relation.relname AS table_name, attribute.attname AS column_name, con.oid AS constraint_oid, - cardinality(con.conkey) AS expected_column_count, key_column.ordinality AS ordinal_position FROM pg_catalog.pg_constraint AS con JOIN pg_catalog.pg_class AS relation @@ -119,7 +121,6 @@ target_namespace.nspname AS foreign_table_schema, target_attribute.attname AS foreign_column_name, con.oid AS constraint_oid, - cardinality(con.conkey) AS expected_column_count, key_column.ordinality AS ordinal_position FROM pg_catalog.pg_constraint AS con JOIN pg_catalog.pg_class AS source_relation @@ -142,6 +143,13 @@ AND NOT target_attribute.attisdropped WHERE con.contype = 'f' AND source_namespace.nspname = %s + -- Skip the per-partition clones Postgres adds for a foreign key that references a partitioned table. + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint AS parent_constraint + WHERE parent_constraint.oid = con.conparentid + AND parent_constraint.conrelid = con.conrelid + ) ORDER BY source_relation.relname, con.oid, key_column.ordinality """ @@ -227,19 +235,34 @@ def import_postgres_from_connector( for table_name in selected_table_names } if catalog_primary_key_rows is not None: - primary_key_rows = _complete_catalog_constraint_rows( + primary_key_rows = _visible_constraint_rows( primary_key_rows, lambda row: row["column_name"] in selected_column_names.get(row["table_name"], set()), ) + foreign_key_rows_from_selected_tables = [ + row for row in foreign_key_rows or [] if row["table_name"] in selected_table_names + ] + omitted_foreign_keys = {} + for row in foreign_key_rows_from_selected_tables: + if row["foreign_table_schema"] != schema or row["foreign_table_name"] not in selected_table_names: + omitted_foreign_keys.setdefault(row["constraint_oid"], row) + for row in omitted_foreign_keys.values(): + logger.warning( + "Omitting foreign key from %s.%s to %s.%s because the target table is not included in the imported contract.", + schema, + row["table_name"], + row["foreign_table_schema"], + row["foreign_table_name"], + ) + selected_foreign_key_rows = [ row - for row in foreign_key_rows or [] - if row["table_name"] in selected_table_names - and row["foreign_table_schema"] == schema + for row in foreign_key_rows_from_selected_tables + if row["foreign_table_schema"] == schema and row["foreign_table_name"] in selected_table_names ] if foreign_key_rows is not None: - selected_foreign_key_rows = _complete_catalog_constraint_rows( + selected_foreign_key_rows = _visible_constraint_rows( selected_foreign_key_rows, lambda row: ( row["column_name"] in selected_column_names.get(row["table_name"], set()) @@ -286,7 +309,7 @@ def _fetch(connection, query: str, params: tuple, optional: bool = False) -> Opt columns = [description[0] for description in cursor.description] return [dict(zip(columns, row)) for row in cursor.fetchall()] except Exception as e: - if optional and getattr(e, "sqlstate", None) == "42501": + if optional: # Roll back so the failed statement doesn't poison the transaction. connection.rollback() return None @@ -307,26 +330,12 @@ def _select_tables(table_rows: List[Dict[str, Any]], tables: Optional[List[str]] return [row for row in table_rows if row["table_name"].lower() in wanted] -def _complete_catalog_constraint_rows( +def _visible_constraint_rows( rows: List[Dict[str, Any]], row_is_visible: Callable[[Dict[str, Any]], bool] ) -> List[Dict[str, Any]]: - """Return only complete, ordinally valid catalog constraints with visible columns.""" - rows_by_constraint = {} - for row in rows: - rows_by_constraint.setdefault(row["constraint_oid"], []).append(row) - - complete_rows = [] - for constraint_rows in rows_by_constraint.values(): - expected_column_count = constraint_rows[0]["expected_column_count"] - ordinals = {row["ordinal_position"] for row in constraint_rows} - if ( - len(constraint_rows) == expected_column_count - and ordinals == set(range(1, expected_column_count + 1)) - and all(row["expected_column_count"] == expected_column_count for row in constraint_rows) - and all(row_is_visible(row) for row in constraint_rows) - ): - complete_rows.extend(sorted(constraint_rows, key=lambda row: row["ordinal_position"])) - return complete_rows + """Drop constraints with a column the role can't read, so its name stays out of the contract.""" + hidden = {row["constraint_oid"] for row in rows if not row_is_visible(row)} + return [row for row in rows if row["constraint_oid"] not in hidden] def _create_schema( diff --git a/docs/docs/imports/postgres.md b/docs/docs/imports/postgres.md index 1ad00d609..92eda451b 100644 --- a/docs/docs/imports/postgres.md +++ b/docs/docs/imports/postgres.md @@ -6,7 +6,7 @@ description: "Create a data contract from a Postgres schema." # Import: Postgres -Creates a data contract from a Postgres schema by reading table metadata from `information_schema` — including column types with length and precision, nullability, primary keys, and the comments stored in `pg_description`. Works with Postgres and Postgres-compatible databases (e.g. RisingWave). +Creates a data contract from a Postgres schema. Tables and columns are read from `information_schema`, comments from `pg_description`, and key metadata from `pg_catalog`. Works with Postgres and Postgres-compatible databases (e.g. RisingWave). ```bash datacontract import postgres \ @@ -24,9 +24,7 @@ Credentials are provided as environment variables and are the same ones `datacon ## Key metadata and privileges -When permitted, the importer reads PostgreSQL catalog metadata to enrich the contract with declared primary keys and foreign-key relationships. If catalog access is unavailable because the role lacks permission, the table and column import still succeeds: primary keys fall back to `information_schema` when available, and otherwise no key metadata is included. - -The supported read-only case is a role with full-table `SELECT`, schema `USAGE`, and ordinary read access to the PostgreSQL catalogs. Column-only grants can produce an incomplete contract, particularly when primary keys use the `information_schema` fallback. Import only the tables that belong together: a foreign-key relationship is included only when both endpoint tables and their properties are present in the selected schema. A declared `NOT VALID` foreign key is imported as a relationship, but its historical validation state is not preserved or reported. +Primary keys and foreign keys are read from the Postgres catalog, so a role with `SELECT` on the tables and `USAGE` on the schema is enough. A foreign key is included only when both of its tables are imported from the same schema. Working from a DDL file instead of a live database? Use [`datacontract import sql --dialect postgres`](./sql.md). diff --git a/tests/fixtures/postgres/data/import.sql b/tests/fixtures/postgres/data/import.sql index 49d8e80f0..1766e1ab2 100644 --- a/tests/fixtures/postgres/data/import.sql +++ b/tests/fixtures/postgres/data/import.sql @@ -24,6 +24,39 @@ INSERT INTO public.orders (order_id, customer_id, order_total, line_count, order CREATE VIEW public.open_orders AS SELECT order_id FROM public.orders WHERE line_count > 1; +CREATE SCHEMA partitioned_foreign_keys; + +CREATE TABLE partitioned_foreign_keys.partitioned_target ( + target_id INT PRIMARY KEY +) PARTITION BY RANGE (target_id); + +CREATE TABLE partitioned_foreign_keys.partitioned_target_low + PARTITION OF partitioned_foreign_keys.partitioned_target + FOR VALUES FROM (0) TO (100); + +CREATE TABLE partitioned_foreign_keys.partitioned_target_high + PARTITION OF partitioned_foreign_keys.partitioned_target + FOR VALUES FROM (100) TO (200); + +CREATE TABLE partitioned_foreign_keys.partitioned_source ( + source_id INT PRIMARY KEY, + target_id INT NOT NULL REFERENCES partitioned_foreign_keys.partitioned_target(target_id) +); + +CREATE TABLE partitioned_foreign_keys.partitioned_source_parent ( + source_id INT NOT NULL, + target_id INT NOT NULL REFERENCES partitioned_foreign_keys.partitioned_target(target_id), + PRIMARY KEY (source_id, target_id) +) PARTITION BY RANGE (source_id); + +CREATE TABLE partitioned_foreign_keys.partitioned_source_low + PARTITION OF partitioned_foreign_keys.partitioned_source_parent + FOR VALUES FROM (0) TO (100); + +CREATE TABLE partitioned_foreign_keys.partitioned_source_high + PARTITION OF partitioned_foreign_keys.partitioned_source_parent + FOR VALUES FROM (100) TO (200); + CREATE SCHEMA select_only_keys; CREATE TABLE select_only_keys.simple_target ( diff --git a/tests/test_import_postgres.py b/tests/test_import_postgres.py index 8f02167cc..a0b039d13 100644 --- a/tests/test_import_postgres.py +++ b/tests/test_import_postgres.py @@ -1,5 +1,6 @@ """Tests for the Postgres importer, run against a real Postgres container.""" +import logging from pathlib import Path from unittest.mock import patch @@ -11,7 +12,7 @@ from datacontract.cli import app from datacontract.data_contract import DataContract -from datacontract.imports.postgres_importer import import_postgres_from_connector +from datacontract.imports import postgres_importer from datacontract.model.exceptions import DataContractException from datacontract.model.run import ResultEnum @@ -41,320 +42,6 @@ def _import(**kwargs): return DataContract.import_from_source("postgres", postgres.get_container_host_ip(), **kwargs) -class _InsufficientPrivilegeError(Exception): - sqlstate = "42501" - - -class _FetchFailure: - def __init__(self, exception): - self.exception = exception - - -class _CatalogConnection: - def __init__(self, results): - self.results = results - self.rollbacks = 0 - self.closed = False - - def cursor(self): - return _CatalogCursor(self) - - def rollback(self): - self.rollbacks += 1 - - def close(self): - self.closed = True - - -class _CatalogCursor: - def __init__(self, connection): - self.connection = connection - self.description = None - self.rows = None - - def __enter__(self): - return self - - def __exit__(self, *args): - return False - - def execute(self, query, params): - result = self.connection.results[query] - if isinstance(result, Exception): - raise result - self.rows = result - if isinstance(result, _FetchFailure): - self.description = [("table_name",)] - return - self.description = [(column,) for column in result[0]] if result else [("table_name",)] - - def fetchall(self): - if isinstance(self.rows, _FetchFailure): - raise self.rows.exception - return [tuple(row[column] for column, *_ in self.description) for row in self.rows] - - -def _catalog_results(catalog_primary_keys, foreign_keys): - from datacontract.imports import postgres_importer - - return { - postgres_importer._TABLES_QUERY: [ - {"table_name": "orders", "table_type": "BASE TABLE", "remarks": None}, - {"table_name": "customers", "table_type": "BASE TABLE", "remarks": None}, - ], - postgres_importer._COLUMNS_QUERY: [ - { - "table_name": "orders", - "column_name": "id", - "data_type": "integer", - "character_maximum_length": None, - "numeric_precision": 32, - "numeric_scale": 0, - "is_nullable": "NO", - "remarks": None, - }, - { - "table_name": "orders", - "column_name": "customer_id", - "data_type": "integer", - "character_maximum_length": None, - "numeric_precision": 32, - "numeric_scale": 0, - "is_nullable": "YES", - "remarks": None, - }, - { - "table_name": "customers", - "column_name": "id", - "data_type": "integer", - "character_maximum_length": None, - "numeric_precision": 32, - "numeric_scale": 0, - "is_nullable": "NO", - "remarks": None, - }, - ], - postgres_importer._INFORMATION_SCHEMA_PRIMARY_KEYS_QUERY: [ - {"table_name": "orders", "column_name": "id", "ordinal_position": 1} - ], - postgres_importer._CATALOG_PRIMARY_KEYS_QUERY: catalog_primary_keys, - postgres_importer._FOREIGN_KEYS_QUERY: foreign_keys, - } - - -def _import_with_catalog_results(monkeypatch, catalog_primary_keys, foreign_keys): - connection = _CatalogConnection(_catalog_results(catalog_primary_keys, foreign_keys)) - monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) - return import_postgres_from_connector(host="localhost", database="postgres"), connection - - -def test_import_postgres_catalog_primary_keys_override_information_schema_when_empty(monkeypatch): - result, connection = _import_with_catalog_results(monkeypatch, [], []) - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - assert orders.properties[0].primaryKey is None - assert connection.rollbacks == 0 - - -def test_import_postgres_falls_back_to_information_schema_primary_keys_after_catalog_failure(monkeypatch): - foreign_key = { - "table_name": "orders", - "column_name": "customer_id", - "foreign_table_name": "customers", - "foreign_table_schema": "public", - "foreign_column_name": "id", - "constraint_oid": 1, - "expected_column_count": 1, - "ordinal_position": 1, - } - result, connection = _import_with_catalog_results( - monkeypatch, _InsufficientPrivilegeError("catalog PK unavailable"), [foreign_key] - ) - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - assert orders.properties[0].primaryKey is True - assert orders.properties[1].relationships[0].to == "customers.id" - assert connection.rollbacks == 1 - - -def test_import_postgres_omits_foreign_keys_after_catalog_failure(monkeypatch): - catalog_primary_key = { - "table_name": "orders", - "column_name": "id", - "constraint_oid": 1, - "expected_column_count": 1, - "ordinal_position": 1, - } - result, connection = _import_with_catalog_results( - monkeypatch, [catalog_primary_key], _InsufficientPrivilegeError("catalog FK unavailable") - ) - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - assert orders.properties[0].primaryKey is True - assert orders.properties[1].relationships is None - assert connection.rollbacks == 1 - - -def test_import_postgres_keeps_composite_catalog_keys_atomic_and_ordered(monkeypatch): - from datacontract.imports import postgres_importer - - results = _catalog_results( - [ - { - "table_name": "orders", - "column_name": "customer_id", - "constraint_oid": 1, - "expected_column_count": 2, - "ordinal_position": 2, - }, - { - "table_name": "orders", - "column_name": "id", - "constraint_oid": 1, - "expected_column_count": 2, - "ordinal_position": 1, - }, - ], - [ - { - "table_name": "orders", - "column_name": "customer_id", - "foreign_table_name": "customers", - "foreign_table_schema": "public", - "foreign_column_name": "id", - "constraint_oid": 2, - "expected_column_count": 2, - "ordinal_position": 2, - }, - { - "table_name": "orders", - "column_name": "id", - "foreign_table_name": "customers", - "foreign_table_schema": "public", - "foreign_column_name": "id", - "constraint_oid": 2, - "expected_column_count": 2, - "ordinal_position": 1, - }, - ], - ) - results[postgres_importer._COLUMNS_QUERY].append( - { - "table_name": "customers", - "column_name": "customer_id", - "data_type": "integer", - "character_maximum_length": None, - "numeric_precision": 32, - "numeric_scale": 0, - "is_nullable": "NO", - "remarks": None, - } - ) - results[postgres_importer._FOREIGN_KEYS_QUERY][0]["foreign_column_name"] = "customer_id" - results[postgres_importer._FOREIGN_KEYS_QUERY][1]["foreign_column_name"] = "id" - connection = _CatalogConnection(results) - monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) - - result = import_postgres_from_connector(host="localhost", database="postgres") - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - assert [property.primaryKeyPosition for property in orders.properties if property.primaryKey] == [1, 2] - assert orders.properties[0].relationships is None - assert orders.properties[1].relationships is None - assert len(orders.relationships) == 1 - assert orders.relationships[0].from_ == ["orders.id", "orders.customer_id"] - assert orders.relationships[0].to == ["customers.id", "customers.customer_id"] - assert DataContract(data_contract_str=result.to_yaml()).lint().result == ResultEnum.passed - - -def test_import_postgres_omits_incomplete_catalog_constraints(monkeypatch): - incomplete_primary_key = { - "table_name": "orders", - "column_name": "id", - "constraint_oid": 1, - "expected_column_count": 2, - "ordinal_position": 1, - } - incomplete_foreign_key = { - "table_name": "orders", - "column_name": "customer_id", - "foreign_table_name": "customers", - "foreign_table_schema": "public", - "foreign_column_name": "id", - "constraint_oid": 2, - "expected_column_count": 2, - "ordinal_position": 1, - } - - result, _ = _import_with_catalog_results(monkeypatch, [incomplete_primary_key], [incomplete_foreign_key]) - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - assert all(property.primaryKey is None for property in orders.properties) - assert orders.properties[1].relationships is None - assert orders.relationships is None - - -def test_import_postgres_keeps_distinct_foreign_keys_on_the_same_source_property(monkeypatch): - foreign_keys = [ - { - "table_name": "orders", - "column_name": "customer_id", - "foreign_table_name": "customers", - "foreign_table_schema": "public", - "foreign_column_name": "id", - "constraint_oid": constraint_oid, - "expected_column_count": 1, - "ordinal_position": 1, - } - for constraint_oid in (1, 2) - ] - - result, _ = _import_with_catalog_results(monkeypatch, [], foreign_keys) - - orders = next(schema for schema in result.schema_ if schema.name == "orders") - customer_id = next(property for property in orders.properties if property.name == "customer_id") - assert [relationship.to for relationship in customer_id.relationships] == ["customers.id", "customers.id"] - - -def test_import_postgres_raises_when_a_required_catalog_query_fails(monkeypatch): - from datacontract.imports import postgres_importer - - connection = _CatalogConnection(_catalog_results([], [])) - connection.results[postgres_importer._TABLES_QUERY] = RuntimeError("tables unavailable") - monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) - - with pytest.raises(DataContractException, match="Could not read the Postgres catalog"): - import_postgres_from_connector(host="localhost", database="postgres") - - assert connection.rollbacks == 0 - - -def test_import_postgres_raises_for_an_optional_catalog_failure_that_is_not_an_access_error(monkeypatch): - connection = _CatalogConnection(_catalog_results(RuntimeError("catalog syntax error"), [])) - monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) - - with pytest.raises(DataContractException, match="Could not read the Postgres catalog") as exc_info: - import_postgres_from_connector(host="localhost", database="postgres") - - assert isinstance(exc_info.value.original_exception, RuntimeError) - assert connection.rollbacks == 0 - - -def test_import_postgres_raises_when_a_required_catalog_row_cannot_be_processed(monkeypatch): - failure = RuntimeError("table rows unavailable") - connection = _CatalogConnection(_catalog_results([], [])) - from datacontract.imports import postgres_importer - - connection.results[postgres_importer._TABLES_QUERY] = _FetchFailure(failure) - monkeypatch.setattr("datacontract.imports.postgres_importer.postgres_connection", lambda **kwargs: connection) - - with pytest.raises(DataContractException, match="Could not read the Postgres catalog") as exc_info: - import_postgres_from_connector(host="localhost", database="postgres") - - assert exc_info.value.original_exception is failure - assert connection.rollbacks == 0 - - def test_import_postgres(): result = _import(schema="public", postgres_table=["orders"]) @@ -430,45 +117,62 @@ def test_import_postgres_imports_foreign_key_relationships(): assert "from" not in customer_id_yaml["relationships"][0] -def test_import_postgres_recovers_keys_for_a_select_only_role(monkeypatch): - import psycopg +def test_import_postgres_omits_target_partition_foreign_key_clones(): + result = _import( + schema="partitioned_foreign_keys", + postgres_table=[ + "partitioned_source", + "partitioned_target", + "partitioned_target_low", + "partitioned_target_high", + ], + ) + + source = next(schema for schema in result.schema_ if schema.name == "partitioned_source") + assert [relationship.to for property in source.properties for relationship in property.relationships or []] == [ + "partitioned_target.target_id" + ] + +def test_import_postgres_keeps_partitioned_source_foreign_key_clones(): + result = _import( + schema="partitioned_foreign_keys", + postgres_table=[ + "partitioned_source_parent", + "partitioned_source_low", + "partitioned_source_high", + "partitioned_target", + "partitioned_target_low", + "partitioned_target_high", + ], + ) + + relationships = { + schema.name: [relationship.to for property in schema.properties for relationship in property.relationships or []] + for schema in result.schema_ + if schema.name.startswith("partitioned_source") + } + + assert relationships == { + "partitioned_source_parent": ["partitioned_target.target_id"], + "partitioned_source_low": ["partitioned_target.target_id"], + "partitioned_source_high": ["partitioned_target.target_id"], + } + + +def test_import_postgres_falls_back_to_information_schema_keys_when_catalog_is_unavailable(monkeypatch): + monkeypatch.setattr(postgres_importer, "_CATALOG_PRIMARY_KEYS_QUERY", "SELECT 1 / 0") + + result = _import(schema="select_only_keys", postgres_table=["simple_target"]) + + primary_key = [property for property in result.schema_[0].properties if property.primaryKey] + assert [(property.name, property.primaryKeyPosition) for property in primary_key] == [("target_id", 1)] + + +def test_import_postgres_recovers_keys_for_a_select_only_role(monkeypatch): reader = "select_only_key_reader" password = "select-only-key-reader-password" schema = "select_only_keys" - host = postgres.get_container_host_ip() - port = postgres.get_exposed_port(5432) - - with psycopg.connect( - dbname=postgres.dbname, user=reader, password=password, host=host, port=port, autocommit=True - ) as connection: - with connection.cursor() as cursor: - cursor.execute( - """ - WITH application_tables AS ( - SELECT relation.oid - FROM pg_catalog.pg_class AS relation - JOIN pg_catalog.pg_namespace AS namespace ON namespace.oid = relation.relnamespace - WHERE namespace.nspname = %s - AND relation.relkind = 'r' - ) - SELECT bool_and(has_schema_privilege(current_user, %s, 'USAGE')), - bool_and(has_table_privilege(current_user, oid, 'SELECT')), - NOT bool_or(has_table_privilege(current_user, oid, 'INSERT')), - NOT bool_or(has_table_privilege(current_user, oid, 'UPDATE')), - NOT bool_or(has_table_privilege(current_user, oid, 'DELETE')), - NOT bool_or(has_table_privilege(current_user, oid, 'REFERENCES')), - bool_and(has_table_privilege(current_user, 'pg_catalog.pg_constraint', 'SELECT')), - bool_and(has_table_privilege(current_user, 'pg_catalog.pg_class', 'SELECT')), - bool_and(has_table_privilege(current_user, 'pg_catalog.pg_namespace', 'SELECT')), - bool_and(has_table_privilege(current_user, 'pg_catalog.pg_attribute', 'SELECT')) - FROM application_tables - """, - (schema, schema), - ) - privileges = cursor.fetchone() - - assert privileges == (True, True, True, True, True, True, True, True, True, True) monkeypatch.setenv("DATACONTRACT_POSTGRES_USERNAME", reader) monkeypatch.setenv("DATACONTRACT_POSTGRES_PASSWORD", password) @@ -510,7 +214,8 @@ def test_import_postgres_recovers_keys_for_a_select_only_role(monkeypatch): assert DataContract(data_contract_str=result.to_yaml()).lint().result == ResultEnum.passed -def test_import_postgres_omits_relationships_outside_selected_tables_and_other_schemas(): +def test_import_postgres_warns_about_relationships_outside_selected_tables_and_other_schemas(caplog): + caplog.set_level(logging.WARNING, logger="datacontract.imports.postgres_importer") source_only = _import(schema="key_visibility", postgres_table=["selected_source", "composite_source"]) selected_source = next(schema for schema in source_only.schema_ if schema.name == "selected_source") composite_source = next(schema for schema in source_only.schema_ if schema.name == "composite_source") @@ -546,6 +251,18 @@ def test_import_postgres_omits_relationships_outside_selected_tables_and_other_s is None ) + warnings = [record.getMessage() for record in caplog.records] + assert len(warnings) == 3 + assert sum("key_visibility.selected_source" in warning for warning in warnings) == 1 + assert sum("key_visibility.composite_source" in warning for warning in warnings) == 1 + assert sum("key_visibility.cross_schema_source" in warning for warning in warnings) == 1 + assert any("key_visibility_other.cross_schema_target" in warning for warning in warnings) + + caplog.clear() + _import(schema="key_visibility", postgres_table=["selected_target"]) + + assert caplog.records == [] + def test_import_postgres_isolates_same_named_objects_in_other_schemas(): result = _import(schema="key_visibility", postgres_table=["same_named"]) @@ -574,35 +291,8 @@ def test_import_postgres_omits_catalog_constraints_with_partially_visible_column def test_import_postgres_catalog_primary_keys_override_partial_column_references(monkeypatch): - import psycopg - reader = "column_reference_key_reader" password = "column-reference-key-reader-password" - host = postgres.get_container_host_ip() - port = postgres.get_exposed_port(5432) - - with psycopg.connect( - dbname=postgres.dbname, user=reader, password=password, host=host, port=port, autocommit=True - ) as connection: - with connection.cursor() as cursor: - cursor.execute( - """ - SELECT has_column_privilege( - current_user, - 'key_visibility.column_reference_primary_key', - 'first_id', - 'REFERENCES' - ), - has_column_privilege( - current_user, - 'key_visibility.column_reference_primary_key', - 'second_id', - 'REFERENCES' - ) - """ - ) - assert cursor.fetchone() == (True, False) - monkeypatch.setenv("DATACONTRACT_POSTGRES_USERNAME", reader) monkeypatch.setenv("DATACONTRACT_POSTGRES_PASSWORD", password) result = _import(schema="key_visibility", postgres_table=["column_reference_primary_key"]) From c23bd43b55a33c1627e06834c5d19f6e78c81a6d Mon Sep 17 00:00:00 2001 From: vtulus <15098903-vtulus@users.noreply.gitlab.com> Date: Fri, 9 Oct 2026 13:17:58 +0200 Subject: [PATCH 3/5] Format Postgres importer changes --- datacontract/imports/postgres_importer.py | 3 +-- tests/test_import_postgres.py | 4 +++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/datacontract/imports/postgres_importer.py b/datacontract/imports/postgres_importer.py index d2fb2bf80..173c33557 100644 --- a/datacontract/imports/postgres_importer.py +++ b/datacontract/imports/postgres_importer.py @@ -258,8 +258,7 @@ def import_postgres_from_connector( selected_foreign_key_rows = [ row for row in foreign_key_rows_from_selected_tables - if row["foreign_table_schema"] == schema - and row["foreign_table_name"] in selected_table_names + if row["foreign_table_schema"] == schema and row["foreign_table_name"] in selected_table_names ] if foreign_key_rows is not None: selected_foreign_key_rows = _visible_constraint_rows( diff --git a/tests/test_import_postgres.py b/tests/test_import_postgres.py index a0b039d13..a11f16539 100644 --- a/tests/test_import_postgres.py +++ b/tests/test_import_postgres.py @@ -148,7 +148,9 @@ def test_import_postgres_keeps_partitioned_source_foreign_key_clones(): ) relationships = { - schema.name: [relationship.to for property in schema.properties for relationship in property.relationships or []] + schema.name: [ + relationship.to for property in schema.properties for relationship in property.relationships or [] + ] for schema in result.schema_ if schema.name.startswith("partitioned_source") } From 0fbdd0985afd3973f21dd5abed67a237b175410f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakob=20Sch=C3=B6dl?= Date: Mon, 12 Oct 2026 09:54:14 +0200 Subject: [PATCH 4/5] Say when an omitted foreign key points to another schema The import warning for a foreign key whose target table is in another schema now says "because the target table is in another schema" instead of "is not included in the imported contract", which suggested --table could include it. The module docstring now mentions that foreign keys come from pg_constraint. --- datacontract/imports/postgres_importer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datacontract/imports/postgres_importer.py b/datacontract/imports/postgres_importer.py index 173c33557..2a12e7ad3 100644 --- a/datacontract/imports/postgres_importer.py +++ b/datacontract/imports/postgres_importer.py @@ -248,11 +248,12 @@ def import_postgres_from_connector( omitted_foreign_keys.setdefault(row["constraint_oid"], row) for row in omitted_foreign_keys.values(): logger.warning( - "Omitting foreign key from %s.%s to %s.%s because the target table is not included in the imported contract.", + "Omitting foreign key from %s.%s to %s.%s because the target table is %s.", schema, row["table_name"], row["foreign_table_schema"], row["foreign_table_name"], + "in another schema" if row["foreign_table_schema"] != schema else "not included in the imported contract", ) selected_foreign_key_rows = [ From 406d041e273c260430f3790c3f34e3a2b1e118e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakob=20Sch=C3=B6dl?= Date: Mon, 12 Oct 2026 09:54:14 +0200 Subject: [PATCH 5/5] Add CHANGELOG entry --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b6861925..9e402681f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `datacontract lint --all-errors`: unknown fields are warnings instead of errors, unless a custom `--schema` rejects them ### Added -- `datacontract import postgres` imports primary and foreign keys for SELECT-only roles with catalog access. +- `datacontract import postgres` imports primary and foreign keys from the Postgres catalog, also for SELECT-only roles (#1709 @vtulus) - `datacontract edit`: enable the editor's AI assistant via `DATACONTRACT_EDITOR_AI_*` environment variables (endpoint, API key, model, provider, auth header) - `datacontract test`: check constraints and quality rules on nested properties on Athena and Trino servers