openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/workflow/base.py

@@ -10,20 +10,12 @@
 #  limitations under the License.
 """
 Base workflow definition.
-
-To be extended by any other workflow:
-- ingestion
-- lineage
-- usage
-- profiler
-- test suite
-- data insights
 """
-import traceback
+
 import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Optional, Tuple, TypeVar, cast
+from typing import Any, Dict, List, Optional, TypeVar, Union
 
 from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
     CreateIngestionPipelineRequest,
@@ -31,23 +23,18 @@ from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPi
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
-from metadata.generated.schema.entity.services.connections.serviceConnection import (
-    ServiceConnection,
-)
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     AirflowConfig,
     IngestionPipeline,
     PipelineState,
 )
-from metadata.generated.schema.metadataIngestion.workflow import (
-    OpenMetadataWorkflowConfig,
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
 )
+from metadata.generated.schema.metadataIngestion.workflow import LogLevels
 from metadata.generated.schema.tests.testSuite import ServiceType
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.api.parser import parse_workflow_config_gracefully
 from metadata.ingestion.api.step import Step
-from metadata.ingestion.api.steps import BulkSink, Processor, Sink, Source, Stage
-from metadata.ingestion.models.custom_types import ServiceWithConnectionType
 from metadata.ingestion.ometa.client_utils import create_ometa_client
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.timer.repeated_timer import RepeatedTimer
@@ -56,15 +43,19 @@ from metadata.utils.class_helper import (
     get_pipeline_type_from_source_config,
     get_reference_type_from_service_type,
     get_service_class_from_service_type,
-    get_service_type_from_source_type,
 )
+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.workflow_output_handler import get_ingestion_status_timer
-from metadata.workflow.workflow_status_mixin import WorkflowStatusMixin
+from metadata.workflow.output_handler import report_ingestion_status
+from metadata.workflow.workflow_status_mixin import (
+    SUCCESS_THRESHOLD_VALUE,
+    WorkflowStatusMixin,
+)
 
 logger = ingestion_logger()
 
+# Type of service linked to the Ingestion Pipeline
 T = TypeVar("T")
 
 REPORTS_INTERVAL_SECONDS = 60
@@ -81,110 +72,47 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
     Base workflow implementation
     """
 
-    config: OpenMetadataWorkflowConfig
+    config: Union[Any, Dict]
     _run_id: Optional[str] = None
-    service_type: ServiceType
-    metadata_config: OpenMetadataConnection
     metadata: OpenMetadata
+    metadata_config: OpenMetadataConnection
+    service_type: ServiceType
 
-    # All workflows require a source as a first step
-    source: Source
-    # All workflows execute a series of steps, aside from the source
-    steps: Tuple[Step]
-
-    def __init__(self, config: OpenMetadataWorkflowConfig):
+    def __init__(
+        self,
+        config: Union[Any, Dict],
+        log_level: LogLevels,
+        metadata_config: OpenMetadataConnection,
+        service_type: ServiceType,
+    ):
         """
         Disabling pylint to wait for workflow reimplementation as a topology
         """
         self.config = config
+        self.service_type = service_type
         self._timer: Optional[RepeatedTimer] = None
         self._ingestion_pipeline: Optional[IngestionPipeline] = None
         self._start_ts = datetime_to_ts(datetime.now())
-
-        set_loggers_level(config.workflowConfig.loggerLevel.value)
-
-        self.service_type: ServiceType = get_service_type_from_source_type(
-            self.config.source.type
+        self._execution_time_tracker = ExecutionTimeTracker(
+            log_level == LogLevels.DEBUG
         )
 
-        self.metadata_config: OpenMetadataConnection = (
-            self.config.workflowConfig.openMetadataServerConfig
-        )
+        set_loggers_level(log_level.value)
 
         # We create the ometa client at the workflow level and pass it to the steps
-        self.metadata = create_ometa_client(self.metadata_config)
-
-        self.set_ingestion_pipeline_status(state=PipelineState.running)
+        self.metadata_config = metadata_config
+        self.metadata = create_ometa_client(metadata_config)
 
-        # Pick up the service connection from the API if needed
-        self._retrieve_service_connection_if_needed(self.service_type)
-
-        # Informs the `source` and the rest of `steps` to execute
-        self.set_steps()
+        self.post_init()
 
     @property
     def ingestion_pipeline(self):
         """Get or create the Ingestion Pipeline from the configuration"""
-        if not self._ingestion_pipeline:
+        if not self._ingestion_pipeline and self.config.ingestionPipelineFQN:
             self._ingestion_pipeline = self.get_or_create_ingestion_pipeline()
 
         return self._ingestion_pipeline
 
-    @abstractmethod
-    def set_steps(self):
-        """
-        initialize the tuple of steps to run for each workflow
-        and the source
-        """
-
-    def _execute_internal(self):
-        """
-        Internal execution that needs to be filled
-        by each ingestion workflow.
-
-        Pass each record from the source down the pipeline:
-        Source -> (Processor) -> Sink
-        or Source -> (Processor) -> Stage -> BulkSink
-
-        Note how the Source class needs to be an Iterator. Specifically,
-        we are defining Sources as Generators.
-        """
-        for record in self.source.run():
-            processed_record = record
-            for step in self.steps:
-                # We only process the records for these Step types
-                if processed_record is not None and isinstance(
-                    step, (Processor, Stage, Sink)
-                ):
-                    processed_record = step.run(processed_record)
-
-        # Try to pick up the BulkSink and execute it, if needed
-        bulk_sink = next(
-            (step for step in self.steps if isinstance(step, BulkSink)), None
-        )
-        if bulk_sink:
-            bulk_sink.run()
-
-    def execute(self) -> None:
-        """
-        Main entrypoint
-        """
-        self.timer.trigger()
-        try:
-            self._execute_internal()
-
-            # If we reach this point, compute the success % and update the associated Ingestion Pipeline status
-            self.update_ingestion_status_at_end()
-
-        # Any unhandled exception breaking the workflow should update the status
-        except Exception as err:
-            self.set_ingestion_pipeline_status(PipelineState.failed)
-            raise err
-
-        # Force resource closing. Required for killing the threading
-        finally:
-            self.stop()
-
     def stop(self) -> None:
         """
         Main stopping logic
@@ -194,30 +122,80 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         self.timer.stop()
         self.metadata.close()
 
-        for step in self.steps:
+        for step in self.workflow_steps():
             try:
                 step.close()
             except Exception as exc:
                 logger.warning(f"Error trying to close the step {step} due to [{exc}]")
 
-        self.source.close()
-
     @property
     def timer(self) -> RepeatedTimer:
         """
         Status timer: It will print the source & sink status every `interval` seconds.
         """
         if not self._timer:
-            self._timer = get_ingestion_status_timer(
-                interval=REPORTS_INTERVAL_SECONDS, logger=logger, workflow=self
+            self._timer = RepeatedTimer(
+                REPORTS_INTERVAL_SECONDS, report_ingestion_status, logger, self
             )
 
         return self._timer
 
     @classmethod
-    def create(cls, config_dict: dict) -> "BaseWorkflow":
-        config = parse_workflow_config_gracefully(config_dict)
-        return cls(config)
+    @abstractmethod
+    def create(cls, config_dict: dict):
+        """Single function to execute to create a Workflow instance"""
+
+    @abstractmethod
+    def post_init(self) -> None:
+        """Method to execute after we have initialized all the internals"""
+
+    @abstractmethod
+    def execute_internal(self) -> None:
+        """Workflow-specific logic to execute safely"""
+
+    @abstractmethod
+    def calculate_success(self) -> float:
+        """Get the success % of the internal execution"""
+
+    @abstractmethod
+    def get_failures(self) -> List[StackTraceError]:
+        """Get the failures to flag whether if the workflow succeeded or not"""
+
+    @abstractmethod
+    def workflow_steps(self) -> List[Step]:
+        """Steps to report status from"""
+
+    @abstractmethod
+    def raise_from_status_internal(self, raise_warnings=False) -> None:
+        """Based on the internal workflow status, raise a WorkflowExecutionError"""
+
+    def execute(self) -> None:
+        """
+        Main entrypoint:
+        1. Start logging timer. It will be closed at `stop`
+        2. Execute the workflow
+        3. Validate the pipeline status
+        4. Update the pipeline status at the end
+        """
+        pipeline_state = PipelineState.success
+        self.timer.trigger()
+        try:
+            self.set_ingestion_pipeline_status(state=PipelineState.running)
+            self.execute_internal()
+
+            if SUCCESS_THRESHOLD_VALUE <= self.calculate_success() < 100:
+                pipeline_state = PipelineState.partialSuccess
+
+        # Any unhandled exception breaking the workflow should update the status
+        except Exception as err:
+            pipeline_state = PipelineState.failed
+            raise err
+
+        # Force resource closing. Required for killing the threading
+        finally:
+            ingestion_status = self.build_ingestion_status()
+            self.set_ingestion_pipeline_status(pipeline_state, ingestion_status)
+            self.stop()
 
     @property
     def run_id(self) -> str:
@@ -233,48 +211,6 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
 
         return self._run_id
 
-    def _retrieve_service_connection_if_needed(self, service_type: ServiceType) -> None:
-        """
-        We override the current `serviceConnection` source config object if source workflow service already exists
-        in OM. When secrets' manager is configured, we retrieve the service connection from the secrets' manager.
-        Otherwise, we get the service connection from the service object itself through the default `SecretsManager`.
-
-        :param service_type: source workflow service type
-        :return:
-        """
-        if (
-            not self.config.source.serviceConnection
-            and not self.metadata.config.forceEntityOverwriting
-        ):
-            service_name = self.config.source.serviceName
-            try:
-                service: ServiceWithConnectionType = cast(
-                    ServiceWithConnectionType,
-                    self.metadata.get_by_name(
-                        get_service_class_from_service_type(service_type),
-                        service_name,
-                    ),
-                )
-                if service:
-                    self.config.source.serviceConnection = ServiceConnection(
-                        __root__=service.connection
-                    )
-                else:
-                    raise InvalidWorkflowJSONException(
-                        f"Error getting the service [{service_name}] from the API. If it exists in OpenMetadata,"
-                        " make sure the ingestion-bot JWT token is valid and that the Workflow is deployed"
-                        " with the latest one. If this error persists, recreate the JWT token and"
-                        " redeploy the Workflow."
-                    )
-            except InvalidWorkflowJSONException as exc:
-                raise exc
-            except Exception as exc:
-                logger.debug(traceback.format_exc())
-                logger.error(
-                    f"Unknown error getting service connection for service name [{service_name}]"
-                    f" using the secrets manager provider [{self.metadata.config.secretsManagerProvider}]: {exc}"
-                )
-
     def get_or_create_ingestion_pipeline(self) -> Optional[IngestionPipeline]:
         """
         If we get the `ingestionPipelineFqn` from the `workflowConfig`, it means we want to
@@ -303,7 +239,6 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
             service = self._get_ingestion_pipeline_service()
 
             if service is not None:
-
                 return self.metadata.create_or_update(
                     CreateIngestionPipelineRequest(
                         name=pipeline_name,

metadata/workflow/data_insight.py

@@ -19,12 +19,12 @@ from metadata.ingestion.api.steps import Sink
 from metadata.utils.importer import import_sink_class
 from metadata.utils.logger import data_insight_logger
 from metadata.utils.time_utils import get_beginning_of_day_timestamp_mill
-from metadata.workflow.base import BaseWorkflow
+from metadata.workflow.ingestion import IngestionWorkflow
 
 logger = data_insight_logger()
 
 
-class DataInsightWorkflow(BaseWorkflow):
+class DataInsightWorkflow(IngestionWorkflow):
     """Data insight ingestion workflow implementation"""
 
     retention_days = 7
@@ -39,7 +39,8 @@ class DataInsightWorkflow(BaseWorkflow):
         kpi_runner = KpiRunner(self.metadata)
 
         for kpi_result in kpi_runner.run():
-            self.sink.run(kpi_result)
+            # Pick up the sink
+            self.steps[0].run(kpi_result)
 
     def _clean_up_web_analytics_events(self) -> None:
         """
@@ -64,14 +65,14 @@ class DataInsightWorkflow(BaseWorkflow):
 
         return sink
 
-    def _execute_internal(self):
+    def execute_internal(self):
         """Use parent logic and add step to process KPIs"""
-        super()._execute_internal()
+        super().execute_internal()
         self._run_kpi_processor()
         self._clean_up_web_analytics_events()
 
     def set_steps(self):
         self.source = DataInsightSource.create(self.metadata)  # type: ignore
 
-        self.sink = self._get_sink()
-        self.steps = (self.sink,)
+        sink = self._get_sink()
+        self.steps = (sink,)

metadata/workflow/data_quality.py

@@ -25,12 +25,13 @@ from metadata.ingestion.api.steps import Processor, Sink
 from metadata.utils import fqn
 from metadata.utils.importer import import_sink_class
 from metadata.utils.logger import test_suite_logger
-from metadata.workflow.base import BaseWorkflow, T
+from metadata.workflow.base import T
+from metadata.workflow.ingestion import IngestionWorkflow
 
 logger = test_suite_logger()
 
 
-class TestSuiteWorkflow(BaseWorkflow):
+class TestSuiteWorkflow(IngestionWorkflow):
     """
     DAta Quality ingestion workflow implementation
 

metadata/workflow/ingestion.py

@@ -0,0 +1,203 @@
+#  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.
+"""
+Generic Workflow entrypoint to execute Ingestions
+
+To be extended by any other workflow:
+- metadata
+- lineage
+- usage
+- profiler
+- test suite
+- data insights
+"""
+import traceback
+from abc import ABC, abstractmethod
+from typing import List, Tuple, cast
+
+from metadata.config.common import WorkflowExecutionError
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.entity.services.connections.serviceConnection import (
+    ServiceConnection,
+)
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.api.parser import parse_workflow_config_gracefully
+from metadata.ingestion.api.step import Step, Summary
+from metadata.ingestion.api.steps import BulkSink, Processor, Sink, Source, Stage
+from metadata.ingestion.models.custom_types import ServiceWithConnectionType
+from metadata.utils.class_helper import (
+    get_service_class_from_service_type,
+    get_service_type_from_source_type,
+)
+from metadata.utils.logger import ingestion_logger
+from metadata.workflow.base import BaseWorkflow, InvalidWorkflowJSONException
+from metadata.workflow.workflow_status_mixin import SUCCESS_THRESHOLD_VALUE
+
+logger = ingestion_logger()
+
+
+class IngestionWorkflow(BaseWorkflow, ABC):
+    """
+    Base Ingestion Workflow implementation
+    """
+
+    config: OpenMetadataWorkflowConfig
+
+    # All workflows require a source as a first step
+    source: Source
+    # All workflows execute a series of steps, aside from the source
+    steps: Tuple[Step]
+
+    def __init__(self, config: OpenMetadataWorkflowConfig):
+        self.config = config
+
+        self.service_type: ServiceType = get_service_type_from_source_type(
+            self.config.source.type
+        )
+
+        metadata_config: OpenMetadataConnection = (
+            self.config.workflowConfig.openMetadataServerConfig
+        )
+
+        super().__init__(
+            config=config,
+            log_level=config.workflowConfig.loggerLevel,
+            metadata_config=metadata_config,
+            service_type=self.service_type,
+        )
+
+    @abstractmethod
+    def set_steps(self):
+        """
+        initialize the tuple of steps to run for each workflow
+        and the source
+        """
+
+    def post_init(self) -> None:
+        # Pick up the service connection from the API if needed
+        self._retrieve_service_connection_if_needed(self.service_type)
+
+        # Informs the `source` and the rest of `steps` to execute
+        self.set_steps()
+
+    @classmethod
+    def create(cls, config_dict: dict):
+        config = parse_workflow_config_gracefully(config_dict)
+        return cls(config)
+
+    def execute_internal(self):
+        """
+        Internal execution that needs to be filled
+        by each ingestion workflow.
+
+        Pass each record from the source down the pipeline:
+        Source -> (Processor) -> Sink
+        or Source -> (Processor) -> Stage -> BulkSink
+
+        Note how the Source class needs to be an Iterator. Specifically,
+        we are defining Sources as Generators.
+        """
+        for record in self.source.run():
+            processed_record = record
+            for step in self.steps:
+                # We only process the records for these Step types
+                if processed_record is not None and isinstance(
+                    step, (Processor, Stage, Sink)
+                ):
+                    processed_record = step.run(processed_record)
+
+        # Try to pick up the BulkSink and execute it, if needed
+        bulk_sink = next(
+            (step for step in self.steps if isinstance(step, BulkSink)), None
+        )
+        if bulk_sink:
+            bulk_sink.run()
+
+    def calculate_success(self) -> float:
+        return self.source.get_status().calculate_success()
+
+    def get_failures(self) -> List[StackTraceError]:
+        return self.source.get_status().failures
+
+    def workflow_steps(self) -> List[Step]:
+        return [self.source] + list(self.steps)
+
+    def raise_from_status_internal(self, raise_warnings=False):
+        """
+        Check the status of all steps
+        """
+        if (
+            self.source.get_status().failures
+            and self.calculate_success() < SUCCESS_THRESHOLD_VALUE
+        ):
+            raise WorkflowExecutionError(
+                f"{self.source.name} reported errors: {Summary.from_step(self.source)}"
+            )
+
+        for step in self.steps:
+            if step.status.failures:
+                raise WorkflowExecutionError(
+                    f"{step.name} reported errors: {Summary.from_step(step)}"
+                )
+            if raise_warnings and step.status.warnings:
+                raise WorkflowExecutionError(
+                    f"{step.name} reported warnings: {Summary.from_step(step)}"
+                )
+
+    def _retrieve_service_connection_if_needed(self, service_type: ServiceType) -> None:
+        """
+        We override the current `serviceConnection` source config object if source workflow service already exists
+        in OM. When secrets' manager is configured, we retrieve the service connection from the secrets' manager.
+        Otherwise, we get the service connection from the service object itself through the default `SecretsManager`.
+
+        :param service_type: source workflow service type
+        :return:
+        """
+        if (
+            not self.config.source.serviceConnection
+            and not self.metadata.config.forceEntityOverwriting
+        ):
+            service_name = self.config.source.serviceName
+            try:
+                service: ServiceWithConnectionType = cast(
+                    ServiceWithConnectionType,
+                    self.metadata.get_by_name(
+                        get_service_class_from_service_type(service_type),
+                        service_name,
+                    ),
+                )
+                if service:
+                    self.config.source.serviceConnection = ServiceConnection(
+                        __root__=service.connection
+                    )
+                else:
+                    raise InvalidWorkflowJSONException(
+                        f"Error getting the service [{service_name}] from the API. If it exists in OpenMetadata,"
+                        " make sure the ingestion-bot JWT token is valid and that the Workflow is deployed"
+                        " with the latest one. If this error persists, recreate the JWT token and"
+                        " redeploy the Workflow."
+                    )
+            except InvalidWorkflowJSONException as exc:
+                raise exc
+            except Exception as exc:
+                logger.debug(traceback.format_exc())
+                logger.error(
+                    f"Unknown error getting service connection for service name [{service_name}]"
+                    f" using the secrets manager provider [{self.metadata.config.secretsManagerProvider}]: {exc}"
+                )

metadata/workflow/metadata.py

@@ -20,18 +20,17 @@ from metadata.utils.importer import (
     import_source_class,
 )
 from metadata.utils.logger import ingestion_logger
-from metadata.workflow.base import BaseWorkflow
+from metadata.workflow.ingestion import IngestionWorkflow
 
 logger = ingestion_logger()
 
 
-class MetadataWorkflow(BaseWorkflow):
+class MetadataWorkflow(IngestionWorkflow):
     """
     Metadata ingestion workflow implementation.
     """
 
     def set_steps(self):
-
         # We keep the source registered in the workflow
         self.source = self._get_source()
         sink = self._get_sink()