acryl-datahub 1.0.0rc4__py3-none-any.whl → 1.0.0rc6__py3-none-any.whl

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()

datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -292,6 +292,11 @@ class BigQuerySchemaApi:
                         if hasattr(d, "_properties") and isinstance(d._properties, dict)
                         else None
                     ),
+                    # TODO: Fetch dataset description individually impacts overall performance if the number of datasets is high (hundreds); instead we should fetch in batch for all datasets.
+                    # TODO: Given we are calling get_dataset for each dataset, we may consume and publish other fields too, such as created, modified, etc...
+                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_get_dataset
+                    # https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.dataset.Dataset
+                    comment=self.bq_client.get_dataset(d.reference).description,
                 )
                 for d in datasets
             ]

datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -296,6 +296,7 @@ class BigQuerySchemaGenerator:
         self,
         dataset: str,
         project_id: str,
+        description: Optional[str] = None,
         tags: Optional[Dict[str, str]] = None,
         extra_properties: Optional[Dict[str, str]] = None,
     ) -> Iterable[MetadataWorkUnit]:
@@ -336,6 +337,7 @@ class BigQuerySchemaGenerator:
             domain_config=self.config.domain,
             schema_container_key=schema_container_key,
             database_container_key=database_container_key,
+            description=description,
             external_url=(
                 BQ_EXTERNAL_DATASET_URL_TEMPLATE.format(
                     project=project_id, dataset=dataset
@@ -471,14 +473,15 @@ class BigQuerySchemaGenerator:
 
         if self.config.include_schema_metadata:
             yield from self.gen_dataset_containers(
-                dataset_name,
-                project_id,
-                bigquery_dataset.labels,
-                (
+                dataset=dataset_name,
+                project_id=project_id,
+                tags=bigquery_dataset.labels,
+                extra_properties=(
                     {"location": bigquery_dataset.location}
                     if bigquery_dataset.location
                     else None
                 ),
+                description=bigquery_dataset.comment,
             )
 
         columns = None