openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

_openmetadata_testutils/dict.py

@@ -0,0 +1,23 @@
+def merge(source: dict, destination: dict):
+    """
+    Merge source dictionary into destination dictionary recursively in place.
+
+    Examples
+    # >>> a = { 'first' : { 'all_rows' : { 'pass' : 'dog', 'number' : '1' } } }
+    # >>> b = { 'first' : { 'all_rows' : { 'fail' : 'cat', 'number' : '5' } } }
+    # >>> merge(b, a) == { 'first' : { 'all_rows' : { 'pass' : 'dog', 'fail' : 'cat', 'number' : '5' } } }
+    True
+
+    Args:
+        source (dict): Source dictionary
+        destination (dict): Destination dictionary
+    """
+    for key, value in source.items():
+        if isinstance(value, dict):
+            # get node or create one
+            node = destination.setdefault(key, {})
+            merge(value, node)
+        else:
+            destination[key] = value
+
+    return destination

_openmetadata_testutils/pydantic/test_utils.py

@@ -1,10 +1,13 @@
 from collections import deque
+from typing import List, Union
 
 from pydantic import BaseModel
 
 
 def assert_equal_pydantic_objects(
-    expected: BaseModel, actual: BaseModel, ignore_none=True
+    expected: Union[BaseModel, List[BaseModel]],
+    actual: Union[BaseModel, List[BaseModel]],
+    ignore_none=True,
 ):
     """Compare 2 pydantic objects recursively and raise an AssertionError if they are not equal along with all
     the differences by field. If `ignore_none` is set to True, expected None values will be ignored. This can be
@@ -32,6 +35,10 @@ def assert_equal_pydantic_objects(
         Traceback (most recent call last):
         ```
         AssertionError: objects mismatched on field: [b.a], expected: [1], actual: [2]
+        >>> assert_equal_pydantic_objects([a1, a2], [a2, a1])
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [0].a, expected: [1], actual: [2]
 
     Args:
         expected (BaseModel): The expected pydantic object.
@@ -69,11 +76,24 @@ def assert_equal_pydantic_objects(
                     errors.append(
                         f"objects mismatched on field: [{new_key_prefix}], expected: [{expected_value}], actual: [{actual_value}]"
                     )
+        elif isinstance(expected, list):
+            if not isinstance(actual, list):
+                errors.append(
+                    f"validation error on field: [{current_key_prefix}], expected: [list], actual: [{type(actual).__name__}]"
+                )
+            elif len(expected) != len(actual):
+                errors.append(
+                    f"mismatch length at {current_key_prefix}: expected: [{len(expected)}], actual: [{len(actual)}]"
+                )
+            else:
+                for i, (expected_item, actual_item) in enumerate(zip(expected, actual)):
+                    queue.append(
+                        (expected_item, actual_item, f"{current_key_prefix}[{i}]")
+                    )
         else:
             if expected != actual:
                 errors.append(
                     f"mismatch at {current_key_prefix}: expected: [{expected}], actual: [{actual}]"
                 )
-
     if errors:
         raise AssertionError("\n".join(errors))

_openmetadata_testutils/pytest_openmetadata/plugin.py

@@ -0,0 +1,86 @@
+from typing import Type
+
+import pytest
+
+from _openmetadata_testutils.ometa import int_admin_ometa
+from metadata.generated.schema.entity.teams.user import AuthenticationMechanism, User
+from metadata.generated.schema.metadataIngestion.workflow import LogLevels
+from metadata.ingestion.api.common import Entity
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.workflow.ingestion import IngestionWorkflow
+
+
+@pytest.fixture(scope="session")
+def metadata():
+    return int_admin_ometa()
+
+
+@pytest.fixture(scope="session")
+def sink_config(metadata):
+    return {
+        "type": "metadata-rest",
+        "config": {},
+    }
+
+
+@pytest.fixture(scope="session")
+def workflow_config(metadata):
+    return {
+        "loggerLevel": LogLevels.DEBUG.value,
+        "openMetadataServerConfig": metadata.config.model_dump(),
+    }
+
+
+@pytest.fixture(scope="session")
+def ingestion_bot_wokflow_config(metadata, workflow_config):
+    ingestion_bot: User = metadata.get_by_name(
+        entity=User, fqn="ingestion-bot", nullable=False
+    )
+    ingestion_bot_auth: AuthenticationMechanism = metadata.get_by_id(
+        entity=AuthenticationMechanism, entity_id=ingestion_bot.id, nullable=False
+    )
+    workflow_config = workflow_config.copy()
+    workflow_config["openMetadataServerConfig"]["securityConfig"][
+        "jwtToken"
+    ] = ingestion_bot_auth.config.JWTToken
+    return workflow_config
+
+
+@pytest.fixture()
+def clean_up_fqn(metadata):
+    fqns = []
+
+    def inner(entity_type: type[Entity], fqn: str):
+        fqns.append((entity_type, fqn))
+
+    yield inner
+    for entity_type, fqn in fqns:
+        entity = metadata.get_by_name(entity_type, fqn, fields=["*"])
+        if entity:
+            metadata.delete(entity_type, entity.id, recursive=True, hard_delete=True)
+
+
+@pytest.fixture(scope="session")
+def ingestion_bot_workflow_config(metadata: OpenMetadata):
+    ingestion_bot: User = metadata.get_by_name(entity=User, fqn="ingestion-bot")
+    ingestion_bot_auth: AuthenticationMechanism = metadata.get_by_id(
+        entity=AuthenticationMechanism, entity_id=ingestion_bot.id
+    )
+    config = metadata.config.model_dump()
+    config["securityConfig"]["jwtToken"] = ingestion_bot_auth.config.JWTToken
+    return {
+        "loggerLevel": LogLevels.DEBUG.value,
+        "openMetadataServerConfig": config,
+    }
+
+
+@pytest.fixture(scope="module")
+def run_workflow():
+    def _run(workflow_type: Type[IngestionWorkflow], config, raise_from_status=True):
+        workflow: IngestionWorkflow = workflow_type.create(config)
+        workflow.execute()
+        if raise_from_status:
+            workflow.raise_from_status()
+        return workflow
+
+    return _run

_openmetadata_testutils/pytest_openmetadata/test_utils.py

@@ -0,0 +1,99 @@
+from collections import deque
+from typing import List, Union
+
+from pydantic import BaseModel
+
+
+def assert_equal_pydantic_objects(
+    expected: Union[BaseModel, List[BaseModel]],
+    actual: Union[BaseModel, List[BaseModel]],
+    ignore_none=True,
+):
+    """Compare 2 pydantic objects recursively and raise an AssertionError if they are not equal along with all
+    the differences by field. If `ignore_none` is set to True, expected None values will be ignored. This can be
+    useful when comparing objects that are partially filled.
+
+    Example:
+        >>> from pydantic import BaseModel
+        >>> class A(BaseModel):
+        ...     a: int
+        >>> class B(BaseModel):
+        ...     b: A
+        >>> a1 = A(a=1)
+        >>> a2 = A(a=2)
+        >>> b1 = B(b=a1)
+        >>> b2 = B(b=a2)
+        >>> assert_equal_pydantic_objects(a1, b1)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on type at : expected: [A], actual: [B]
+        >>> assert_equal_pydantic_objects(a1, a2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [a], expected: [1], actual: [2]
+        >>> assert_equal_pydantic_objects(b1, b2)
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [b.a], expected: [1], actual: [2]
+        >>> assert_equal_pydantic_objects([a1, a2], [a2, a1])
+        Traceback (most recent call last):
+        ```
+        AssertionError: objects mismatched on field: [0].a, expected: [1], actual: [2]
+
+    Args:
+        expected (BaseModel): The expected pydantic object.
+        actual (BaseModel): The actual pydantic object.
+        ignore_none (bool, optional): Whether to ignore None values. Defaults to True.
+
+    Raises:
+        AssertionError: If the objects are not equal. The error message will contain all the differences.
+    """
+    errors = []
+    queue = deque([(expected, actual, "")])
+    while queue:
+        expected, actual, current_key_prefix = queue.popleft()
+        if not isinstance(expected, actual.__class__):
+            errors.append(
+                f"objects mismatched on type at {current_key_prefix}: "
+                f"expected: [{type(expected).__name__}], actual: [{type(actual).__name__}]"
+            )
+            continue
+        if issubclass(expected.__class__, BaseModel) and isinstance(
+            expected.model_dump(), dict
+        ):
+            for key, expected_value in expected.model_dump().items():
+                if expected_value is None and ignore_none:
+                    continue
+                actual_value = actual.model_dump().get(key)
+                new_key_prefix = (
+                    f"{current_key_prefix}.{key}" if current_key_prefix else key
+                )
+                if issubclass(getattr(expected, key).__class__, BaseModel):
+                    queue.append(
+                        (getattr(expected, key), getattr(actual, key), new_key_prefix)
+                    )
+                elif expected_value != actual_value:
+                    errors.append(
+                        f"objects mismatched on field: [{new_key_prefix}], expected: [{expected_value}], actual: [{actual_value}]"
+                    )
+        elif isinstance(expected, list):
+            if not isinstance(actual, list):
+                errors.append(
+                    f"validation error on field: [{current_key_prefix}], expected: [list], actual: [{type(actual).__name__}]"
+                )
+            elif len(expected) != len(actual):
+                errors.append(
+                    f"mismatch length at {current_key_prefix}: expected: [{len(expected)}], actual: [{len(actual)}]"
+                )
+            else:
+                for i, (expected_item, actual_item) in enumerate(zip(expected, actual)):
+                    queue.append(
+                        (expected_item, actual_item, f"{current_key_prefix}[{i}]")
+                    )
+        else:
+            if expected != actual:
+                errors.append(
+                    f"mismatch at {current_key_prefix}: expected: [{expected}], actual: [{actual}]"
+                )
+    if errors:
+        raise AssertionError("\n".join(errors))

metadata/cli/app.py

@@ -32,6 +32,7 @@ def run_app(config_path: Path) -> None:
 
     try:
         config_dict = load_config_file(config_path)
+        # no logging for config because apps might have custom secrets
         workflow = ApplicationWorkflow.create(config_dict)
     except Exception as exc:
         logger.error(f"Error running the application {exc}")

metadata/cli/dataquality.py

@@ -20,7 +20,7 @@ from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     PipelineType,
 )
-from metadata.utils.logger import cli_logger
+from metadata.utils.logger import cli_logger, redacted_config
 from metadata.workflow.data_quality import TestSuiteWorkflow
 from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
@@ -37,6 +37,9 @@ def run_test(config_path: Path) -> None:
     workflow_config_dict = None
     try:
         workflow_config_dict = load_config_file(config_path)
+        logger.debug(
+            "Using workflow config:\n%s", redacted_config(workflow_config_dict)
+        )
         workflow = TestSuiteWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/cli/ingest.py

@@ -20,7 +20,7 @@ from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     PipelineType,
 )
-from metadata.utils.logger import cli_logger
+from metadata.utils.logger import cli_logger, redacted_config
 from metadata.workflow.metadata import MetadataWorkflow
 from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
@@ -37,6 +37,7 @@ def run_ingest(config_path: Path) -> None:
     config_dict = None
     try:
         config_dict = load_config_file(config_path)
+        logger.debug("Using workflow config:\n%s", redacted_config(config_dict))
         workflow = MetadataWorkflow.create(config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/cli/lineage.py

@@ -27,7 +27,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
 from metadata.generated.schema.metadataIngestion.workflow import WorkflowConfig
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.constants import UTF_8
-from metadata.utils.logger import cli_logger
+from metadata.utils.logger import cli_logger, redacted_config
 from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
@@ -52,6 +52,7 @@ def run_lineage(config_path: Path) -> None:
     config_dict = None
     try:
         config_dict = load_config_file(config_path)
+        logger.debug("Using workflow config:\n%s", redacted_config(config_dict))
         workflow = LineageWorkflow.model_validate(config_dict)
 
     except Exception as exc:

metadata/cli/profile.py

@@ -20,7 +20,7 @@ from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     PipelineType,
 )
-from metadata.utils.logger import cli_logger
+from metadata.utils.logger import cli_logger, redacted_config
 from metadata.workflow.profiler import ProfilerWorkflow
 from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
@@ -37,6 +37,9 @@ def run_profiler(config_path: Path) -> None:
     workflow_config_dict = None
     try:
         workflow_config_dict = load_config_file(config_path)
+        logger.debug(
+            "Using workflow config:\n%s", redacted_config(workflow_config_dict)
+        )
         workflow = ProfilerWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/cli/usage.py

@@ -20,7 +20,7 @@ from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     PipelineType,
 )
-from metadata.utils.logger import cli_logger
+from metadata.utils.logger import cli_logger, redacted_config
 from metadata.workflow.usage import UsageWorkflow
 from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
@@ -37,6 +37,7 @@ def run_usage(config_path: Path) -> None:
     config_dict = None
     try:
         config_dict = load_config_file(config_path)
+        logger.debug("Using workflow config:\n%s", redacted_config(config_dict))
         workflow = UsageWorkflow.create(config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())

metadata/config/common.py

@@ -19,7 +19,9 @@ from abc import ABC, abstractmethod
 from typing import IO, Any, Optional
 
 import yaml
-from pydantic import BaseModel, ConfigDict
+from pydantic import ConfigDict
+
+from metadata.ingestion.models.custom_pydantic import BaseModel
 
 
 class ConfigModel(BaseModel):

metadata/data_quality/api/models.py

@@ -18,13 +18,14 @@ multiple test cases per workflow.
 
 from typing import List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteRequest
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.tests.basic import TestCaseResult
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
+from metadata.ingestion.models.custom_pydantic import BaseModel
 
 
 class TestCaseDefinition(ConfigModel):

metadata/data_quality/builders/i_validator_builder.py

@@ -15,7 +15,7 @@ Validators are test classes (e.g. columnValuesToBeBetween, etc.)
 """
 
 from abc import ABC, abstractmethod
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Optional, Type, Union
 
 from metadata.data_quality.validations.base_test_handler import BaseTestValidator
@@ -23,6 +23,7 @@ from metadata.data_quality.validations.runtime_param_setter.param_setter import
     RuntimeParameterSetter,
 )
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
+from metadata.generated.schema.type.basic import Timestamp
 from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.importer import import_test_case_class
 
@@ -91,7 +92,9 @@ class IValidatorBuilder(ABC):
         self._validator = self.validator_cls(
             self.runner,
             test_case=self.test_case,
-            execution_date=int(datetime.now().timestamp() * 1000),
+            execution_date=Timestamp(
+                int(datetime.now(tz=timezone.utc).timestamp() * 1000)
+            ),
         )
 
     @abstractmethod

metadata/data_quality/interface/pandas/pandas_test_suite_interface.py

@@ -46,6 +46,7 @@ class PandasTestSuiteInterface(TestSuiteInterface, PandasInterfaceMixin):
         table_entity: Table = None,
         **kwargs,  # pylint: disable=unused-argument
     ):
+        super().__init__()
         self.table_entity = table_entity
 
         self.ometa_client = ometa_client

metadata/data_quality/interface/sqlalchemy/sqa_test_suite_interface.py

@@ -52,6 +52,7 @@ class SQATestSuiteInterface(SQAInterfaceMixin, TestSuiteInterface):
         table_entity: Table = None,
         sqa_metadata=None,
     ):
+        super().__init__()
         self.ometa_client = ometa_client
         self.table_entity = table_entity
         self.service_connection_config = service_connection_config

metadata/data_quality/interface/test_suite_interface.py

@@ -43,17 +43,12 @@ class TestSuiteInterface(ABC):
 
     runtime_params_setter_fact = RuntimeParameterSetterFactory
 
-    @abstractmethod
-    def __init__(
-        self,
-        ometa_client: OpenMetadata,
-        service_connection_config: DatabaseConnection,
-        table_entity: Table,
-    ):
-        """Required attribute for the interface"""
-        self.ometa_client = ometa_client
-        self.service_connection_config = service_connection_config
-        self.table_entity = table_entity
+    def __init__(self):
+        """This constructor exists to staisfy the linter. Its fields should be set in the implementation class"""
+        self.ometa_client: DatabaseConnection = None
+        self.ometa_client: OpenMetadata = None
+        self.table_entity: Table = None
+        self.service_connection_config: DatabaseConnection = None
 
     @property
     def sampler(self):

metadata/data_quality/interface/test_suite_interface_factory.py

@@ -8,27 +8,14 @@
 #  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.
+#  pylint: disable=import-outside-toplevel
 """
 Interface factory
 """
 import traceback
 from logging import Logger
+from typing import Callable, Dict, Type
 
-from metadata.data_quality.interface.pandas.pandas_test_suite_interface import (
-    PandasTestSuiteInterface,
-)
-from metadata.data_quality.interface.sqlalchemy.databricks.test_suite_interface import (
-    DatabricksTestSuiteInterface,
-)
-from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import (
-    SnowflakeTestSuiteInterface,
-)
-from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
-    SQATestSuiteInterface,
-)
-from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
-    UnityCatalogTestSuiteInterface,
-)
 from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
@@ -55,19 +42,18 @@ class TestSuiteInterfaceFactory:
 
     def __init__(self):
         """Initialize the interface factory"""
-        self._interface_type = {
-            "base": SQATestSuiteInterface,
-            DatalakeConnection.__name__: PandasTestSuiteInterface,
+        self._interface_type: Dict[str, Callable[[], Type[TestSuiteInterface]]] = {
+            "base": self.sqa,
         }
 
-    def register(self, interface_type: str, interface: TestSuiteInterface):
+    def register(self, interface_type: str, fn: Callable[[], Type[TestSuiteInterface]]):
         """Register the interface
 
         Args:
             interface_type (str): type of the interface
-            interface (TestSuiteInterface): a class that implements the TestSuiteInterface
+            interface (callable): a class that implements the TestSuiteInterface
         """
-        self._interface_type[interface_type] = interface
+        self._interface_type[interface_type] = fn
 
     def register_many(self, interface_dict):
         """
@@ -77,8 +63,8 @@ class TestSuiteInterfaceFactory:
             interface_dict: A dictionary mapping connection class names (strings) to their
             corresponding profiler interface classes.
         """
-        for interface_type, interface_class in interface_dict.items():
-            self.register(interface_type, interface_class)
+        for interface_type, interface_fn in interface_dict.items():
+            self.register(interface_type, interface_fn)
 
     def create(
         self,
@@ -104,25 +90,69 @@ class TestSuiteInterfaceFactory:
         except AttributeError as err:
             logger.debug(traceback.format_exc())
             raise AttributeError(f"Could not instantiate interface class: {err}")
-        interface = self._interface_type.get(connection_type)
+        interface_fn = self._interface_type.get(connection_type)
 
-        if not interface:
-            interface = self._interface_type["base"]
+        if not interface_fn:
+            interface_fn = self._interface_type["base"]
 
-        return interface(
+        interface_class = interface_fn()
+        return interface_class(
             service_connection_config, ometa_client, table_entity, *args, **kwargs
         )
 
+    @staticmethod
+    def sqa() -> Type[TestSuiteInterface]:
+        """Lazy load the SQATestSuiteInterface"""
+        from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
+            SQATestSuiteInterface,
+        )
 
-test_suite_interface_factory = TestSuiteInterfaceFactory()
+        return SQATestSuiteInterface
+
+    @staticmethod
+    def pandas() -> Type[TestSuiteInterface]:
+        """Lazy load the PandasTestSuiteInterface"""
+        from metadata.data_quality.interface.pandas.pandas_test_suite_interface import (
+            PandasTestSuiteInterface,
+        )
+
+        return PandasTestSuiteInterface
+
+    @staticmethod
+    def snowflake() -> Type[TestSuiteInterface]:
+        """Lazy load the SnowflakeTestSuiteInterface"""
+        from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import (
+            SnowflakeTestSuiteInterface,
+        )
+
+        return SnowflakeTestSuiteInterface
+
+    @staticmethod
+    def unity_catalog() -> Type[TestSuiteInterface]:
+        """Lazy load the UnityCatalogTestSuiteInterface"""
+        from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
+            UnityCatalogTestSuiteInterface,
+        )
+
+        return UnityCatalogTestSuiteInterface
+
+    @staticmethod
+    def databricks() -> Type[TestSuiteInterface]:
+        """Lazy load the DatabricksTestSuiteInterface"""
+        from metadata.data_quality.interface.sqlalchemy.databricks.test_suite_interface import (
+            DatabricksTestSuiteInterface,
+        )
+
+        return DatabricksTestSuiteInterface
 
 
 test_suite_interface = {
-    DatabaseConnection.__name__: SQATestSuiteInterface,
-    DatalakeConnection.__name__: PandasTestSuiteInterface,
-    SnowflakeConnection.__name__: SnowflakeTestSuiteInterface,
-    UnityCatalogConnection.__name__: UnityCatalogTestSuiteInterface,
-    DatabricksConnection.__name__: DatabricksTestSuiteInterface,
+    DatabaseConnection.__name__: TestSuiteInterfaceFactory.sqa,
+    DatalakeConnection.__name__: TestSuiteInterfaceFactory.pandas,
+    SnowflakeConnection.__name__: TestSuiteInterfaceFactory.snowflake,
+    UnityCatalogConnection.__name__: TestSuiteInterfaceFactory.unity_catalog,
+    DatabricksConnection.__name__: TestSuiteInterfaceFactory.databricks,
 }
 
+test_suite_interface_factory = TestSuiteInterfaceFactory()
 test_suite_interface_factory.register_many(test_suite_interface)

metadata/data_quality/source/test_suite.py

@@ -38,8 +38,8 @@ from metadata.ingestion.api.steps import Source
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils import fqn
 from metadata.utils.constants import CUSTOM_CONNECTOR_PREFIX
-from metadata.utils.importer import import_source_class
 from metadata.utils.logger import test_suite_logger
+from metadata.utils.service_spec.service_spec import import_source_class
 
 logger = test_suite_logger()
 

metadata/data_quality/validations/base_test_handler.py

@@ -17,7 +17,6 @@ from __future__ import annotations
 
 import reprlib
 from abc import ABC, abstractmethod
-from datetime import datetime
 from typing import TYPE_CHECKING, Callable, List, Optional, Type, TypeVar, Union
 
 from metadata.data_quality.validations import utils
@@ -30,6 +29,7 @@ from metadata.generated.schema.tests.basic import (
     TestResultValue,
 )
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
+from metadata.generated.schema.type.basic import Timestamp
 from metadata.profiler.processor.runner import QueryRunner
 
 if TYPE_CHECKING:
@@ -51,7 +51,7 @@ class BaseTestValidator(ABC):
         self,
         runner: Union[QueryRunner, List["DataFrame"]],
         test_case: TestCase,
-        execution_date: Union[datetime, float],
+        execution_date: Timestamp,
     ) -> None:
         self.runner = runner
         self.test_case = test_case
@@ -80,7 +80,7 @@ class BaseTestValidator(ABC):
 
     def get_test_case_result_object(  # pylint: disable=too-many-arguments
         self,
-        execution_date: Union[datetime, float],
+        execution_date: Timestamp,
         status: TestCaseStatus,
         result: str,
         test_result_value: List[TestResultValue],