openmetadata-ingestion 1.4.8.0__py3-none-any.whl → 1.5.0.0rc2__py3-none-any.whl

metadata/workflow/base.py

@@ -12,6 +12,7 @@
 Base workflow definition.
 """
 
+import traceback
 import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime
@@ -47,7 +48,7 @@ from metadata.utils.class_helper import (
 from metadata.utils.execution_time_tracker import ExecutionTimeTracker
 from metadata.utils.helpers import datetime_to_ts
 from metadata.utils.logger import ingestion_logger, set_loggers_level
-from metadata.workflow.output_handler import report_ingestion_status
+from metadata.workflow.workflow_output_handler import WorkflowOutputHandler
 from metadata.workflow.workflow_status_mixin import (
     SUCCESS_THRESHOLD_VALUE,
     WorkflowStatusMixin,
@@ -84,10 +85,12 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         log_level: LogLevels,
         metadata_config: OpenMetadataConnection,
         service_type: ServiceType,
+        output_handler: WorkflowOutputHandler = WorkflowOutputHandler(),
     ):
         """
         Disabling pylint to wait for workflow reimplementation as a topology
         """
+        self.output_handler = output_handler
         self.config = config
         self.service_type = service_type
         self._timer: Optional[RepeatedTimer] = None
@@ -107,7 +110,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         self.post_init()
 
     @property
-    def ingestion_pipeline(self):
+    def ingestion_pipeline(self) -> Optional[IngestionPipeline]:
         """Get or create the Ingestion Pipeline from the configuration"""
         if not self._ingestion_pipeline and self.config.ingestionPipelineFQN:
             self._ingestion_pipeline = self.get_or_create_ingestion_pipeline()
@@ -136,7 +139,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         """
         if not self._timer:
             self._timer = RepeatedTimer(
-                REPORTS_INTERVAL_SECONDS, report_ingestion_status, logger, self
+                REPORTS_INTERVAL_SECONDS, self._report_ingestion_status
             )
 
         return self._timer
@@ -205,7 +208,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         """
         if not self._run_id:
             if self.config.pipelineRunId:
-                self._run_id = str(self.config.pipelineRunId.__root__)
+                self._run_id = str(self.config.pipelineRunId.root)
             else:
                 self._run_id = str(uuid.uuid4())
 
@@ -249,7 +252,7 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
                             ),
                         ),
                         pipelineType=get_pipeline_type_from_source_config(
-                            self.config.source.sourceConfig.config
+                            self.config.source.sourceConfig
                         ),
                         sourceConfig=self.config.source.sourceConfig,
                         airflowConfig=AirflowConfig(),
@@ -277,3 +280,37 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
             entity=get_service_class_from_service_type(self.service_type),
             fqn=self.config.source.serviceName,
         )
+
+    def _report_ingestion_status(self):
+        """
+        Given a logger, use it to INFO the workflow status
+        """
+        try:
+            for step in self.workflow_steps():
+                logger.info(
+                    f"{step.name}: Processed {len(step.status.records)} records,"
+                    f" filtered {len(step.status.filtered)} records,"
+                    f" found {len(step.status.failures)} errors"
+                )
+
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Wild exception reporting status - {exc}")
+
+    def _is_debug_enabled(self) -> bool:
+        return (
+            hasattr(self, "config")
+            and hasattr(self.config, "workflowConfig")
+            and hasattr(self.config.workflowConfig, "loggerLevel")
+            and self.config.workflowConfig.loggerLevel is LogLevels.DEBUG
+        )
+
+    def print_status(self):
+        start_time = self.workflow_steps()[0].get_status().source_start_time
+
+        self.output_handler.print_status(
+            self.result_status(),
+            self.workflow_steps(),
+            start_time,
+            self._is_debug_enabled(),
+        )

metadata/workflow/data_quality.py

@@ -42,7 +42,7 @@ class TestSuiteWorkflow(IngestionWorkflow):
     __test__ = False
 
     def set_steps(self):
-        self.source = TestSuiteSource.create(self.config.dict(), self.metadata)
+        self.source = TestSuiteSource.create(self.config.model_dump(), self.metadata)
 
         test_runner_processor = self._get_test_runner_processor()
         sink = self._get_sink()
@@ -52,14 +52,14 @@ class TestSuiteWorkflow(IngestionWorkflow):
     def _get_sink(self) -> Sink:
         sink_type = self.config.sink.type
         sink_class = import_sink_class(sink_type=sink_type)
-        sink_config = self.config.sink.dict().get("config", {})
+        sink_config = self.config.sink.model_dump().get("config", {})
         sink: Sink = sink_class.create(sink_config, self.metadata)
         logger.debug(f"Sink type:{self.config.sink.type}, {sink_class} configured")
 
         return sink
 
     def _get_test_runner_processor(self) -> Processor:
-        return TestCaseRunner.create(self.config.dict(), self.metadata)
+        return TestCaseRunner.create(self.config.model_dump(), self.metadata)
 
     def _retrieve_service_connection_if_needed(self, service_type: ServiceType) -> None:
         """Get service object from source config `entityFullyQualifiedName`"""
@@ -68,7 +68,7 @@ class TestSuiteWorkflow(IngestionWorkflow):
             and not self.metadata.config.forceEntityOverwriting
         ):
             fully_qualified_name = (
-                self.config.source.sourceConfig.config.entityFullyQualifiedName.__root__
+                self.config.source.sourceConfig.config.entityFullyQualifiedName.root
             )
             try:
                 service_name = fqn.split(fully_qualified_name)[0]
@@ -89,7 +89,7 @@ class TestSuiteWorkflow(IngestionWorkflow):
                     )
 
                 self.config.source.serviceConnection = ServiceConnection(
-                    __root__=service.connection
+                    service.connection
                 )
 
             except Exception as exc:

metadata/workflow/ingestion.py

@@ -185,7 +185,7 @@ class IngestionWorkflow(BaseWorkflow, ABC):
                 )
                 if service:
                     self.config.source.serviceConnection = ServiceConnection(
-                        __root__=service.connection
+                        service.connection
                     )
                 else:
                     raise InvalidWorkflowJSONException(
@@ -205,13 +205,11 @@ class IngestionWorkflow(BaseWorkflow, ABC):
 
     def validate(self):
         try:
-            if (
-                not self.config.source.serviceConnection.__root__.config.supportsProfiler
-            ):
+            if not self.config.source.serviceConnection.root.config.supportsProfiler:
                 raise AttributeError()
         except AttributeError:
-            if ProfilerProcessorConfig.parse_obj(
-                self.config.processor.dict().get("config")
+            if ProfilerProcessorConfig.model_validate(
+                self.config.processor.model_dump().get("config")
             ).ignoreValidation:
                 logger.debug(
                     f"Profiler is not supported for the service connection: {self.config.source.serviceConnection}"

metadata/workflow/metadata.py

@@ -49,7 +49,7 @@ class MetadataWorkflow(IngestionWorkflow):
 
         source_class = (
             import_from_module(
-                self.config.source.serviceConnection.__root__.config.sourcePythonClass
+                self.config.source.serviceConnection.root.config.sourcePythonClass
             )
             if source_type.startswith("custom")
             else import_source_class(
@@ -58,13 +58,13 @@ class MetadataWorkflow(IngestionWorkflow):
         )
 
         pipeline_name = (
-            self.ingestion_pipeline.fullyQualifiedName.__root__
+            self.ingestion_pipeline.fullyQualifiedName.root
             if self.ingestion_pipeline
             else None
         )
 
         source: Source = source_class.create(
-            self.config.source.dict(), self.metadata, pipeline_name
+            self.config.source.model_dump(), self.metadata, pipeline_name
         )
         logger.debug(f"Source type:{source_type},{source_class} configured")
         source.prepare()
@@ -75,7 +75,7 @@ class MetadataWorkflow(IngestionWorkflow):
     def _get_sink(self) -> Sink:
         sink_type = self.config.sink.type
         sink_class = import_sink_class(sink_type=sink_type)
-        sink_config = self.config.sink.dict().get("config", {})
+        sink_config = self.config.sink.model_dump().get("config", {})
         sink: Sink = sink_class.create(sink_config, self.metadata)
         logger.debug(f"Sink type:{self.config.sink.type}, {sink_class} configured")
 

metadata/workflow/output_handler.py

@@ -10,41 +10,21 @@
 #  limitations under the License.
 
 """
-Common Output Handling methods
+Module that handles the legacy WorkflowType until deprecation
 """
-import traceback
 from enum import Enum
-from logging import Logger
-from pathlib import Path
-from typing import Dict, List
+from typing import Optional
 
-from pydantic import BaseModel
-from tabulate import tabulate
-
-from metadata.generated.schema.entity.services.ingestionPipelines.status import (
-    StackTraceError,
+from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
+    PipelineType,
 )
-from metadata.generated.schema.metadataIngestion.workflow import LogLevels
-from metadata.ingestion.api.step import Summary
-from metadata.ingestion.lineage.models import QueryParsingFailures
-from metadata.utils.execution_time_tracker import ExecutionTimeTracker
-from metadata.utils.helpers import pretty_print_time_duration
-from metadata.utils.logger import ANSI, log_ansi_encoded_string
-
-WORKFLOW_FAILURE_MESSAGE = "Workflow finished with failures"
-WORKFLOW_WARNING_MESSAGE = "Workflow finished with warnings"
-WORKFLOW_SUCCESS_MESSAGE = "Workflow finished successfully"
-
-
-class Failure(BaseModel):
-    """
-    Auxiliary class to print the error per status
-    """
-
-    name: str
-    failures: List[StackTraceError]
+from metadata.utils.deprecation import deprecated
 
 
+@deprecated(
+    message="Use 'PipelineType' in 'metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline'",
+    release="1.6",
+)
 class WorkflowType(Enum):
     """
     Workflow type enums based on the `metadata` CLI commands
@@ -59,190 +39,31 @@ class WorkflowType(Enum):
     APP = "application"
 
 
-EXAMPLES_WORKFLOW_PATH: Path = Path(__file__).parent / "../examples" / "workflows"
-
-URLS = {
-    WorkflowType.INGEST: "https://docs.open-metadata.org/connectors/ingestion/workflows/metadata",
-    WorkflowType.PROFILE: "https://docs.open-metadata.org/connectors/ingestion/workflows/profiler",
-    WorkflowType.TEST: "https://docs.open-metadata.org/connectors/ingestion/workflows/data-quality",
-    WorkflowType.LINEAGE: "https://docs.open-metadata.org/connectors/ingestion/workflows/lineage",
-    WorkflowType.USAGE: "https://docs.open-metadata.org/connectors/ingestion/workflows/usage",
-}
-
-DEFAULT_EXAMPLE_FILE = {
-    WorkflowType.INGEST: "bigquery",
-    WorkflowType.PROFILE: "bigquery_profiler",
-    WorkflowType.TEST: "test_suite",
-    WorkflowType.LINEAGE: "bigquery_lineage",
-    WorkflowType.USAGE: "bigquery_usage",
-}
-
-
-def print_more_info(workflow_type: WorkflowType) -> None:
-    """
-    Print more information message
-    """
-    log_ansi_encoded_string(
-        message=f"\nFor more information, please visit: {URLS[workflow_type]}"
-        "\nOr join us in Slack: https://slack.open-metadata.org/"
-    )
-
-
-def print_error_msg(msg: str) -> None:
-    """
-    Print message with error style
-    """
-    log_ansi_encoded_string(color=ANSI.BRIGHT_RED, bold=False, message=f"{msg}")
-
-
-def get_failures(failure: Failure) -> List[Dict[str, str]]:
-    return [
-        {
-            "From": failure.name,
-            "Entity Name": f.name,
-            "Message": f.error,
-            "Stack Trace": f.stackTrace,
-        }
-        for f in failure.failures
-    ]
-
-
-def print_failures_if_apply(failures: List[Failure]) -> None:
-    # take only the ones that contain failures
-    failures = [f for f in failures if f.failures]
-    if failures:
-        # create a list of dictionaries' list
-        all_data = [get_failures(failure) for failure in failures]
-        # create a single of dictionaries
-        data = [f for fs in all_data for f in fs]
-        # create a dictionary with a key and a list of values from the list
-        error_table = {k: [dic[k] for dic in data] for k in data[0]}
-        if len(list(error_table.items())[0][1]) > 100:
-            log_ansi_encoded_string(
-                bold=True, message="Showing only the first 100 failures:"
-            )
-            # truncate list if number of values are over 100
-            error_table = {k: v[:100] for k, v in error_table.items()}
-        else:
-            log_ansi_encoded_string(bold=True, message="List of failures:")
-
-        log_ansi_encoded_string(
-            message=f"\n{tabulate(error_table, headers='keys', tablefmt='grid')}"
-        )
-
-
-def is_debug_enabled(workflow) -> bool:
-    return (
-        hasattr(workflow, "config")
-        and hasattr(workflow.config, "workflowConfig")
-        and hasattr(workflow.config.workflowConfig, "loggerLevel")
-        and workflow.config.workflowConfig.loggerLevel is LogLevels.DEBUG
-    )
-
-
-def print_execution_time_summary():
-    """Log the ExecutionTimeTracker Summary."""
-    tracker = ExecutionTimeTracker()
-
-    summary_table = {
-        "Context": [],
-        "Execution Time Aggregate": [],
+# TODO: Delete this method after the removal of WorkflowType in release 1.6
+# Remember to remove it where it is being used
+def workflow_type_to_pipeline_type(
+    workflow_type: WorkflowType, source_type_name: Optional[str]
+) -> PipelineType:
+    """Helper Function to Map between the Deprecated WorkflowType to PipelineType."""
+
+    def _fix_ingest_type() -> PipelineType:
+        """Helper Function to Map between the Deprecated WorkflowType.INGESTION and the
+        correct PipelineType."""
+        if source_type_name:
+            if source_type_name.endswith("lineage"):
+                return PipelineType.lineage
+            if source_type_name.endswith("usage"):
+                return PipelineType.usage
+        return PipelineType.metadata
+
+    map_ = {
+        WorkflowType.INGEST: _fix_ingest_type(),
+        WorkflowType.PROFILE: PipelineType.profiler,
+        WorkflowType.TEST: PipelineType.TestSuite,
+        WorkflowType.LINEAGE: PipelineType.lineage,
+        WorkflowType.USAGE: PipelineType.usage,
+        WorkflowType.INSIGHT: PipelineType.dataInsight,
+        WorkflowType.APP: PipelineType.application,
     }
 
-    for key in sorted(tracker.state.state.keys()):
-        summary_table["Context"].append(key)
-        summary_table["Execution Time Aggregate"].append(
-            pretty_print_time_duration(tracker.state.state[key])
-        )
-
-    log_ansi_encoded_string(bold=True, message="Execution Time Summary")
-    log_ansi_encoded_string(message=f"\n{tabulate(summary_table, tablefmt='grid')}")
-
-
-def print_query_parsing_issues():
-    """Log the QueryParsingFailures Summary."""
-    query_failures = QueryParsingFailures()
-
-    summary_table = {
-        "Query": [],
-        "Error": [],
-    }
-
-    for failure in query_failures:
-        summary_table["Query"].append(failure.query)
-        summary_table["Error"].append(failure.error)
-
-    if summary_table["Query"]:
-        log_ansi_encoded_string(bold=True, message="Query Parsing Error Summary")
-        log_ansi_encoded_string(
-            message=f"\n{tabulate(summary_table, tablefmt='grid', headers=summary_table.keys())}"
-        )
-
-
-def print_workflow_summary(workflow: "BaseWorkflow") -> None:
-    """
-    Args:
-        workflow: the workflow status to be printed
-
-    Returns:
-        Print Workflow status when the workflow logger level is DEBUG
-    """
-
-    if is_debug_enabled(workflow):
-        print_workflow_status_debug(workflow)
-        print_execution_time_summary()
-        print_query_parsing_issues()
-
-    failures = []
-    total_records = 0
-    total_errors = 0
-    for step in workflow.workflow_steps():
-        step_summary = Summary.from_step(step)
-        total_records += step_summary.records
-        total_errors += step_summary.errors
-        failures.append(Failure(name=step.name, failures=step.get_status().failures))
-
-        log_ansi_encoded_string(bold=True, message=f"Workflow {step.name} Summary:")
-        log_ansi_encoded_string(message=f"Processed records: {step_summary.records}")
-        log_ansi_encoded_string(
-            message=f"Updated records: {step_summary.updated_records}"
-        )
-        log_ansi_encoded_string(message=f"Warnings: {step_summary.warnings}")
-        if step_summary.filtered:
-            log_ansi_encoded_string(message=f"Filtered: {step_summary.filtered}")
-        log_ansi_encoded_string(message=f"Errors: {step_summary.errors}")
-
-    print_failures_if_apply(failures)
-
-    total_success = max(total_records, 1)
-    log_ansi_encoded_string(
-        color=ANSI.BRIGHT_CYAN,
-        bold=True,
-        message=f"Success %: "
-        f"{round(total_success * 100 / (total_success + total_errors), 2)}",
-    )
-
-
-def print_workflow_status_debug(workflow: "BaseWorkflow") -> None:
-    """Print the statuses from each workflow step"""
-    log_ansi_encoded_string(bold=True, message="Statuses detailed info:")
-    for step in workflow.workflow_steps():
-        log_ansi_encoded_string(bold=True, message=f"{step.name} Status:")
-        log_ansi_encoded_string(message=step.get_status().as_string())
-
-
-def report_ingestion_status(logger: Logger, workflow: "BaseWorkflow") -> None:
-    """
-    Given a logger, use it to INFO the workflow status
-    """
-    try:
-        for step in workflow.workflow_steps():
-            logger.info(
-                f"{step.name}: Processed {len(step.status.records)} records,"
-                f" filtered {len(step.status.filtered)} records,"
-                f" found {len(step.status.failures)} errors"
-            )
-
-    except Exception as exc:
-        logger.debug(traceback.format_exc())
-        logger.error(f"Wild exception reporting status - {exc}")
+    return map_.get(workflow_type, PipelineType.metadata)

metadata/workflow/profiler.py

@@ -54,7 +54,7 @@ class ProfilerWorkflow(IngestionWorkflow):
 
     def set_steps(self):
         source_class = self._get_source_class()
-        self.source = source_class.create(self.config.dict(), self.metadata)
+        self.source = source_class.create(self.config.model_dump(), self.metadata)
 
         profiler_processor = self._get_profiler_processor()
         pii_processor = self._get_pii_processor()
@@ -62,7 +62,7 @@ class ProfilerWorkflow(IngestionWorkflow):
         self.steps = (profiler_processor, pii_processor, sink)
 
     def test_connection(self):
-        service_config = self.config.source.serviceConnection.__root__.config
+        service_config = self.config.source.serviceConnection.root.config
         conn = get_ssl_connection(service_config)
 
         test_connection_fn = get_test_connection_fn(service_config)
@@ -71,14 +71,14 @@ class ProfilerWorkflow(IngestionWorkflow):
     def _get_sink(self) -> Sink:
         sink_type = self.config.sink.type
         sink_class = import_sink_class(sink_type=sink_type)
-        sink_config = self.config.sink.dict().get("config", {})
+        sink_config = self.config.sink.model_dump().get("config", {})
         sink: Sink = sink_class.create(sink_config, self.metadata)
         logger.debug(f"Sink type:{self.config.sink.type}, {sink_class} configured")
 
         return sink
 
     def _get_profiler_processor(self) -> Processor:
-        return ProfilerProcessor.create(self.config.dict(), self.metadata)
+        return ProfilerProcessor.create(self.config.model_dump(), self.metadata)
 
     def _get_pii_processor(self) -> Processor:
-        return PIIProcessor.create(self.config.dict(), self.metadata)
+        return PIIProcessor.create(self.config.model_dump(), self.metadata)

metadata/workflow/usage.py

@@ -53,7 +53,7 @@ class UsageWorkflow(IngestionWorkflow):
 
         source_class = (
             import_from_module(
-                self.config.source.serviceConnection.__root__.config.sourcePythonClass
+                self.config.source.serviceConnection.root.config.sourcePythonClass
             )
             if source_type.startswith("custom")
             else import_source_class(
@@ -61,7 +61,9 @@ class UsageWorkflow(IngestionWorkflow):
             )
         )
 
-        source: Source = source_class.create(self.config.source.dict(), self.metadata)
+        source: Source = source_class.create(
+            self.config.source.model_dump(), self.metadata
+        )
         logger.debug(f"Source type:{source_type},{source_class} configured")
         source.prepare()
         logger.debug(f"Source type:{source_type},{source_class}  prepared")
@@ -72,12 +74,12 @@ class UsageWorkflow(IngestionWorkflow):
         """Load the processor class"""
         processor_type = self.config.processor.type
         processor_class = import_processor_class(processor_type=processor_type)
-        processor_config = self.config.processor.dict().get("config", {})
+        processor_config = self.config.processor.model_dump().get("config", {})
         processor: Processor = processor_class.create(
             processor_config,
             self.metadata,
             connection_type=str(
-                self.config.source.serviceConnection.__root__.config.type.value
+                self.config.source.serviceConnection.root.config.type.value
             ),
         )
         logger.debug(f"Processor Type: {processor_type}, {processor_class} configured")
@@ -88,7 +90,7 @@ class UsageWorkflow(IngestionWorkflow):
         """Load the Stage class"""
         stage_type = self.config.stage.type
         stage_class = import_stage_class(stage_type=stage_type)
-        stage_config = self.config.stage.dict().get("config", {})
+        stage_config = self.config.stage.model_dump().get("config", {})
         stage: Stage = stage_class.create(stage_config, self.metadata)
         logger.debug(f"Stage Type: {stage_type}, {stage_class} configured")
 
@@ -98,7 +100,7 @@ class UsageWorkflow(IngestionWorkflow):
         """Load the BulkSink class"""
         bulk_sink_type = self.config.bulkSink.type
         bulk_sink_class = import_bulk_sink_type(bulk_sink_type=bulk_sink_type)
-        bulk_sink_config = self.config.bulkSink.dict().get("config", {})
+        bulk_sink_config = self.config.bulkSink.model_dump().get("config", {})
         bulk_sink: BulkSink = bulk_sink_class.create(bulk_sink_config, self.metadata)
         logger.info(
             f"BulkSink type:{self.config.bulkSink.type},{bulk_sink_class} configured"