openmetadata-ingestion 1.4.7.2__py3-none-any.whl → 1.5.0.0__py3-none-any.whl

metadata/cli/ingest.py

@@ -17,13 +17,12 @@ 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
 from metadata.workflow.metadata import MetadataWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -39,13 +38,14 @@ def run_ingest(config_path: Path) -> None:
     try:
         config_dict = load_config_file(config_path)
         workflow = MetadataWorkflow.create(config_dict)
-        logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, config_dict, PipelineType.metadata
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/lineage.py

@@ -21,18 +21,21 @@ from pydantic import BaseModel
 
 from metadata.config.common import load_config_file
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
+)
 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.workflow.workflow_output_handler import WorkflowType, print_init_error
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
 
 class LineageWorkflow(BaseModel):
-    filePath: Optional[str]
-    query: Optional[str]
+    filePath: Optional[str] = None
+    query: Optional[str] = None
     checkPatch: Optional[bool] = True
     serviceName: str
     workflowConfig: WorkflowConfig
@@ -49,11 +52,13 @@ def run_lineage(config_path: Path) -> None:
     config_dict = None
     try:
         config_dict = load_config_file(config_path)
-        workflow = LineageWorkflow.parse_obj(config_dict)
+        workflow = LineageWorkflow.model_validate(config_dict)
 
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, config_dict, PipelineType.lineage
+        )
         sys.exit(1)
 
     if workflow.filePath:

metadata/cli/profile.py

@@ -17,13 +17,12 @@ 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
 from metadata.workflow.profiler import ProfilerWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -38,14 +37,15 @@ def run_profiler(config_path: Path) -> None:
     workflow_config_dict = None
     try:
         workflow_config_dict = load_config_file(config_path)
-        logger.debug(f"Using config: {workflow_config_dict}")
         workflow = ProfilerWorkflow.create(workflow_config_dict)
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, workflow_config_dict, WorkflowType.PROFILE)
+        WorkflowInitErrorHandler.print_init_error(
+            exc, workflow_config_dict, PipelineType.profiler
+        )
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/cli/usage.py

@@ -17,13 +17,12 @@ 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
 from metadata.workflow.usage import UsageWorkflow
-from metadata.workflow.workflow_output_handler import (
-    WorkflowType,
-    print_init_error,
-    print_status,
-)
+from metadata.workflow.workflow_init_error_handler import WorkflowInitErrorHandler
 
 logger = cli_logger()
 
@@ -39,13 +38,12 @@ def run_usage(config_path: Path) -> None:
     try:
         config_dict = load_config_file(config_path)
         workflow = UsageWorkflow.create(config_dict)
-        logger.debug(f"Using config: {workflow.config}")
     except Exception as exc:
         logger.debug(traceback.format_exc())
-        print_init_error(exc, config_dict, WorkflowType.INGEST)
+        WorkflowInitErrorHandler.print_init_error(exc, config_dict, PipelineType.usage)
         sys.exit(1)
 
     workflow.execute()
     workflow.stop()
-    print_status(workflow)
+    workflow.print_status()
     workflow.raise_from_status()

metadata/clients/aws_client.py

@@ -47,7 +47,7 @@ class AWSAssumeRoleException(Exception):
 class AWSAssumeRoleCredentialWrapper(BaseModel):
     accessKeyId: str
     secretAccessKey: CustomSecretStr
-    sessionToken: Optional[str]
+    sessionToken: Optional[str] = None
 
 
 class AWSClient:
@@ -59,7 +59,7 @@ class AWSClient:
         self.config = (
             config
             if isinstance(config, AWSCredentials)
-            else (AWSCredentials.parse_obj(config) if config else config)
+            else (AWSCredentials.model_validate(config) if config else config)
         )
 
     @staticmethod
@@ -148,7 +148,7 @@ class AWSClient:
             session = self.create_session()
             if self.config.endPointURL is not None:
                 return session.client(
-                    service_name=service_name, endpoint_url=self.config.endPointURL
+                    service_name=service_name, endpoint_url=str(self.config.endPointURL)
                 )
             return session.client(service_name=service_name)
 
@@ -160,7 +160,7 @@ class AWSClient:
         session = self.create_session()
         if self.config.endPointURL is not None:
             return session.resource(
-                service_name=service_name, endpoint_url=self.config.endPointURL
+                service_name=service_name, endpoint_url=str(self.config.endPointURL)
             )
         return session.resource(service_name=service_name)
 

metadata/clients/azure_client.py

@@ -28,7 +28,7 @@ class AzureClient:
     def __init__(self, credentials: "AzureCredentials"):
         self.credentials = credentials
         if not isinstance(credentials, AzureCredentials):
-            self.credentials = AzureCredentials.parse_obj(credentials)
+            self.credentials = AzureCredentials.model_validate(credentials)
 
     def create_client(
         self,

metadata/clients/domo_client.py

@@ -17,7 +17,7 @@ import traceback
 from dataclasses import dataclass
 from typing import List, Optional, Union
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, ConfigDict
 from pydomo import Domo
 
 from metadata.generated.schema.entity.services.connections.dashboard.domoDashboardConnection import (
@@ -43,8 +43,7 @@ class DomoBaseModel(BaseModel):
     Domo basic configurations
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     id: str
     name: str
@@ -64,10 +63,10 @@ class DomoDashboardDetails(DomoBaseModel):
     Response from Domo API
     """
 
-    cardIds: Optional[List[int]]
-    collectionIds: Optional[List[int]]
-    description: Optional[str]
-    owners: Optional[List[DomoOwner]]
+    cardIds: Optional[List[int]] = None
+    collectionIds: Optional[List[int]] = None
+    description: Optional[str] = None
+    owners: Optional[List[DomoOwner]] = None
 
 
 class DomoChartMetadataDetails(BaseModel):
@@ -75,10 +74,9 @@ class DomoChartMetadataDetails(BaseModel):
     Metadata Details in chart
     """
 
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
-    chartType: Optional[str]
+    chartType: Optional[str] = None
 
 
 class DomoChartDetails(DomoBaseModel):
@@ -87,7 +85,7 @@ class DomoChartDetails(DomoBaseModel):
     """
 
     metadata: DomoChartMetadataDetails
-    description: Optional[str]
+    description: Optional[str] = None
 
 
 class DomoClient:
@@ -103,14 +101,10 @@ class DomoClient:
         ],
     ):
         self.config = config
-        self.config.instanceDomain = (
-            self.config.instanceDomain[:-1]
-            if self.config.instanceDomain.endswith("/")
-            else self.config.instanceDomain
-        )
         HEADERS.update({"X-DOMO-Developer-Token": self.config.accessToken})
         client_config: ClientConfig = ClientConfig(
-            base_url=self.config.instanceDomain,
+            # AnyUrl string ends with / and the domo API does not respond properly if it has 2 // at the end
+            base_url=str(self.config.instanceDomain)[:-1],
             api_version="api/",
             auth_header="Authorization",
             auth_token=lambda: ("no_token", 0),
@@ -132,7 +126,7 @@ class DomoClient:
 
             if isinstance(response, list) and len(response) > 0:
                 return DomoChartDetails(
-                    id=response[0]["id"],
+                    id=str(response[0]["id"]),
                     name=response[0]["title"],
                     metadata=DomoChartMetadataDetails(
                         chartType=response[0].get("metadata", {}).get("chartType", "")

metadata/cmd.py

@@ -17,11 +17,13 @@ from enum import Enum
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
 
+# pyright: reportUnusedCallResult=false
+from typing import List, Optional, Union
+
 from metadata.__version__ import get_metadata_version
 from metadata.cli.app import run_app
 from metadata.cli.dataquality import run_test
 from metadata.cli.ingest import run_ingest
-from metadata.cli.insight import run_insight
 from metadata.cli.lineage import run_lineage
 from metadata.cli.profile import run_profiler
 from metadata.cli.usage import run_usage
@@ -36,7 +38,6 @@ class MetadataCommands(Enum):
     PROFILE = "profile"
     TEST = "test"
     WEBHOOK = "webhook"
-    INSIGHT = "insight"
     LINEAGE = "lineage"
     APP = "app"
 
@@ -45,7 +46,6 @@ RUN_PATH_METHODS = {
     MetadataCommands.INGEST.value: run_ingest,
     MetadataCommands.USAGE.value: run_usage,
     MetadataCommands.LINEAGE.value: run_lineage,
-    MetadataCommands.INSIGHT.value: run_insight,
     MetadataCommands.PROFILE.value: run_profiler,
     MetadataCommands.TEST.value: run_test,
     MetadataCommands.APP.value: run_app,
@@ -88,7 +88,7 @@ def add_metadata_args(parser: argparse.ArgumentParser):
     )
 
 
-def get_parser(args=None):
+def get_parser(args: Optional[List[str]] = None):
     """
     Parser method that returns parsed_args
     """
@@ -130,35 +130,29 @@ def get_parser(args=None):
             help="Simple Webserver to test webhook metadata events",
         )
     )
-    create_common_config_parser_args(
-        sub_parser.add_parser(
-            MetadataCommands.INSIGHT.value, help="Data Insights Workflow"
-        )
-    )
 
     add_metadata_args(parser)
     parser.add_argument("--debug", help="Debug Mode", action="store_true")
     return parser.parse_args(args)
 
 
-def metadata(args=None):
+def metadata(args: Optional[List[str]] = None):
     """
     This method implements parsing of the arguments passed from CLI
     """
     contains_args = vars(get_parser(args))
     metadata_workflow = contains_args.get("command")
-    config_file = contains_args.get("config")
+    config_file: Optional[Path] = contains_args.get("config")
     path = None
     if config_file:
-        path = Path(config_file).expanduser()
+        path = config_file.expanduser()
     if contains_args.get("debug"):
         set_loggers_level(logging.DEBUG)
-    elif contains_args.get("log_level"):
-        set_loggers_level(contains_args.get("log_level"))
     else:
-        set_loggers_level(logging.INFO)
+        log_level: Union[str, int] = contains_args.get("log_level") or logging.INFO
+        set_loggers_level(log_level)
 
-    if metadata_workflow in RUN_PATH_METHODS:
+    if path and metadata_workflow and metadata_workflow in RUN_PATH_METHODS:
         RUN_PATH_METHODS[metadata_workflow](path)
 
     if metadata_workflow == MetadataCommands.WEBHOOK.value:
@@ -171,17 +165,19 @@ def metadata(args=None):
                 self.wfile.write(bytes("Hello, World! Here is a GET response", "utf8"))
 
             def do_POST(self):  # pylint: disable=invalid-name
-                content_len = int(self.headers.get("Content-Length"))
-                post_body = self.rfile.read(content_len)
+                if self.headers.get("Content-Length"):
+                    content_len = int(self.headers["Content-Length"])
+                    post_body = self.rfile.read(content_len)
+                    logger.info(post_body)
+
                 self.send_response(200)
                 self.send_header("Content-type", "application/json")
                 self.end_headers()
-                logger.info(post_body)
 
         logger.info(
             f"Starting server at {contains_args.get('host')}:{contains_args.get('port')}"
         )
         with HTTPServer(
-            (contains_args.get("host"), contains_args.get("port")), WebhookHandler
+            (contains_args["host"], contains_args["port"]), WebhookHandler
         ) as server:
             server.serve_forever()

metadata/config/common.py

@@ -19,21 +19,20 @@ from abc import ABC, abstractmethod
 from typing import IO, Any, Optional
 
 import yaml
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 
 class ConfigModel(BaseModel):
     """Class definition for config model"""
 
-    class Config:
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class DynamicTypedConfig(ConfigModel):
     """Class definition for Dynamic Typed Config"""
 
     type: str
-    config: Optional[Any]
+    config: Optional[Any] = None
 
 
 class WorkflowExecutionError(Exception):

metadata/data_quality/api/models.py

@@ -35,7 +35,7 @@ class TestCaseDefinition(ConfigModel):
     description: Optional[str] = None
     testDefinitionName: str
     columnName: Optional[str] = None
-    parameterValues: Optional[List[TestCaseParameterValue]]
+    parameterValues: Optional[List[TestCaseParameterValue]] = None
     computePassedFailedRowCount: Optional[bool] = False
 
 
@@ -56,7 +56,7 @@ class TableAndTests(BaseModel):
 
     table: Table = Field(None, description="Table being processed by the DQ workflow")
     service_type: str = Field(..., description="Service type the table belongs to")
-    test_cases: Optional[List[TestCase]] = Field(
+    test_cases: List[TestCase] = Field(
         None, description="Test Cases already existing in the Test Suite, if any"
     )
     executable_test_suite: Optional[CreateTestSuiteRequest] = Field(

metadata/data_quality/builders/i_validator_builder.py

@@ -0,0 +1,100 @@
+#  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.
+
+"""
+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 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.profiler.processor.runner import QueryRunner
+from metadata.utils.importer import import_test_case_class
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
+
+class IValidatorBuilder(ABC):
+    """Interface for validator builders"""
+
+    @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
+
+    def __init__(
+        self,
+        runner: Union[QueryRunner, "DataFrame"],
+        test_case: TestCase,
+        entity_type: str,
+    ) -> None:
+        """Builder object for SQA validators. This builder is used to create a validator object
+
+        Args:
+            runner (QueryRunner): The runner object
+            test_case (TestCase): The test case object
+            entity_type (str): one of COLUMN or TABLE -- fetched from the test definition
+        """
+        self._test_case = test_case
+        self.runner = runner
+        self.validator_cls: Type[BaseTestValidator] = import_test_case_class(
+            entity_type,
+            self._get_source_type(),
+            self.test_case.testDefinition.fullyQualifiedName,  # type: ignore
+        )
+        self.reset()
+
+    def set_runtime_params(
+        self, runtime_params_setter: Optional[RuntimeParameterSetter]
+    ):
+        """Set the runtime parameters for the validator object
+
+        # TODO: We should support setting n runtime parameters
+
+        Args:
+            runtime_params_setter (Optional[RuntimeParameterSetter]): The runtime parameter setter
+        """
+        if runtime_params_setter:
+            params = runtime_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()
+                )
+            )
+
+    def reset(self):
+        """Reset the builder"""
+        self._validator = self.validator_cls(
+            self.runner,
+            test_case=self.test_case,
+            execution_date=int(datetime.now().timestamp() * 1000),
+        )
+
+    @abstractmethod
+    def _get_source_type(self):
+        """Get the source type"""
+        raise NotImplementedError

metadata/{data_insight/producer/producer_interface.py → data_quality/builders/pandas_validator_builder.py}

@@ -8,22 +8,22 @@
 #  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.
+
 """
-OpenMetadata data insight producer interface. This is an abstract class that 
-defines the interface for all data insight producers.
+Builder defining the structure of builders for validators for Pandas sources
 """
 
-from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
 
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
+if TYPE_CHECKING:
+    pass
 
 
-class ProducerInterface(ABC):
-    def __init__(self, metadata: OpenMetadata):
-        """instantiate a producer object"""
-        self.metadata = metadata
+class PandasValidatorBuilder(IValidatorBuilder):
+    """Builder for Pandas validators"""
 
-    @abstractmethod
-    def fetch_data(self, limit, fields, entities_cache=None):
-        """fetch data from source"""
-        raise NotImplementedError
+    def _get_source_type(self) -> str:
+        """Return the test case"""
+        return "pandas"

metadata/data_quality/{validations/validator.py → builders/sqa_validator_builder.py}

@@ -10,17 +10,16 @@
 #  limitations under the License.
 
 """
-Test case validator object
+Builder defining the structure of builders for validators for SQA sources
 """
 
-from metadata.generated.schema.tests.basic import TestCaseResult
 
+from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
 
-class Validator:
-    """Test case validator object. it take test handler obkect and run validation"""
 
-    def __init__(self, validator_obj):
-        self.validator_obj = validator_obj
+class SQAValidatorBuilder(IValidatorBuilder):
+    """Builder for SQA validators"""
 
-    def validate(self) -> TestCaseResult:
-        return self.validator_obj.run_validation()
+    def _get_source_type(self) -> str:
+        """Return the test case"""
+        return "sqlalchemy"

metadata/data_quality/interface/pandas/pandas_test_suite_interface.py

@@ -13,22 +13,20 @@
 Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
-from datetime import datetime, timezone
-from typing import Optional
 
+from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
+from metadata.data_quality.builders.pandas_validator_builder import (
+    PandasValidatorBuilder,
+)
 from metadata.data_quality.interface.test_suite_interface import TestSuiteInterface
-from metadata.data_quality.validations.validator import Validator
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
-from metadata.generated.schema.tests.basic import TestCaseResult
 from metadata.generated.schema.tests.testCase import TestCase
-from metadata.generated.schema.tests.testDefinition import TestDefinition
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.mixins.pandas.pandas_mixin import PandasInterfaceMixin
-from metadata.utils.importer import import_test_case_class
 from metadata.utils.logger import test_suite_logger
 
 logger = test_suite_logger()
@@ -62,45 +60,14 @@ class PandasTestSuiteInterface(TestSuiteInterface, PandasInterfaceMixin):
         # add partition logic to test suite
         self.dfs = self.return_ometa_dataframes_sampled(
             service_connection_config=self.service_connection_config,
-            client=get_connection(self.service_connection_config).client,
+            client=get_connection(self.service_connection_config).client._client,
             table=self.table_entity,
             profile_sample_config=self.table_sample_config,
         )
         if self.dfs and self.table_partition_config:
             self.dfs = self.get_partitioned_df(self.dfs)
 
-    def run_test_case(
-        self,
-        test_case: TestCase,
-    ) -> Optional[TestCaseResult]:
-        """Run table tests where platformsTest=OpenMetadata
-
-        Args:
-            test_case: test case object to execute
-
-        Returns:
-            TestCaseResult object
-        """
-
-        try:
-            TestHandler = import_test_case_class(  # pylint: disable=invalid-name
-                self.ometa_client.get_by_id(
-                    TestDefinition, test_case.testDefinition.id
-                ).entityType.value,
-                "pandas",
-                test_case.testDefinition.fullyQualifiedName,
-            )
-
-            test_handler = TestHandler(
-                self.dfs,
-                test_case=test_case,
-                execution_date=int(datetime.now(tz=timezone.utc).timestamp() * 1000),
-            )
-
-            return Validator(validator_obj=test_handler).validate()
-        except Exception as err:
-            logger.error(
-                f"Error executing {test_case.testDefinition.fullyQualifiedName} - {err}"
-            )
-
-            raise RuntimeError(err)
+    def _get_validator_builder(
+        self, test_case: TestCase, entity_type: str
+    ) -> IValidatorBuilder:
+        return PandasValidatorBuilder(self.dfs, test_case, entity_type)