diff --git a/CHANGELOG.md b/CHANGELOG.md index b9f10ec86..9e402681f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +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 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 diff --git a/datacontract/imports/postgres_importer.py b/datacontract/imports/postgres_importer.py index 7dca8d838..2a12e7ad3 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,15 @@ from __future__ import annotations -from typing import Any, Dict, List, Optional +import logging +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 @@ -33,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 = """ @@ -66,17 +76,83 @@ 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, + 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, + 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 + -- 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 +""" + class PostgresImporter(Importer): def import_source(self, source: str, import_args: dict, config: "Config | None" = None) -> OpenDataContractStandard: @@ -119,9 +195,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 +229,48 @@ 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 = _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 %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 = [ + 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 foreign_key_rows is not None: + 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()) + 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 +301,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: + # 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 +330,67 @@ 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 _visible_constraint_rows( + rows: List[Dict[str, Any]], row_is_visible: Callable[[Dict[str, Any]], bool] +) -> List[Dict[str, Any]]: + """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( 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 +416,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..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 \ @@ -22,6 +22,10 @@ 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 + +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). 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..1766e1ab2 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,149 @@ 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 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 ( + 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..a11f16539 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,6 +12,7 @@ from datacontract.cli import app from datacontract.data_contract import DataContract +from datacontract.imports import postgres_importer from datacontract.model.exceptions import DataContractException from datacontract.model.run import ResultEnum @@ -73,6 +75,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 +103,209 @@ 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_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" + 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_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") + + 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 + ) + + 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"]) + + 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): + reader = "column_reference_key_reader" + password = "column-reference-key-reader-password" + 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 +327,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():