diff --git a/Makefile b/Makefile index 4915eb42be47..42b38c1ace63 100644 --- a/Makefile +++ b/Makefile @@ -263,8 +263,8 @@ build-ingestion-base-local: ## Builds the ingestion DEV docker operator with th $(MAKE) install_dev generate docker build -f ingestion/operators/docker/Dockerfile.ci . -t openmetadata/ingestion-base:local -.PHONY: build-ingestion-base-slim-local -build-ingestion-base-local: ## Builds the ingestion DEV docker operator with the local ingestion files +.PHONY: build-ingestion-slim-local +build-ingestion-slim-local: ## Builds the SLIM ingestion DEV docker operator with the local ingestion files $(MAKE) install_dev generate docker build -f ingestion/operators/docker/Dockerfile.ci . -t openmetadata/ingestion-base-slim:local --build-arg INGESTION_DEPENDENCY=slim diff --git a/ingestion/src/metadata/data_quality/validations/models.py b/ingestion/src/metadata/data_quality/validations/models.py index 3f0bb0c869af..469d93712d8b 100644 --- a/ingestion/src/metadata/data_quality/validations/models.py +++ b/ingestion/src/metadata/data_quality/validations/models.py @@ -3,7 +3,9 @@ from typing import List, Optional, Union # noqa: UP035 from pydantic import BaseModel, Field +from sqlalchemy.engine import make_url +from metadata.data_quality.validations.utils import render_url_for_data_diff from metadata.generated.schema.entity.data.table import ( Column, Table, @@ -28,6 +30,17 @@ class TableParameter(BaseModel): key_columns: Optional[list[str]] = None # noqa: UP045 extra_columns: Optional[list[str]] = None # noqa: UP045 + @property + def data_diff_service_url(self) -> Union[str, dict]: # noqa: UP007 + """`serviceUrl` rendered for data-diff's own URI parser. + + `serviceUrl` is a canonical SQLAlchemy URL, which encodes more than data-diff decodes. + Connection dicts are passed through: data-diff reads their values verbatim. + """ + if isinstance(self.serviceUrl, dict): + return self.serviceUrl + return render_url_for_data_diff(make_url(self.serviceUrl)) + class TableDiffRuntimeParameters(BaseModel): table1: TableParameter diff --git a/ingestion/src/metadata/data_quality/validations/table/sqlalchemy/tableDiff.py b/ingestion/src/metadata/data_quality/validations/table/sqlalchemy/tableDiff.py index 09b224655632..146d638170e0 100644 --- a/ingestion/src/metadata/data_quality/validations/table/sqlalchemy/tableDiff.py +++ b/ingestion/src/metadata/data_quality/validations/table/sqlalchemy/tableDiff.py @@ -13,10 +13,11 @@ import random import string import traceback +from contextlib import contextmanager from decimal import Decimal from functools import reduce from itertools import islice -from typing import Dict, Iterable, List, Optional, Tuple, cast # noqa: UP035 +from typing import Dict, Iterable, Iterator, List, Optional, Tuple, cast # noqa: UP035 from urllib.parse import urlparse import data_diff @@ -78,6 +79,8 @@ Dialects.UnityCatalog, ] +DUPLICATE_KEY_MESSAGE = "Duplicate primary keys" + class SchemaDiffResult(BaseModel): class Config: @@ -164,6 +167,29 @@ def __init__(self, param: str, dialect: str): super().__init__(f"Unsupported dialect in param {param}: {dialect}") +class DuplicateKeyError(Exception): + """A diff key column is not unique, which makes a row-level diff undefined. + + The message names the key columns only. The duplicated values themselves are row data, which + this test result is not the place to publish, and finding them costs a scan of the table. + """ + + def __init__( + self, + key_columns: List[str], # noqa: UP006 + table: Optional[str] = None, # noqa: UP045 + ): + if len(key_columns) == 1: + subject = f"Key column '{key_columns[0]}' is" + else: + subject = f"Key columns ({', '.join(repr(c) for c in key_columns)}) are" + location = f"in {table}" if table else "in one of the compared tables" + super().__init__( + f"{subject} not unique {location}. A row-level diff needs a unique key: " + "pick a unique column, or add more columns to the key." + ) + + def masked(s: str, mask: bool = True) -> str: """Mask a string if masked is True otherwise return the string. Only for development purposes, do not use in production. @@ -204,6 +230,14 @@ def _run_validation(self): try: self._validate_dialects() return self._run() + except DuplicateKeyError as e: + logger.error(f"[Data Diff]: {e}") + result = TestCaseResult( + timestamp=self.execution_date, # type: ignore + testCaseStatus=TestCaseStatus.Aborted, + result=str(e), + ) + return result # noqa: RET504 except DataDiffMismatchingKeyTypesError as e: result = TestCaseResult( timestamp=self.execution_date, # type: ignore @@ -260,36 +294,37 @@ def _run(self) -> TestCaseResult: self.runtime_params.table2.extra_columns = common_columns table_diff_iter = self.get_table_diff() - if not threshold or self.test_case.computePassedFailedRowCount: - stats = table_diff_iter.get_stats_dict() - if stats["total"] > 0: - logger.debug("Sample of failed rows:") - # depending on the data, this require scanning a lot of data - # so we only log the sample in debug mode. data can be sensitive - # so it is masked by default - for s in islice( - self.safe_table_diff_iterator(), - 10 if logger.level <= logging.DEBUG else 0, - ): - logger.debug("%s", str([s[0]] + [masked(st) for st in s[1]])) - test_case_result = self.get_row_diff_test_case_result( + with self._duplicate_keys_named(table_diff_iter): + if not threshold or self.test_case.computePassedFailedRowCount: + stats = table_diff_iter.get_stats_dict() + if stats["total"] > 0: + logger.debug("Sample of failed rows:") + # depending on the data, this require scanning a lot of data + # so we only log the sample in debug mode. data can be sensitive + # so it is masked by default + for s in islice( + self.safe_table_diff_iterator(), + 10 if logger.level <= logging.DEBUG else 0, + ): + logger.debug("%s", str([s[0]] + [masked(st) for st in s[1]])) + test_case_result = self.get_row_diff_test_case_result( + threshold, + stats["total"], + stats["updated"], + stats["exclusive_A"], + stats["exclusive_B"], + column_diff, + ) + count = self._compute_row_count(self.runner, None) # type: ignore + test_case_result.passedRows = stats["unchanged"] + test_case_result.passedRowsPercentage = test_case_result.passedRows / count * 100 + test_case_result.failedRowsPercentage = test_case_result.failedRows / count * 100 + return test_case_result + return self.get_row_diff_test_case_result( threshold, - stats["total"], - stats["updated"], - stats["exclusive_A"], - stats["exclusive_B"], + self.calculate_diffs_with_limit(table_diff_iter, threshold), column_diff, ) - count = self._compute_row_count(self.runner, None) # type: ignore - test_case_result.passedRows = stats["unchanged"] - test_case_result.passedRowsPercentage = test_case_result.passedRows / count * 100 - test_case_result.failedRowsPercentage = test_case_result.failedRows / count * 100 - return test_case_result - return self.get_row_diff_test_case_result( - threshold, - self.calculate_diffs_with_limit(table_diff_iter, threshold), - column_diff, - ) def get_incomparable_columns(self) -> List[str]: # noqa: UP006 """Get the columns that have types that are not comparable between the two tables. For example @@ -300,7 +335,7 @@ def get_incomparable_columns(self) -> List[str]: # noqa: UP006 """ table1 = data_diff.connect_to_table( - self.runtime_params.table1.serviceUrl, + self.runtime_params.table1.data_diff_service_url, self.runtime_params.table1.path, self.runtime_params.table1.key_columns, extra_columns=self.runtime_params.extraColumns, @@ -313,7 +348,7 @@ def get_incomparable_columns(self) -> List[str]: # noqa: UP006 else None, ).with_schema() table2 = data_diff.connect_to_table( - self.runtime_params.table2.serviceUrl, + self.runtime_params.table2.data_diff_service_url, self.runtime_params.table2.path, self.runtime_params.table2.key_columns, extra_columns=self.runtime_params.extraColumns, @@ -377,7 +412,7 @@ def get_table_diff(self) -> DiffResultWrapper: """Calls data_diff.diff_tables with the parameters from the test case.""" left_where, right_where = self.sample_where_clause() table1 = data_diff.connect_to_table( - self.runtime_params.table1.serviceUrl, + self.runtime_params.table1.data_diff_service_url, self.runtime_params.table1.path, self.runtime_params.table1.key_columns, # type: ignore extra_columns=self.runtime_params.table1.extra_columns, @@ -391,7 +426,7 @@ def get_table_diff(self) -> DiffResultWrapper: else None, ) table2 = data_diff.connect_to_table( - self.runtime_params.table2.serviceUrl, + self.runtime_params.table2.data_diff_service_url, self.runtime_params.table2.path, self.runtime_params.table2.key_columns, # type: ignore extra_columns=self.runtime_params.table2.extra_columns, @@ -586,6 +621,60 @@ def _validate_dialects(self): if dialect not in SUPPORTED_DIALECTS: raise UnsupportedDialectError(name, dialect) + @contextmanager + def _duplicate_keys_named(self, table_diff_iter: DiffResultWrapper) -> Iterator[None]: + """Re-raise data-diff's duplicate-key failures with the key columns named. + + A key that is not unique makes a row-level diff undefined, and data-diff reports it in two + equally opaque ways: joindiff validates the key itself and raises + `ValueError("Duplicate primary keys")`, while hashdiff only trips over it in `_get_stats`, + which folds rows into a `{key: sign}` map and asserts a key never repeats with the same + sign - a bare `AssertionError`. Neither names the key, so we do. + + We deliberately do not check uniqueness up front: that is a COUNT/COUNT DISTINCT over both + tables on every run, far too expensive on a large table to pay for an error that only + happens when the key is misconfigured. Both paths here are reached only once the diff has + already failed, and neither queries anything. + """ + try: + yield + except ValueError as exc: + if DUPLICATE_KEY_MESSAGE not in str(exc): + raise + # joindiff does not say which of the two tables it found the duplicates in. + raise DuplicateKeyError(self._diff_key_columns()) from exc + except AssertionError as exc: + table_param = self._table_with_duplicate_keys(table_diff_iter) + if table_param is None: + # Some other invariant inside data-diff; nothing useful to add. + raise + raise DuplicateKeyError( + list(table_param.key_columns or []), + table_param.fullyQualifiedName or table_param.path, + ) from exc + + def _diff_key_columns(self) -> List[str]: # noqa: UP006 + """The key columns the diff runs on. Both sides diff on the same key.""" + return list(self.runtime_params.table1.key_columns or self.runtime_params.table2.key_columns or []) + + def _table_with_duplicate_keys(self, table_diff_iter: DiffResultWrapper) -> Optional[TableParameter]: # noqa: UP045 + """The table whose rows repeat a key, or None if the diffed rows show no duplicate. + + Read off the rows `_get_stats` had already materialised into `result_list` when it failed, + so this costs no query. A row is `(sign, values)` with the key first; "-" rows come from + table1 and "+" rows from table2. + """ + key_length = len(self._diff_key_columns()) + if not key_length: + return None + seen = set() + for sign, values in table_diff_iter.result_list: + marker = (sign, tuple(values[:key_length])) + if marker in seen: + return self.runtime_params.table1 if sign == "-" else self.runtime_params.table2 + seen.add(marker) + return None + def get_column_diff(self) -> Optional[ColumnDiffResult]: # noqa: UP045 """Get the column diff between the two tables. If there are no differences, return None.""" removed, added = self.get_changed_added_columns( diff --git a/ingestion/src/metadata/data_quality/validations/utils.py b/ingestion/src/metadata/data_quality/validations/utils.py index 93cfa69e4cf7..4680086b6eb8 100644 --- a/ingestion/src/metadata/data_quality/validations/utils.py +++ b/ingestion/src/metadata/data_quality/validations/utils.py @@ -3,12 +3,22 @@ """ from typing import Any, Callable, List, Optional, TypeVar, Union # noqa: UP035 +from urllib.parse import quote + +from sqlalchemy.engine import URL from metadata.generated.schema.tests.testCase import TestCaseParameterValue +from metadata.utils.logger import test_suite_logger + +logger = test_suite_logger() T = TypeVar("T", bound=Callable) R = TypeVar("R") +# Characters that terminate the userinfo (or the whole authority) while data-diff parses the URI. +# They must stay percent-encoded even though data-diff will not decode them back. +USERNAME_RESERVED_CHARACTERS = ":/?#" + def get_test_case_param_value( test_case_param_vals: List[TestCaseParameterValue], # noqa: UP006 @@ -54,6 +64,44 @@ def get_bool_test_case_param( return str_val.lower() == "true" +def _encode_username_for_data_diff(username: str) -> str: + """Percent-encode only what data-diff's URI parser needs to locate the userinfo boundaries.""" + return "".join(f"%{ord(char):02X}" if char in USERNAME_RESERVED_CHARACTERS else char for char in username) + + +def render_url_for_data_diff(url: URL) -> str: + """Render `url` so that data-diff reads back the values it was built from. + + `URL.render_as_string` percent-encodes the username, but data-diff only decodes the password + (`CustomParseResult`), the host and the query string when it parses a URI. An encoded username + therefore reaches the driver still encoded, and `user@corp.com` tries to authenticate as + `user%40corp.com`. We hand data-diff the decoded username and keep the password encoded, so that + every component survives exactly one encode/decode round trip. + """ + if url.username is None: + return url.render_as_string(hide_password=False) + + userinfo = _encode_username_for_data_diff(url.username) + if userinfo != url.username: + logger.warning( + "[Data Diff]: The username contains characters reserved by the connection URI (%s). " + "data-diff does not decode them, so authentication may fail.", + ", ".join(sorted(set(url.username) & set(USERNAME_RESERVED_CHARACTERS))), + ) + if url.password is not None: + userinfo += f":{quote(str(url.password), safe=' +')}" + + authority = URL.create( + drivername=url.drivername, + host=url.host, + port=url.port, + database=url.database, + query=url.query, + ).render_as_string(hide_password=False) + scheme, _, rest = authority.partition("://") + return f"{scheme}://{userinfo}@{rest}" + + def casefold_if_string(value: Any) -> Any: """Case fold the value if it is a string. diff --git a/ingestion/tests/unit/observability/data_quality/validations/table/sqlalchemy/test_table_diff.py b/ingestion/tests/unit/observability/data_quality/validations/table/sqlalchemy/test_table_diff.py index 947285093a07..7ab04fa61e65 100644 --- a/ingestion/tests/unit/observability/data_quality/validations/table/sqlalchemy/test_table_diff.py +++ b/ingestion/tests/unit/observability/data_quality/validations/table/sqlalchemy/test_table_diff.py @@ -2,7 +2,10 @@ from typing import Generator # noqa: UP035 from unittest.mock import MagicMock, Mock, patch +import dsnparse import pytest +from data_diff.databases._connect import CustomParseResult +from data_diff.diff_tables import DiffResultWrapper from dirty_equals import Contains, DirtyEquals, HasAttributes, IsList from metadata.data_quality.validations.models import ( @@ -10,6 +13,7 @@ TableParameter, ) from metadata.data_quality.validations.table.sqlalchemy.tableDiff import ( + DuplicateKeyError, TableDiffValidator, ) from metadata.generated.schema.entity.data.table import ( @@ -187,3 +191,213 @@ def test_it_returns_the_expected_result( expected: DirtyEquals, ) -> None: assert validator.get_column_diff() == expected + + +ENCODED_SERVICE_URL = "postgresql://svc_user%40corp.com:p%40ssw0rd@service{n}:5432/my_db" + + +@pytest.fixture +def encoded_credentials_validator() -> Generator[tuple[TableDiffValidator, Mock], None, None]: + """A validator whose two service urls carry percent-encoded credentials.""" + runtime_params = TableDiffRuntimeParameters( + table1=build_table_parameter( + build_column("id", constraint=Constraint.PRIMARY_KEY), + key_columns=["id"], + extra_columns=[], + service_url=ENCODED_SERVICE_URL.format(n=1), + ), + table2=build_table_parameter( + build_column("id", constraint=Constraint.PRIMARY_KEY), + key_columns=["id"], + extra_columns=[], + service_url=ENCODED_SERVICE_URL.format(n=2), + ), + table_profile_config=None, + whereClause=None, + keyColumns=["id"], + extraColumns=[], + ) + with patch("metadata.data_quality.validations.table.sqlalchemy.tableDiff.data_diff") as data_diff: + mock_table = MagicMock() + mock_table.key_columns = [] + mock_table.extra_columns = [] + data_diff.connect_to_table = Mock(return_value=mock_table) + + validator = TableDiffValidator( + runner=[], + test_case=TestCase.model_construct(parameterValues=[]), + execution_date=Timestamp(root=int(datetime.datetime.now().timestamp())), + ) + validator.runtime_params = runtime_params + yield validator, data_diff.connect_to_table + + +class TestServiceUrlHandedToDataDiff: + """data-diff never decodes the username, so it must not receive a percent-encoded one. + + See https://github.com/open-metadata/OpenMetadata/issues/31124. + """ + + @staticmethod + def assert_credentials_are_decoded(connect_to_table: Mock) -> None: + urls = [call.args[0] for call in connect_to_table.call_args_list] + assert len(urls) == 2 + + for url, host in zip(urls, ["service1", "service2"], strict=True): + parsed = dsnparse.parse(url, parse_class=CustomParseResult) + assert parsed.username == "svc_user@corp.com" + # the password stays encoded on the wire: data-diff decodes that one itself + assert parsed.password == "p@ssw0rd" + assert parsed.host == host + assert parsed.paths == ["my_db"] + + def test_get_table_diff_passes_a_decoded_username( + self, encoded_credentials_validator: tuple[TableDiffValidator, Mock] + ) -> None: + validator, connect_to_table = encoded_credentials_validator + + validator.get_table_diff() + + self.assert_credentials_are_decoded(connect_to_table) + + def test_get_incomparable_columns_passes_a_decoded_username( + self, encoded_credentials_validator: tuple[TableDiffValidator, Mock] + ) -> None: + validator, connect_to_table = encoded_credentials_validator + + validator.get_incomparable_columns() + + self.assert_credentials_are_decoded(connect_to_table) + + def test_it_does_not_mutate_the_stored_service_url( + self, encoded_credentials_validator: tuple[TableDiffValidator, Mock] + ) -> None: + """`serviceUrl` is reported and re-parsed elsewhere, so it stays a canonical SQLAlchemy url.""" + validator, _ = encoded_credentials_validator + + validator.get_table_diff() + + assert validator.runtime_params.table1.serviceUrl == ENCODED_SERVICE_URL.format(n=1) + assert validator.runtime_params.table2.serviceUrl == ENCODED_SERVICE_URL.format(n=2) + + +class TestDuplicateKeyErrorMessage: + """The message is the whole point: it names the key columns the user has to change.""" + + def test_it_names_a_single_key_column_and_the_table(self) -> None: + error = DuplicateKeyError(["OrderRef"], "mssql.SalesDB.dbo.OrderEvents_Source") + + assert str(error) == ( + "Key column 'OrderRef' is not unique in mssql.SalesDB.dbo.OrderEvents_Source. " + "A row-level diff needs a unique key: pick a unique column, or add more columns to the key." + ) + + def test_it_pluralises_for_a_composite_key(self) -> None: + error = DuplicateKeyError(["OrderRef", "Region"], "db.schema.tbl") + + assert "Key columns ('OrderRef', 'Region') are not unique in db.schema.tbl" in str(error) + + def test_it_hedges_when_the_table_is_unknown(self) -> None: + """joindiff validates both tables at once and does not say which one failed.""" + error = DuplicateKeyError(["id"]) + + assert "Key column 'id' is not unique in one of the compared tables." in str(error) + + +def build_duplicate_key_validator() -> TableDiffValidator: + validator = TableDiffValidator( + runner=[], + test_case=TestCase.model_construct(parameterValues=[]), + execution_date=Timestamp(root=int(datetime.datetime.now().timestamp())), + ) + table1 = build_table_parameter( + build_column("id", constraint=Constraint.PRIMARY_KEY), key_columns=["id"], extra_columns=["name"] + ) + table1.fullyQualifiedName = "pg.db.schema.table1" + table2 = build_table_parameter( + build_column("id", constraint=Constraint.PRIMARY_KEY), key_columns=["id"], extra_columns=["name"] + ) + table2.fullyQualifiedName = "pg.db.schema.table2" + validator.runtime_params = TableDiffRuntimeParameters( + table1=table1, + table2=table2, + table_profile_config=None, + whereClause=None, + keyColumns=["id"], + extraColumns=["name"], + ) + return validator + + +def build_diff_result(*rows: tuple[str, tuple[str, ...]]) -> DiffResultWrapper: + """A wrapper whose rows are already materialised, as they are when `_get_stats` fails.""" + return DiffResultWrapper(diff=iter(()), info_tree=None, stats={}, result_list=list(rows)) + + +class TestDuplicateKeysNamed: + """data-diff reports a non-unique key opaquely, and only after the diff has already failed. + + See the notes in tableDiff.TableDiffValidator._duplicate_keys_named. + """ + + def test_it_names_the_key_when_joindiff_rejects_it(self) -> None: + validator = build_duplicate_key_validator() + + with pytest.raises(DuplicateKeyError) as excinfo, validator._duplicate_keys_named(build_diff_result()): + raise ValueError("Duplicate primary keys") + + assert "Key column 'id' is not unique in one of the compared tables" in str(excinfo.value) + + def test_it_leaves_an_unrelated_value_error_alone(self) -> None: + validator = build_duplicate_key_validator() + + with ( + pytest.raises(ValueError, match="Cannot apply key types") as excinfo, + validator._duplicate_keys_named(build_diff_result()), + ): + raise ValueError("Cannot apply key types") + + assert not isinstance(excinfo.value, DuplicateKeyError) + + def test_it_blames_table1_for_a_repeated_left_row(self) -> None: + """hashdiff only trips on the duplicate in `_get_stats`; the rows are already in memory.""" + validator = build_duplicate_key_validator() + diff_result = build_diff_result(("-", ("7", "alice")), ("+", ("7", "alicia")), ("-", ("7", "alice"))) + + with pytest.raises(DuplicateKeyError) as excinfo, validator._duplicate_keys_named(diff_result): + raise AssertionError + + assert "Key column 'id' is not unique in pg.db.schema.table1" in str(excinfo.value) + + def test_it_blames_table2_for_a_repeated_right_row(self) -> None: + validator = build_duplicate_key_validator() + diff_result = build_diff_result(("+", ("7", "alicia")), ("+", ("7", "alicia"))) + + with pytest.raises(DuplicateKeyError) as excinfo, validator._duplicate_keys_named(diff_result): + raise AssertionError + + assert "in pg.db.schema.table2" in str(excinfo.value) + + def test_it_re_raises_an_assertion_that_is_not_about_duplicate_keys(self) -> None: + """Other data-diff invariants assert too; we must not mislabel them.""" + validator = build_duplicate_key_validator() + diff_result = build_diff_result(("-", ("7", "alice")), ("+", ("8", "bob"))) + + with pytest.raises(AssertionError) as excinfo, validator._duplicate_keys_named(diff_result): + raise AssertionError("table1.is_bounded") + + assert not isinstance(excinfo.value, DuplicateKeyError) + + def test_the_same_key_on_both_sides_is_an_ordinary_diff(self) -> None: + """One '-' and one '+' for a key is how every changed row is reported.""" + validator = build_duplicate_key_validator() + diff_result = build_diff_result(("-", ("7", "alice")), ("+", ("7", "alicia"))) + + with pytest.raises(AssertionError), validator._duplicate_keys_named(diff_result): + raise AssertionError + + def test_it_passes_a_successful_diff_through(self) -> None: + validator = build_duplicate_key_validator() + + with validator._duplicate_keys_named(build_diff_result()): + pass diff --git a/ingestion/tests/unit/observability/data_quality/validations/test_data_diff_url.py b/ingestion/tests/unit/observability/data_quality/validations/test_data_diff_url.py new file mode 100644 index 000000000000..54519bfc7bda --- /dev/null +++ b/ingestion/tests/unit/observability/data_quality/validations/test_data_diff_url.py @@ -0,0 +1,368 @@ +# Copyright 2026 Collate +# Licensed under the Collate Community License, Version 1.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the URL handed to the data-diff package. + +`serviceUrl` is a canonical SQLAlchemy URL, which percent-encodes the username, the password and +the query string. data-diff parses a URI with `dsnparse` and only decodes the password, the host +and the query string, so a percent-encoded username reaches the driver still encoded and +`svc_user@corp.com` authenticates as `svc_user%40corp.com`. + +See https://github.com/open-metadata/OpenMetadata/issues/31124. + +Every assertion below goes through data-diff's *own* parser rather than a hand-rolled expectation, +so the tests fail if data-diff ever changes what it decodes. +""" + +import uuid + +import dsnparse +import pytest +from data_diff.databases._connect import CustomParseResult +from sqlalchemy.engine import URL, make_url + +from metadata.data_quality.validations.models import TableParameter +from metadata.data_quality.validations.utils import render_url_for_data_diff +from metadata.generated.schema.entity.data.table import ( + Column, + ColumnName, + DataType, + Table, +) +from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import ( + SnowflakeConnection, +) +from metadata.generated.schema.entity.services.databaseService import ( + DatabaseConnection, + DatabaseService, + DatabaseServiceType, +) +from metadata.ingestion.source.database.snowflake.data_diff.data_diff import ( + SnowflakeTableParameter, +) + + +def data_diff_parse(url: str) -> CustomParseResult: + """Parse a URI exactly the way `data_diff.connect_to_table` does.""" + return dsnparse.parse(url, parse_class=CustomParseResult) + + +def assert_round_trips(url: URL) -> str: + """Assert data-diff reads back every component of `url` unchanged. Returns the rendered URI.""" + rendered = render_url_for_data_diff(url) + parsed = data_diff_parse(rendered) + + assert parsed.username == url.username + assert parsed.password == (str(url.password) if url.password is not None else None) + assert parsed.host == url.host + assert parsed.port == url.port + assert "/".join(parsed.paths) == (url.database or "") + assert parsed.query == dict(url.query) + return rendered + + +class TestRenderUrlForDataDiff: + """The renderer must survive exactly one encode/decode round trip through data-diff.""" + + @pytest.mark.parametrize( + "url", + [ + pytest.param( + URL.create( + drivername="snowflake", + username="svc_user@corp.com", + password="p@ssw0rd/x", + host="my_account", + database="my_db/my_schema", + query={"account": "my_account", "warehouse": "MY WAREHOUSE", "role": "MY ROLE"}, + ), + id="snowflake-email-username-and-special-password", + ), + pytest.param( + URL.create( + drivername="snowflake", + username="svc_user@corp.com", + host="my_account", + database="my_db/my_schema", + query={"account": "my_account"}, + ), + id="username-with-at-sign-and-no-password", + ), + pytest.param( + URL.create( + drivername="mssql", + username="CORP\\svc_user", + password="p@ss:w/rd?#!", + host="sql.example.com", + port=1433, + database="my_db/my_schema", + ), + id="windows-domain-username-and-delimiter-heavy-password", + ), + pytest.param( + URL.create( + drivername="postgresql", + username="100%pure", + password="a%b", + host="pg.example.com", + port=5432, + database="my_db", + ), + id="literal-percent-in-username-and-password", + ), + pytest.param( + URL.create( + drivername="mysql", + username="my user", + password="my pass", + host="mysql.example.com", + port=3306, + database="my_schema", + ), + id="spaces-in-credentials", + ), + pytest.param( + URL.create( + drivername="postgresql", + username="plain_user", + password="plain_password", + host="pg.example.com", + port=5432, + database="my_db", + ), + id="nothing-to-encode", + ), + pytest.param( + URL.create(drivername="postgresql", host="pg.example.com", port=5432, database="my_db"), + id="no-credentials-at-all", + ), + pytest.param( + URL.create(drivername="postgresql", username="only_user", host="pg.example.com"), + id="username-only-no-password-no-database", + ), + ], + ) + def test_data_diff_reads_back_every_component(self, url: URL) -> None: + assert_round_trips(url) + + def test_it_stops_double_encoding_the_username(self) -> None: + """The regression from #31124: data-diff never decodes the username.""" + url = URL.create( + drivername="snowflake", + username="svc_user@corp.com", + password="my_password", + host="my_account", + database="my_db/my_schema", + ) + + # What the code used to send: SQLAlchemy encoded the username, data-diff kept it encoded + assert data_diff_parse(url.render_as_string(hide_password=False)).username == "svc_user%40corp.com" + + assert data_diff_parse(render_url_for_data_diff(url)).username == "svc_user@corp.com" + + def test_it_keeps_the_password_encoded(self) -> None: + """Decoding the whole URI is the tempting-but-wrong fix: data-diff decodes the password itself. + + A raw `@` or `/` in the password would break `dsnparse`'s credential regex, so the password + has to stay percent-encoded on the wire. + """ + url = URL.create( + drivername="snowflake", + username="svc_user", + password="p@ssw0rd/x", + host="my_account", + database="my_db/my_schema", + ) + + rendered = render_url_for_data_diff(url) + + assert "p%40ssw0rd%2Fx" in rendered + assert data_diff_parse(rendered).password == "p@ssw0rd/x" + + @pytest.mark.parametrize("reserved", [":", "/", "?", "#"]) + def test_it_keeps_username_delimiters_encoded_so_the_uri_still_parses(self, reserved: str) -> None: + """A raw delimiter in the username would move the userinfo/authority boundary. + + data-diff cannot decode these back, but a mangled username beats an unparseable URI. + """ + url = URL.create( + drivername="postgresql", + username=f"svc{reserved}user", + password="my_password", + host="pg.example.com", + port=5432, + database="my_db", + ) + + parsed = data_diff_parse(render_url_for_data_diff(url)) + + assert reserved not in parsed.username + assert parsed.host == "pg.example.com" + assert parsed.port == 5432 + assert parsed.password == "my_password" + + def test_it_warns_when_the_username_cannot_be_decoded(self, caplog: pytest.LogCaptureFixture) -> None: + url = URL.create(drivername="postgresql", username="svc:user", host="pg.example.com", database="my_db") + + with caplog.at_level("WARNING"): + render_url_for_data_diff(url) + + assert "reserved by the connection URI" in caplog.text + + def test_it_stays_quiet_for_an_ordinary_username(self, caplog: pytest.LogCaptureFixture) -> None: + url = URL.create(drivername="postgresql", username="svc_user@corp.com", host="pg.example.com") + + with caplog.at_level("WARNING"): + render_url_for_data_diff(url) + + assert caplog.text == "" + + +class TestTableParameterDataDiffServiceUrl: + def test_it_decodes_the_username_of_a_url(self) -> None: + table_parameter = TableParameter.model_construct( + serviceUrl="snowflake://svc_user%40corp.com:my_password@my_account/my_db/my_schema", + path="my_schema.my_table", + database_service_type=DatabaseServiceType.Snowflake, + columns=[], + privateKey=None, + passPhrase=None, + ) + + parsed = data_diff_parse(table_parameter.data_diff_service_url) + + assert parsed.username == "svc_user@corp.com" + assert parsed.password == "my_password" + + def test_it_passes_a_connection_dict_through_untouched(self) -> None: + """Connection dicts skip the URI parser entirely, so data-diff reads their values verbatim.""" + connection_dict = { + "driver": "mssql", + "host": "sql.example.com", + "port": 1433, + "user": "user@example.com", + "password": "p@ss/word", + "database": "my_db", + "schema": "my_schema", + } + table_parameter = TableParameter.model_construct( + serviceUrl=connection_dict, + path="my_schema.my_table", + database_service_type=DatabaseServiceType.AzureSQL, + columns=[], + privateKey=None, + passPhrase=None, + ) + + assert table_parameter.data_diff_service_url is connection_dict + + def test_it_leaves_the_stored_service_url_a_canonical_sqlalchemy_url(self) -> None: + """`serviceUrl` is re-parsed by `make_url` in the service-specific setters, so it must stay encoded.""" + table_parameter = TableParameter.model_construct( + serviceUrl="snowflake://svc_user%40corp.com:p%40ss@my_account/my_db/my_schema", + path="my_schema.my_table", + database_service_type=DatabaseServiceType.Snowflake, + columns=[], + privateKey=None, + passPhrase=None, + ) + + assert data_diff_parse(table_parameter.data_diff_service_url).username == "svc_user@corp.com" + + assert make_url(table_parameter.serviceUrl).username == "svc_user@corp.com" + assert make_url(table_parameter.serviceUrl).password == "p@ss" + + +class TestSnowflakeServiceUrlEndToEnd: + """The path that produced the `JWT token is invalid` failure reported in #31124.""" + + @staticmethod + def build_service(**overrides) -> DatabaseService: + connection = SnowflakeConnection( + username="svc_user@corp.com", + account="my_account", + warehouse="MY WAREHOUSE", + role="MY ROLE", + **overrides, + ) + return DatabaseService( + id=uuid.uuid4(), + name="snowflake_service", + serviceType=DatabaseServiceType.Snowflake, + connection=DatabaseConnection(config=connection), + ) + + @staticmethod + def build_table() -> Table: + return Table( + id=uuid.uuid4(), + name="my_table", + fullyQualifiedName="snowflake_service.my_db.my_schema.my_table", + columns=[Column(name=ColumnName("id"), dataType=DataType.INT)], + ) + + def get_table_parameter(self, service_url: str | None = None, **overrides) -> TableParameter: + return SnowflakeTableParameter().get( + self.build_service(**overrides), + self.build_table(), + {"id"}, + set(), + False, + service_url, + ) + + def test_password_authentication_reaches_data_diff_decoded(self) -> None: + table_parameter = self.get_table_parameter(password="p@ssw0rd/x") + + parsed = data_diff_parse(table_parameter.data_diff_service_url) + + assert parsed.username == "svc_user@corp.com" + assert parsed.password == "p@ssw0rd/x" + assert parsed.host == "my_account" + assert parsed.paths == ["my_db", "my_schema"] + assert parsed.query == { + "account": "my_account", + "warehouse": "MY WAREHOUSE", + "role": "MY ROLE", + } + + def test_private_key_authentication_reaches_data_diff_decoded(self) -> None: + """Key-pair auth signs a JWT over `ACCOUNT.USER`, so an encoded username is rejected by Snowflake. + + This is the `JWT token is invalid` failure from the issue: the private key is valid, the + username is not the one the key was registered for. + """ + table_parameter = self.get_table_parameter( + service_url="snowflake://svc_user%40corp.com:p%40ssw0rd@my_account/my_default_db", + password="p@ssw0rd", + privateKey="-----BEGIN PRIVATE KEY-----\nmy_key\n-----END PRIVATE KEY-----", + snowflakePrivatekeyPassphrase="my_passphrase", + ) + + parsed = data_diff_parse(table_parameter.data_diff_service_url) + + assert parsed.username == "svc_user@corp.com" + # the private key wins over the password, which the setter strips from the url + assert parsed.password is None + assert table_parameter.privateKey is not None + + def test_an_overridden_service_url_is_decoded_too(self) -> None: + """`serviceUrl` can be supplied as a test case parameter; it takes the same route.""" + table_parameter = self.get_table_parameter( + service_url="snowflake://other_user%40corp.com:other%40password@other_account/other_db", + password="my_password", + ) + + parsed = data_diff_parse(table_parameter.data_diff_service_url) + + assert parsed.username == "other_user@corp.com" + assert parsed.password == "other@password" + assert parsed.host == "other_account" diff --git a/scripts/datamodel_generation.py b/scripts/datamodel_generation.py index 1e3bfd69490c..aea6d0adfc75 100644 --- a/scripts/datamodel_generation.py +++ b/scripts/datamodel_generation.py @@ -16,6 +16,7 @@ import glob import os import re +import warnings from datamodel_code_generator.imports import Import # `model.pydantic` held the Pydantic v1 models and was removed in datamodel-code-generator 0.60+. @@ -40,7 +41,34 @@ f"{ingestion_path}src/metadata/generated/schema/type/basic.py", ] -args = f"--input {directory_root}openmetadata-spec/src/main/resources/json/schema --output-model-type pydantic_v2.BaseModel --use-annotated --base-class metadata.ingestion.models.custom_pydantic.BaseModel --input-file-type jsonschema --output {ingestion_path}src/metadata/generated/schema --set-default-enum-member".split(" ") +# OpenMetadata uses `format` as its own vocabulary rather than only the JSON Schema standard one: +# the UI form builder picks a widget from it (`FormBuilder.tsx` maps "queryBuilder" to +# QueryBuilderWidget, "password" to a masked input), and the rest are semantic hints for readers and +# for the Java/TS generators. datamodel-code-generator only knows the standard formats, so it warns +# and falls back to the base type -- which is exactly the behaviour we want here. +# +# Silence only the vocabulary we own, so a genuinely new or misspelled format still surfaces. +# Changing these in the schemas is NOT a safe cleanup: `format` drives UI widget selection. +KNOWN_CUSTOM_FORMATS = ( + "int64", + "json", + "queryBuilder", + "string", + "timezone", + "URI", + "url", + "utc-millisec", +) +warnings.filterwarnings( + "ignore", + message=rf"format of '(?:{'|'.join(map(re.escape, KNOWN_CUSTOM_FORMATS))})' not understood", + category=UserWarning, +) + +# `--formatters` is passed explicitly because the external formatters (black, isort) are about to +# become opt-in upstream. Naming them keeps today's output shape, which the post-processing below +# depends on: SOURCE_CONFIG_BLOCK matches black's parenthesised-annotation layout. +args = f"--input {directory_root}openmetadata-spec/src/main/resources/json/schema --output-model-type pydantic_v2.BaseModel --use-annotated --base-class metadata.ingestion.models.custom_pydantic.BaseModel --input-file-type jsonschema --output {ingestion_path}src/metadata/generated/schema --set-default-enum-member --formatters black isort".split(" ") main(args)