openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__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/classify.py

@@ -0,0 +1,52 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+
+"""
+Sampler utility for the metadata CLI
+"""
+import sys
+import traceback
+from pathlib import Path
+
+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, redacted_config
+from metadata.workflow.classification import AutoClassificationWorkflow
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
+
+logger = cli_logger()
+
+
+def run_classification(config_path: Path) -> None:
+    """
+    Run the sampler workflow from a config path
+    to a JSON or YAML file
+    :param config_path: Path to load JSON config
+    """
+
+    config_dict = None
+    try:
+        config_dict = load_config_file(config_path)
+        logger.debug("Using workflow config:\n%s", redacted_config(config_dict))
+        workflow = AutoClassificationWorkflow.create(config_dict)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        WorkflowInitErrorHandler.print_init_error(
+            exc, config_dict, PipelineType.metadata
+        )
+        sys.exit(1)
+
+    workflow.execute()
+    workflow.stop()
+    workflow.print_status()
+    workflow.raise_from_status()

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/clients/aws_client.py

@@ -144,7 +144,7 @@ class AWSClient:
     def get_client(self, service_name: str) -> Any:
         # initialize the client depending on the AWSCredentials passed
         if self.config is not None:
-            logger.info(f"Getting AWS client for service [{service_name}]")
+            logger.debug(f"Getting AWS client for service [{service_name}]")
             session = self.create_session()
             if self.config.endPointURL is not None:
                 return session.client(
@@ -152,7 +152,7 @@ class AWSClient:
                 )
             return session.client(service_name=service_name)
 
-        logger.info(f"Getting AWS default client for service [{service_name}]")
+        logger.debug(f"Getting AWS default client for service [{service_name}]")
         # initialized with the credentials loaded from running machine
         return boto3.client(service_name=service_name)
 

metadata/cmd.py

@@ -22,6 +22,7 @@ from typing import List, Optional, Union
 
 from metadata.__version__ import get_metadata_version
 from metadata.cli.app import run_app
+from metadata.cli.classify import run_classification
 from metadata.cli.dataquality import run_test
 from metadata.cli.ingest import run_ingest
 from metadata.cli.lineage import run_lineage
@@ -40,6 +41,7 @@ class MetadataCommands(Enum):
     WEBHOOK = "webhook"
     LINEAGE = "lineage"
     APP = "app"
+    AUTO_CLASSIFICATION = "classify"
 
 
 RUN_PATH_METHODS = {
@@ -49,6 +51,7 @@ RUN_PATH_METHODS = {
     MetadataCommands.PROFILE.value: run_profiler,
     MetadataCommands.TEST.value: run_test,
     MetadataCommands.APP.value: run_app,
+    MetadataCommands.AUTO_CLASSIFICATION.value: run_classification,
 }
 
 
@@ -124,6 +127,12 @@ def get_parser(args: Optional[List[str]] = None):
             help="Workflow for running external applications",
         )
     )
+    create_common_config_parser_args(
+        sub_parser.add_parser(
+            MetadataCommands.AUTO_CLASSIFICATION.value,
+            help="Workflow for running auto classification",
+        )
+    )
     webhook_args(
         sub_parser.add_parser(
             MetadataCommands.WEBHOOK.value,

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 → validator_builder.py}

@@ -14,15 +14,16 @@ Builder interface defining the structure of builders for validators.
 Validators are test classes (e.g. columnValuesToBeBetween, etc.)
 """
 
-from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import TYPE_CHECKING, Optional, Type, Union
+from datetime import datetime, timezone
+from enum import Enum
+from typing import TYPE_CHECKING, Set, Type, Union
 
 from metadata.data_quality.validations.base_test_handler import BaseTestValidator
 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
 
@@ -30,23 +31,29 @@ if TYPE_CHECKING:
     from pandas import DataFrame
 
 
-class IValidatorBuilder(ABC):
-    """Interface for validator builders"""
+class TestCaseImporter:
+    def import_test_case_validator(
+        self,
+        test_type: str,
+        runner_type: str,
+        test_definition: str,
+    ) -> Type[BaseTestValidator]:
+        return import_test_case_class(test_type, runner_type, test_definition)
 
-    @property
-    def test_case(self):
-        """Return the test case object"""
-        return self._test_case
 
-    @property
-    def validator(self):
-        """Return the validator object"""
-        return self._validator
+class SourceType(Enum):
+    PANDAS = "pandas"
+    SQL = "sqlalchemy"
+
+
+class ValidatorBuilder(TestCaseImporter):
+    """Interface for validator builders"""
 
     def __init__(
         self,
         runner: Union[QueryRunner, "DataFrame"],
         test_case: TestCase,
+        source_type: SourceType,
         entity_type: str,
     ) -> None:
         """Builder object for SQA validators. This builder is used to create a validator object
@@ -54,35 +61,45 @@ class IValidatorBuilder(ABC):
         Args:
             runner (QueryRunner): The runner object
             test_case (TestCase): The test case object
+            source_type (SourceType): The source type
             entity_type (str): one of COLUMN or TABLE -- fetched from the test definition
         """
+        super().__init__()
         self._test_case = test_case
         self.runner = runner
-        self.validator_cls: Type[BaseTestValidator] = import_test_case_class(
+        self.validator_cls: Type[
+            BaseTestValidator
+        ] = super().import_test_case_validator(
             entity_type,
-            self._get_source_type(),
-            self.test_case.testDefinition.fullyQualifiedName,  # type: ignore
+            source_type.value,
+            self.test_case.testDefinition.fullyQualifiedName,
         )
         self.reset()
 
-    def set_runtime_params(
-        self, runtime_params_setter: Optional[RuntimeParameterSetter]
-    ):
-        """Set the runtime parameters for the validator object
+    @property
+    def test_case(self):
+        """Return the test case object"""
+        return self._test_case
+
+    @property
+    def validator(self):
+        """Return the validator object"""
+        return self._validator
 
-        # TODO: We should support setting n runtime parameters
+    def set_runtime_params(self, runtime_params_setters: Set[RuntimeParameterSetter]):
+        """Set the runtime parameters for the validator object
 
         Args:
-            runtime_params_setter (Optional[RuntimeParameterSetter]): The runtime parameter setter
+            runtime_params_setters (Optional[RuntimeParameterSetter]): The runtime parameter setter
         """
-        if runtime_params_setter:
-            params = runtime_params_setter.get_parameters(self.test_case)
+        for setter in runtime_params_setters:
+            params = setter.get_parameters(self.test_case)
             if not self.test_case.parameterValues:
                 # If there are no parameters, create a new list
                 self.test_case.parameterValues = []
             self.test_case.parameterValues.append(
                 TestCaseParameterValue(
-                    name="runtimeParams", value=params.model_dump_json()
+                    name=type(params).__name__, value=params.model_dump_json()
                 )
             )
 
@@ -91,10 +108,7 @@ 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
-    def _get_source_type(self):
-        """Get the source type"""
-        raise NotImplementedError