acryl-datahub 1.0.0rc5__py3-none-any.whl → 1.0.0rc7__py3-none-any.whl

datahub/ingestion/graph/entity_versioning.py

@@ -93,7 +93,7 @@ class EntityVersioningAPI(DataHubGraphProtocol):
         try:
             return response["linkAssetVersion"]["urn"]
         except KeyError:
-            raise ValueError(f"Unexpected response: {response}")
+            raise ValueError(f"Unexpected response: {response}") from None
 
     def link_asset_to_versioned_asset(
         self,
@@ -165,7 +165,7 @@ class EntityVersioningAPI(DataHubGraphProtocol):
         try:
             return response["unlinkAssetVersion"]["urn"]
         except KeyError:
-            raise ValueError(f"Unexpected response: {response}")
+            raise ValueError(f"Unexpected response: {response}") from None
 
     def unlink_latest_asset_from_version_set(
         self, version_set_urn: str
@@ -198,4 +198,4 @@ class EntityVersioningAPI(DataHubGraphProtocol):
         try:
             return response["unlinkAssetVersion"]["urn"]
         except KeyError:
-            raise ValueError(f"Unexpected response: {response}")
+            raise ValueError(f"Unexpected response: {response}") from None

datahub/ingestion/run/pipeline.py

@@ -9,7 +9,7 @@ import sys
 import threading
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, Iterator, List, Optional, cast
+from typing import Any, Dict, Iterable, Iterator, List, Optional
 
 import click
 import humanfriendly
@@ -26,7 +26,7 @@ from datahub.ingestion.api.common import EndOfStream, PipelineContext, RecordEnv
 from datahub.ingestion.api.global_context import set_graph_context
 from datahub.ingestion.api.pipeline_run_listener import PipelineRunListener
 from datahub.ingestion.api.report import Report
-from datahub.ingestion.api.sink import Sink, SinkReport, WriteCallback
+from datahub.ingestion.api.sink import Sink, SinkReport
 from datahub.ingestion.api.source import Extractor, Source
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.extractor.extractor_registry import extractor_registry
@@ -35,15 +35,15 @@ from datahub.ingestion.reporting.reporting_provider_registry import (
     reporting_provider_registry,
 )
 from datahub.ingestion.run.pipeline_config import PipelineConfig, ReporterConfig
+from datahub.ingestion.run.sink_callback import DeadLetterQueueCallback, LoggingCallback
 from datahub.ingestion.sink.datahub_rest import DatahubRestSink
-from datahub.ingestion.sink.file import FileSink, FileSinkConfig
 from datahub.ingestion.sink.sink_registry import sink_registry
 from datahub.ingestion.source.source_registry import source_registry
 from datahub.ingestion.transformer.system_metadata_transformer import (
     SystemMetadataTransformer,
 )
 from datahub.ingestion.transformer.transform_registry import transform_registry
-from datahub.metadata.schema_classes import MetadataChangeProposalClass
+from datahub.sdk._attribution import KnownAttribution, change_default_attribution
 from datahub.telemetry import stats
 from datahub.telemetry.telemetry import telemetry_instance
 from datahub.utilities._custom_package_loader import model_version_name
@@ -57,68 +57,6 @@ logger = logging.getLogger(__name__)
 _REPORT_PRINT_INTERVAL_SECONDS = 60
 
 
-class LoggingCallback(WriteCallback):
-    def __init__(self, name: str = "") -> None:
-        super().__init__()
-        self.name = name
-
-    def on_success(
-        self, record_envelope: RecordEnvelope, success_metadata: dict
-    ) -> None:
-        logger.debug(
-            f"{self.name} sink wrote workunit {record_envelope.metadata['workunit_id']}"
-        )
-
-    def on_failure(
-        self,
-        record_envelope: RecordEnvelope,
-        failure_exception: Exception,
-        failure_metadata: dict,
-    ) -> None:
-        logger.error(
-            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}",
-            extra={"failure_metadata": failure_metadata},
-            exc_info=failure_exception,
-        )
-
-
-class DeadLetterQueueCallback(WriteCallback):
-    def __init__(self, ctx: PipelineContext, config: Optional[FileSinkConfig]) -> None:
-        if not config:
-            config = FileSinkConfig.parse_obj({"filename": "failed_events.json"})
-        self.file_sink: FileSink = FileSink(ctx, config)
-        self.logging_callback = LoggingCallback(name="failure-queue")
-        logger.info(f"Failure logging enabled. Will log to {config.filename}.")
-
-    def on_success(
-        self, record_envelope: RecordEnvelope, success_metadata: dict
-    ) -> None:
-        pass
-
-    def on_failure(
-        self,
-        record_envelope: RecordEnvelope,
-        failure_exception: Exception,
-        failure_metadata: dict,
-    ) -> None:
-        if "workunit_id" in record_envelope.metadata:
-            if isinstance(record_envelope.record, MetadataChangeProposalClass):
-                mcp = cast(MetadataChangeProposalClass, record_envelope.record)
-                if mcp.systemMetadata:
-                    if not mcp.systemMetadata.properties:
-                        mcp.systemMetadata.properties = {}
-                    if "workunit_id" not in mcp.systemMetadata.properties:
-                        # update the workunit id
-                        mcp.systemMetadata.properties["workunit_id"] = (
-                            record_envelope.metadata["workunit_id"]
-                        )
-                record_envelope.record = mcp
-        self.file_sink.write_record_async(record_envelope, self.logging_callback)
-
-    def close(self) -> None:
-        self.file_sink.close()
-
-
 class PipelineInitError(Exception):
     pass
 
@@ -236,76 +174,99 @@ class Pipeline:
         self.last_time_printed = int(time.time())
         self.cli_report = CliReport()
 
-        self.graph = None
-        with _add_init_error_context("connect to DataHub"):
-            if self.config.datahub_api:
-                self.graph = DataHubGraph(self.config.datahub_api)
-                self.graph.test_connection()
-
-        with _add_init_error_context("set up framework context"):
-            self.ctx = PipelineContext(
-                run_id=self.config.run_id,
-                graph=self.graph,
-                pipeline_name=self.config.pipeline_name,
-                dry_run=dry_run,
-                preview_mode=preview_mode,
-                pipeline_config=self.config,
-            )
-
-        if self.config.sink is None:
-            logger.info(
-                "No sink configured, attempting to use the default datahub-rest sink."
-            )
-            with _add_init_error_context("configure the default rest sink"):
-                self.sink_type = "datahub-rest"
-                self.sink = _make_default_rest_sink(self.ctx)
-        else:
-            self.sink_type = self.config.sink.type
-            with _add_init_error_context(
-                f"find a registered sink for type {self.sink_type}"
-            ):
-                sink_class = sink_registry.get(self.sink_type)
-
-            with _add_init_error_context(f"configure the sink ({self.sink_type})"):
-                sink_config = self.config.sink.dict().get("config") or {}
-                self.sink = sink_class.create(sink_config, self.ctx)
-                logger.debug(f"Sink type {self.sink_type} ({sink_class}) configured")
-        logger.info(f"Sink configured successfully. {self.sink.configured()}")
-
-        if self.graph is None and isinstance(self.sink, DatahubRestSink):
-            with _add_init_error_context("setup default datahub client"):
-                self.graph = self.sink.emitter.to_graph()
-                self.graph.test_connection()
-        self.ctx.graph = self.graph
-        telemetry_instance.set_context(server=self.graph)
-
-        with set_graph_context(self.graph):
-            with _add_init_error_context("configure reporters"):
-                self._configure_reporting(report_to)
-
-            with _add_init_error_context(
-                f"find a registered source for type {self.source_type}"
-            ):
-                source_class = source_registry.get(self.source_type)
-
-            with _add_init_error_context(f"configure the source ({self.source_type})"):
-                self.source = source_class.create(
-                    self.config.source.dict().get("config", {}), self.ctx
-                )
-                logger.debug(
-                    f"Source type {self.source_type} ({source_class}) configured"
+        with contextlib.ExitStack() as exit_stack, contextlib.ExitStack() as inner_exit_stack:
+            self.graph: Optional[DataHubGraph] = None
+            with _add_init_error_context("connect to DataHub"):
+                if self.config.datahub_api:
+                    self.graph = exit_stack.enter_context(
+                        DataHubGraph(self.config.datahub_api)
+                    )
+                    self.graph.test_connection()
+
+            with _add_init_error_context("set up framework context"):
+                self.ctx = PipelineContext(
+                    run_id=self.config.run_id,
+                    graph=self.graph,
+                    pipeline_name=self.config.pipeline_name,
+                    dry_run=dry_run,
+                    preview_mode=preview_mode,
+                    pipeline_config=self.config,
                 )
-                logger.info("Source configured successfully.")
 
-            extractor_type = self.config.source.extractor
-            with _add_init_error_context(f"configure the extractor ({extractor_type})"):
-                extractor_class = extractor_registry.get(extractor_type)
-                self.extractor = extractor_class(
-                    self.config.source.extractor_config, self.ctx
+            if self.config.sink is None:
+                logger.info(
+                    "No sink configured, attempting to use the default datahub-rest sink."
                 )
+                with _add_init_error_context("configure the default rest sink"):
+                    self.sink_type = "datahub-rest"
+                    self.sink = exit_stack.enter_context(
+                        _make_default_rest_sink(self.ctx)
+                    )
+            else:
+                self.sink_type = self.config.sink.type
+                with _add_init_error_context(
+                    f"find a registered sink for type {self.sink_type}"
+                ):
+                    sink_class = sink_registry.get(self.sink_type)
+
+                with _add_init_error_context(f"configure the sink ({self.sink_type})"):
+                    sink_config = self.config.sink.dict().get("config") or {}
+                    self.sink = exit_stack.enter_context(
+                        sink_class.create(sink_config, self.ctx)
+                    )
+                    logger.debug(
+                        f"Sink type {self.sink_type} ({sink_class}) configured"
+                    )
+            logger.info(f"Sink configured successfully. {self.sink.configured()}")
+
+            if self.graph is None and isinstance(self.sink, DatahubRestSink):
+                with _add_init_error_context("setup default datahub client"):
+                    self.graph = self.sink.emitter.to_graph()
+                    self.graph.test_connection()
+            self.ctx.graph = self.graph
+            telemetry_instance.set_context(server=self.graph)
 
-            with _add_init_error_context("configure transformers"):
-                self._configure_transforms()
+            with set_graph_context(self.graph):
+                with _add_init_error_context("configure reporters"):
+                    self._configure_reporting(report_to)
+
+                with _add_init_error_context(
+                    f"find a registered source for type {self.source_type}"
+                ):
+                    source_class = source_registry.get(self.source_type)
+
+                with _add_init_error_context(
+                    f"configure the source ({self.source_type})"
+                ):
+                    self.source = inner_exit_stack.enter_context(
+                        source_class.create(
+                            self.config.source.dict().get("config", {}), self.ctx
+                        )
+                    )
+                    logger.debug(
+                        f"Source type {self.source_type} ({source_class}) configured"
+                    )
+                    logger.info("Source configured successfully.")
+
+                extractor_type = self.config.source.extractor
+                with _add_init_error_context(
+                    f"configure the extractor ({extractor_type})"
+                ):
+                    extractor_class = extractor_registry.get(extractor_type)
+                    self.extractor = inner_exit_stack.enter_context(
+                        extractor_class(self.config.source.extractor_config, self.ctx)
+                    )
+
+                with _add_init_error_context("configure transformers"):
+                    self._configure_transforms()
+
+            # If all of the initialization succeeds, we can preserve the exit stack until the pipeline run.
+            # We need to use an exit stack so that if we have an exception during initialization,
+            # things that were already initialized are still cleaned up.
+            # We need to separate the source/extractor from the rest because stateful
+            # ingestion requires the source to be closed before the state can be updated.
+            self.inner_exit_stack = inner_exit_stack.pop_all()
+            self.exit_stack = exit_stack.pop_all()
 
     @property
     def source_type(self) -> str:
@@ -440,17 +401,19 @@ class Pipeline:
         return False
 
     def run(self) -> None:
-        with contextlib.ExitStack() as stack:
+        with self.exit_stack, self.inner_exit_stack:
             if self.config.flags.generate_memory_profiles:
                 import memray
 
-                stack.enter_context(
+                self.exit_stack.enter_context(
                     memray.Tracker(
                         f"{self.config.flags.generate_memory_profiles}/{self.config.run_id}.bin"
                     )
                 )
 
-            stack.enter_context(self.sink)
+            self.exit_stack.enter_context(
+                change_default_attribution(KnownAttribution.INGESTION)
+            )
 
             self.final_status = PipelineStatus.UNKNOWN
             self._notify_reporters_on_ingestion_start()
@@ -459,8 +422,10 @@ class Pipeline:
                 callback = (
                     LoggingCallback()
                     if not self.config.failure_log.enabled
-                    else DeadLetterQueueCallback(
-                        self.ctx, self.config.failure_log.log_config
+                    else self.exit_stack.enter_context(
+                        DeadLetterQueueCallback(
+                            self.ctx, self.config.failure_log.log_config
+                        )
                     )
                 )
                 for wu in itertools.islice(
@@ -506,12 +471,11 @@ class Pipeline:
                             "Failed to process some records. Continuing.",
                             exc_info=e,
                         )
-                        # TODO: Transformer errors should cause the pipeline to fail.
+                        # TODO: Transformer errors should be reported more loudly / as part of the pipeline report.
 
                     if not self.dry_run:
                         self.sink.handle_work_unit_end(wu)
-                self.extractor.close()
-                self.source.close()
+
                 # no more data is coming, we need to let the transformers produce any additional records if they are holding on to state
                 for record_envelope in self.transform(
                     [
@@ -527,6 +491,11 @@ class Pipeline:
                         # TODO: propagate EndOfStream and other control events to sinks, to allow them to flush etc.
                         self.sink.write_record_async(record_envelope, callback)
 
+                # Stateful ingestion generates the updated state objects as part of the
+                # source's close method. Because of that, we need to close the source
+                # before we call process_commits.
+                self.inner_exit_stack.close()
+
                 self.process_commits()
                 self.final_status = PipelineStatus.COMPLETED
             except (SystemExit, KeyboardInterrupt) as e:
@@ -539,9 +508,6 @@ class Pipeline:
             finally:
                 clear_global_warnings()
 
-                if callback and hasattr(callback, "close"):
-                    callback.close()  # type: ignore
-
                 self._notify_reporters_on_ingestion_completion()
 
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:

datahub/ingestion/run/sink_callback.py

@@ -0,0 +1,77 @@
+import logging
+import threading
+from typing import Optional
+
+from datahub.ingestion.api.closeable import Closeable
+from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
+from datahub.ingestion.api.sink import WriteCallback
+from datahub.ingestion.sink.file import FileSink, FileSinkConfig
+from datahub.metadata.schema_classes import MetadataChangeProposalClass
+
+logger = logging.getLogger(__name__)
+
+
+class LoggingCallback(WriteCallback):
+    def __init__(self, name: str = "") -> None:
+        super().__init__()
+        self.name = name
+
+    def on_success(
+        self, record_envelope: RecordEnvelope, success_metadata: dict
+    ) -> None:
+        logger.debug(
+            f"{self.name} sink wrote workunit {record_envelope.metadata['workunit_id']}"
+        )
+
+    def on_failure(
+        self,
+        record_envelope: RecordEnvelope,
+        failure_exception: Exception,
+        failure_metadata: dict,
+    ) -> None:
+        logger.error(
+            f"{self.name} failed to write record with workunit {record_envelope.metadata['workunit_id']}",
+            extra={"failure_metadata": failure_metadata},
+            exc_info=failure_exception,
+        )
+
+
+class DeadLetterQueueCallback(WriteCallback, Closeable):
+    def __init__(self, ctx: PipelineContext, config: Optional[FileSinkConfig]) -> None:
+        if not config:
+            config = FileSinkConfig.parse_obj({"filename": "failed_events.json"})
+        self.file_sink: FileSink = FileSink(ctx, config)
+        self.file_sink_lock = threading.Lock()
+        self.logging_callback = LoggingCallback(name="failure-queue")
+        logger.info(f"Failure logging enabled. Will log to {config.filename}.")
+
+    def on_success(
+        self, record_envelope: RecordEnvelope, success_metadata: dict
+    ) -> None:
+        pass
+
+    def on_failure(
+        self,
+        record_envelope: RecordEnvelope,
+        failure_exception: Exception,
+        failure_metadata: dict,
+    ) -> None:
+        if "workunit_id" in record_envelope.metadata and isinstance(
+            record_envelope.record, MetadataChangeProposalClass
+        ):
+            mcp: MetadataChangeProposalClass = record_envelope.record
+            if mcp.systemMetadata:
+                if not mcp.systemMetadata.properties:
+                    mcp.systemMetadata.properties = {}
+                if "workunit_id" not in mcp.systemMetadata.properties:
+                    # update the workunit id
+                    mcp.systemMetadata.properties["workunit_id"] = (
+                        record_envelope.metadata["workunit_id"]
+                    )
+            record_envelope.record = mcp
+        with self.file_sink_lock:
+            self.file_sink.write_record_async(record_envelope, self.logging_callback)
+
+    def close(self) -> None:
+        with self.file_sink_lock:
+            self.file_sink.close()