openmetadata-ingestion 1.6.9.0__py3-none-any.whl → 1.6.11.0__py3-none-any.whl

metadata/data_quality/validations/models.py

@@ -8,6 +8,7 @@ from metadata.generated.schema.entity.data.table import Column, TableProfilerCon
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
+from metadata.ingestion.models.custom_pydantic import CustomSecretStr
 
 
 class TableParameter(BaseModel):
@@ -15,6 +16,8 @@ class TableParameter(BaseModel):
     path: str
     columns: List[Column]
     database_service_type: DatabaseServiceType
+    privateKey: Optional[CustomSecretStr]
+    passPhrase: Optional[CustomSecretStr]
 
 
 class TableDiffRuntimeParameters(BaseModel):

metadata/data_quality/validations/runtime_param_setter/base_diff_params_setter.py

@@ -0,0 +1,120 @@
+"""Base class for param setter logic for table data diff"""
+
+from typing import List, Optional, Set
+from urllib.parse import urlparse
+
+from metadata.data_quality.validations.models import Column, TableParameter
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.ingestion.source.connections import get_connection
+from metadata.profiler.orm.registry import Dialects
+from metadata.utils import fqn
+from metadata.utils.collections import CaseInsensitiveList
+
+
+class BaseTableParameter:
+    """Base table parameter setter for the table diff test."""
+
+    def get(
+        self,
+        service: DatabaseService,
+        entity: Table,
+        key_columns,
+        extra_columns,
+        case_sensitive_columns,
+        service_url: Optional[str],
+    ) -> TableParameter:
+        """Getter table parameter for the table diff test.
+
+        Returns:
+            TableParameter
+        """
+        return TableParameter(
+            database_service_type=service.serviceType,
+            path=self.get_data_diff_table_path(entity.fullyQualifiedName.root),
+            serviceUrl=self.get_data_diff_url(
+                service,
+                entity.fullyQualifiedName.root,
+                override_url=service_url,
+            ),
+            columns=self.filter_relevant_columns(
+                entity.columns,
+                key_columns,
+                extra_columns,
+                case_sensitive=case_sensitive_columns,
+            ),
+            privateKey=None,
+            passPhrase=None,
+        )
+
+    @staticmethod
+    def get_data_diff_table_path(table_fqn: str) -> str:
+        """Get the data diff table path.
+
+        Args:
+            table_fqn (str): The fully qualified name of the table
+
+        Returns:
+            str
+        """
+        _, _, schema, table = fqn.split(table_fqn)
+        return fqn._build(  # pylint: disable=protected-access
+            "___SERVICE___", "__DATABASE__", schema, table
+        ).replace("___SERVICE___.__DATABASE__.", "")
+
+    @staticmethod
+    def get_data_diff_url(
+        db_service: DatabaseService, table_fqn, override_url: Optional[str] = None
+    ) -> str:
+        """Get the url for the data diff service.
+
+        Args:
+            db_service (DatabaseService): The database service entity
+            table_fqn (str): The fully qualified name of the table
+            override_url (Optional[str], optional): Override the url. Defaults to None.
+
+        Returns:
+            str: The url for the data diff service
+        """
+        source_url = (
+            str(get_connection(db_service.connection.config).url)
+            if not override_url
+            else override_url
+        )
+        url = urlparse(source_url)
+        # remove the driver name from the url because table-diff doesn't support it
+        kwargs = {"scheme": url.scheme.split("+")[0]}
+        service, database, schema, table = fqn.split(  # pylint: disable=unused-variable
+            table_fqn
+        )
+        # path needs to include the database AND schema in some of the connectors
+        if hasattr(db_service.connection.config, "supportsDatabase"):
+            kwargs["path"] = f"/{database}"
+        if kwargs["scheme"] in {Dialects.MSSQL, Dialects.Snowflake, Dialects.Trino}:
+            kwargs["path"] = f"/{database}/{schema}"
+        return url._replace(**kwargs).geturl()
+
+    @staticmethod
+    def filter_relevant_columns(
+        columns: List[Column],
+        key_columns: Set[str],
+        extra_columns: Set[str],
+        case_sensitive: bool,
+    ) -> List[Column]:
+        """Filter relevant columns.
+
+        Args:
+            columns (List[Column]): list of columns
+            key_columns (Set[str]): set of key columns
+            extra_columns (Set[str]): set of extra columns
+            case_sensitive (bool): case sensitive flag
+
+        Returns:
+            List[Column]
+        """
+        validated_columns = (
+            [*key_columns, *extra_columns]
+            if case_sensitive
+            else CaseInsensitiveList([*key_columns, *extra_columns])
+        )
+        return [c for c in columns if c.name.root in validated_columns]

metadata/data_quality/validations/runtime_param_setter/table_diff_params_setter.py

@@ -14,21 +14,32 @@ from typing import List, Optional, Set
 from urllib.parse import urlparse
 
 from metadata.data_quality.validations import utils
-from metadata.data_quality.validations.models import (
-    Column,
-    TableDiffRuntimeParameters,
-    TableParameter,
-)
+from metadata.data_quality.validations.models import Column, TableDiffRuntimeParameters
 from metadata.data_quality.validations.runtime_param_setter.param_setter import (
     RuntimeParameterSetter,
 )
 from metadata.generated.schema.entity.data.table import Constraint, Table
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.tests.testCase import TestCase
 from metadata.ingestion.source.connections import get_connection
 from metadata.profiler.orm.registry import Dialects
 from metadata.utils import fqn
 from metadata.utils.collections import CaseInsensitiveList
+from metadata.utils.importer import get_module_dir, import_from_module
+
+
+def get_for_source(
+    service_type: ServiceType, source_type: str, from_: str = "ingestion"
+):
+    return import_from_module(
+        "metadata.{}.source.{}.{}.{}.ServiceSpec".format(  # pylint: disable=C0209
+            from_,
+            service_type.name.lower(),
+            get_module_dir(source_type),
+            "service_spec",
+        )
+    )
 
 
 class TableDiffParamsSetter(RuntimeParameterSetter):
@@ -51,67 +62,70 @@ class TableDiffParamsSetter(RuntimeParameterSetter):
         }
 
     def get_parameters(self, test_case) -> TableDiffRuntimeParameters:
+        # Using the specs class method causes circular import as TestSuiteInterface
+        # imports RuntimeParameterSetter
+        cls_path = get_for_source(
+            ServiceType.Database,
+            source_type=self.service_connection_config.type.value.lower(),
+        ).data_diff
+        cls = import_from_module(cls_path)()
+
+        service1: DatabaseService = self.ometa_client.get_by_id(
+            DatabaseService, self.table_entity.service.id, nullable=False
+        )
+
         service1_url = (
             str(get_connection(self.service_connection_config).url)
             if self.service_connection_config
             else None
         )
-        service1: DatabaseService = self.ometa_client.get_by_id(
-            DatabaseService, self.table_entity.service.id, nullable=False
-        )
+
         table2_fqn = self.get_parameter(test_case, "table2")
-        case_sensitive_columns: bool = utils.get_bool_test_case_param(
-            test_case.parameterValues, "caseSensitiveColumns"
-        )
         if table2_fqn is None:
             raise ValueError("table2 not set")
         table2: Table = self.ometa_client.get_by_name(
             Table, fqn=table2_fqn, nullable=False
         )
-        service2_url = (
-            service1_url if table2.service == self.table_entity.service else None
-        )
         service2: DatabaseService = self.ometa_client.get_by_id(
             DatabaseService, table2.service.id, nullable=False
         )
+        service2_url = (
+            self.get_parameter(test_case, "service2Url") or service1_url
+            if table2.service == self.table_entity.service
+            else None
+        )
+
         key_columns = self.get_key_columns(test_case)
-        extra_columns = self.get_extra_columns(
-            key_columns, test_case, self.table_entity.columns, table2.columns
+        extra_columns = (
+            self.get_extra_columns(
+                key_columns, test_case, self.table_entity.columns, table2.columns
+            )
+            or set()
         )
+        case_sensitive_columns: bool = (
+            utils.get_bool_test_case_param(
+                test_case.parameterValues, "caseSensitiveColumns"
+            )
+            or False
+        )
+
         return TableDiffRuntimeParameters(
             table_profile_config=self.table_entity.tableProfilerConfig,
-            table1=TableParameter(
-                database_service_type=service1.serviceType,
-                path=self.get_data_diff_table_path(
-                    self.table_entity.fullyQualifiedName.root
-                ),
-                serviceUrl=self.get_data_diff_url(
-                    service1,
-                    self.table_entity.fullyQualifiedName.root,
-                    override_url=service1_url,
-                ),
-                columns=self.filter_relevant_columns(
-                    self.table_entity.columns,
-                    key_columns,
-                    extra_columns,
-                    case_sensitive=case_sensitive_columns,
-                ),
+            table1=cls.get(
+                service1,
+                self.table_entity,
+                key_columns,
+                extra_columns,
+                case_sensitive_columns,
+                service1_url,
             ),
-            table2=TableParameter(
-                database_service_type=service2.serviceType,
-                path=self.get_data_diff_table_path(table2_fqn),
-                serviceUrl=self.get_data_diff_url(
-                    service2,
-                    table2_fqn,
-                    override_url=self.get_parameter(test_case, "service2Url")
-                    or service2_url,
-                ),
-                columns=self.filter_relevant_columns(
-                    table2.columns,
-                    key_columns,
-                    extra_columns,
-                    case_sensitive=case_sensitive_columns,
-                ),
+            table2=cls.get(
+                service2,
+                table2,
+                key_columns,
+                extra_columns,
+                case_sensitive_columns,
+                service2_url,
             ),
             keyColumns=list(key_columns),
             extraColumns=list(extra_columns),

metadata/data_quality/validations/table/sqlalchemy/tableDiff.py

@@ -265,6 +265,12 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             self.runtime_params.keyColumns,
             extra_columns=self.runtime_params.extraColumns,
             case_sensitive=self.get_case_sensitive(),
+            key_content=self.runtime_params.table1.privateKey.get_secret_value()
+            if self.runtime_params.table1.privateKey
+            else None,
+            private_key_passphrase=self.runtime_params.table1.passPhrase.get_secret_value()
+            if self.runtime_params.table1.passPhrase
+            else None,
         ).with_schema()
         table2 = data_diff.connect_to_table(
             self.runtime_params.table2.serviceUrl,
@@ -272,6 +278,12 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             self.runtime_params.keyColumns,
             extra_columns=self.runtime_params.extraColumns,
             case_sensitive=self.get_case_sensitive(),
+            key_content=self.runtime_params.table2.privateKey.get_secret_value()
+            if self.runtime_params.table2.privateKey
+            else None,
+            private_key_passphrase=self.runtime_params.table2.passPhrase.get_secret_value()
+            if self.runtime_params.table2.passPhrase
+            else None,
         ).with_schema()
         result = []
         for column in table1.key_columns + table1.extra_columns:
@@ -332,6 +344,12 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             self.runtime_params.keyColumns,  # type: ignore
             case_sensitive=self.get_case_sensitive(),
             where=left_where,
+            key_content=self.runtime_params.table1.privateKey.get_secret_value()
+            if self.runtime_params.table1.privateKey
+            else None,
+            private_key_passphrase=self.runtime_params.table1.passPhrase.get_secret_value()
+            if self.runtime_params.table1.passPhrase
+            else None,
         )
         table2 = data_diff.connect_to_table(
             self.runtime_params.table2.serviceUrl,
@@ -339,6 +357,12 @@ class TableDiffValidator(BaseTestValidator, SQAValidatorMixin):
             self.runtime_params.keyColumns,  # type: ignore
             case_sensitive=self.get_case_sensitive(),
             where=right_where,
+            key_content=self.runtime_params.table1.privateKey.get_secret_value()
+            if self.runtime_params.table1.privateKey
+            else None,
+            private_key_passphrase=self.runtime_params.table1.passPhrase.get_secret_value()
+            if self.runtime_params.table1.passPhrase
+            else None,
         )
         data_diff_kwargs = {
             "key_columns": self.runtime_params.keyColumns,

metadata/generated/schema/analytics/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/analytics/basic.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/basic.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/analytics/reportDataType/aggregatedCostAnalysisReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/aggregatedCostAnalysisReportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/entityReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/entityReportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/rawCostAnalysisReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/rawCostAnalysisReportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/webAnalyticEntityViewReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/webAnalyticEntityViewReportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/reportDataType/webAnalyticUserActivityReportData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/reportDataType/webAnalyticUserActivityReportData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEvent.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventData.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventType/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/analytics/webAnalyticEventType/customEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventType/customEvent.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/analytics/webAnalyticEventType/pageViewEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  analytics/webAnalyticEventType/pageViewEvent.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/api/addGlossaryToAssetsRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/addGlossaryToAssetsRequest.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/addTagToAssetsRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/addTagToAssetsRequest.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/analytics/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/api/analytics/createWebAnalyticEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/analytics/createWebAnalyticEvent.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/automations/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/api/automations/createWorkflow.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/automations/createWorkflow.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/bulkAssets.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/bulkAssets.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/classification/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/api/classification/createClassification.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/classification/createClassification.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/classification/createTag.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/classification/createTag.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/classification/loadTags.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/classification/loadTags.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/createBot.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/createBot.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/createEventPublisherJob.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/createEventPublisherJob.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/createType.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/createType.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00

metadata/generated/schema/api/data/createAPICollection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createAPICollection.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/createAPIEndpoint.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createAPIEndpoint.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/createChart.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createChart.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/createContainer.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createContainer.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/createCustomProperty.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createCustomProperty.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/api/data/createDashboard.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  api/data/createDashboard.json
-#   timestamp: 2025-04-17T06:23:01+00:00
+#   timestamp: 2025-04-29T03:36:00+00:00
 
 from __future__ import annotations