acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -13,6 +13,7 @@ from datahub.configuration.common import (
 from datahub.emitter.aspect import JSON_CONTENT_TYPE
 from datahub.emitter.mce_builder import datahub_guid, make_data_platform_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.rest_emitter import EmitMode
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.pipeline_run_listener import PipelineRunListener
 from datahub.ingestion.api.sink import NoopWriteCallback, Sink
@@ -111,6 +112,7 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
     def __init__(self, sink: Sink, report_recipe: bool, ctx: PipelineContext) -> None:
         assert ctx.pipeline_config is not None
 
+        self.ctx = ctx
         self.sink: Sink = sink
         self.report_recipe = report_recipe
         ingestion_source_key = self.generate_unique_key(ctx.pipeline_config)
@@ -191,18 +193,25 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
             )
             return json.dumps(converted_recipe)
 
-    def _emit_aspect(self, entity_urn: Urn, aspect_value: _Aspect) -> None:
-        self.sink.write_record_async(
-            RecordEnvelope(
-                record=MetadataChangeProposalWrapper(
-                    entityUrn=str(entity_urn),
-                    aspect=aspect_value,
-                ),
-                metadata={},
-            ),
-            NoopWriteCallback(),
+    def _emit_aspect(
+        self, entity_urn: Urn, aspect_value: _Aspect, try_sync: bool = False
+    ) -> None:
+        mcp = MetadataChangeProposalWrapper(
+            entityUrn=str(entity_urn),
+            aspect=aspect_value,
         )
 
+        if try_sync and self.ctx.graph:
+            self.ctx.graph.emit_mcp(mcp, emit_mode=EmitMode.SYNC_PRIMARY)
+        else:
+            self.sink.write_record_async(
+                RecordEnvelope(
+                    record=mcp,
+                    metadata={},
+                ),
+                NoopWriteCallback(),
+            )
+
     def on_start(self, ctx: PipelineContext) -> None:
         assert ctx.pipeline_config is not None
         # Construct the dataHubExecutionRequestInput aspect
@@ -223,6 +232,7 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
         self._emit_aspect(
             entity_urn=self.execution_request_input_urn,
             aspect_value=execution_input_aspect,
+            try_sync=True,
         )
 
     def on_completion(
@@ -258,4 +268,4 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
             entity_urn=self.execution_request_input_urn,
             aspect_value=execution_result_aspect,
         )
-        self.sink.close()
+        # Note: sink.close() is handled by the pipeline's context manager

datahub/ingestion/run/pipeline.py

@@ -44,6 +44,10 @@ from datahub.ingestion.transformer.transform_registry import transform_registry
 from datahub.sdk._attribution import KnownAttribution, change_default_attribution
 from datahub.telemetry import stats
 from datahub.telemetry.telemetry import telemetry_instance
+from datahub.upgrade.upgrade import (
+    is_server_default_cli_ahead,
+    retrieve_version_stats,
+)
 from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.global_warning_util import (
     clear_global_warnings,
@@ -171,7 +175,10 @@ class Pipeline:
         self.last_time_printed = int(time.time())
         self.cli_report = CliReport()
 
-        with contextlib.ExitStack() as exit_stack, contextlib.ExitStack() as inner_exit_stack:
+        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:
@@ -258,6 +265,11 @@ class Pipeline:
                 with _add_init_error_context("configure transformers"):
                     self._configure_transforms()
 
+                # Register completion callback with sink to handle final reporting
+                self.sink.register_pre_shutdown_callback(
+                    self._notify_reporters_on_ingestion_completion
+                )
+
             # 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.
@@ -337,8 +349,48 @@ class Pipeline:
         for reporter in self.reporters:
             try:
                 reporter.on_start(ctx=self.ctx)
-            except Exception as e:
-                logger.warning("Reporting failed on start", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on start", exc_info=True)
+
+    def _warn_old_cli_version(self) -> None:
+        """
+        Check if the server default CLI version is ahead of the CLI version being used.
+        If so, add a warning to the report.
+        """
+
+        try:
+            version_stats = retrieve_version_stats(timeout=2.0, graph=self.graph)
+        except RuntimeError as e:
+            # Handle case where there's no event loop available (e.g., in ThreadPoolExecutor)
+            if "no current event loop" in str(e):
+                logger.debug("Skipping version check - no event loop available")
+                return
+            raise
+
+        if not version_stats or not self.graph:
+            return
+
+        if is_server_default_cli_ahead(version_stats):
+            server_default_version = (
+                version_stats.server.current_server_default_cli_version.version
+                if version_stats.server.current_server_default_cli_version
+                else None
+            )
+            current_version = version_stats.client.current.version
+
+            logger.debug(
+                f"""
+                client_version: {current_version}
+                server_default_version: {server_default_version}
+                server_default_cli_ahead: True
+            """
+            )
+
+            self.source.get_report().warning(
+                title="Server default CLI version is ahead of CLI version",
+                message="Please upgrade the CLI version being used",
+                context=f"Server Default CLI version: {server_default_version}, Used CLI version: {current_version}",
+            )
 
     def _notify_reporters_on_ingestion_completion(self) -> None:
         for reporter in self.reporters:
@@ -360,8 +412,8 @@ class Pipeline:
                     report=self._get_structured_report(),
                     ctx=self.ctx,
                 )
-            except Exception as e:
-                logger.warning("Reporting failed on completion", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on completion", exc_info=True)
 
     @classmethod
     def create(
@@ -395,7 +447,20 @@ class Pipeline:
             return True
         return False
 
+    def _set_platform(self) -> None:
+        platform = self.source.infer_platform()
+        if platform:
+            self.source.get_report().set_platform(platform)
+        else:
+            self.source.get_report().warning(
+                message="Platform not found",
+                title="Platform not found",
+                context="Platform not found",
+            )
+
     def run(self) -> None:
+        self._set_platform()
+        self._warn_old_cli_version()
         with self.exit_stack, self.inner_exit_stack:
             if self.config.flags.generate_memory_profiles:
                 import memray
@@ -461,10 +526,10 @@ class Pipeline:
 
                     except (RuntimeError, SystemExit):
                         raise
-                    except Exception as e:
+                    except Exception:
                         logger.error(
                             "Failed to process some records. Continuing.",
-                            exc_info=e,
+                            exc_info=True,
                         )
                         # TODO: Transformer errors should be reported more loudly / as part of the pipeline report.
 
@@ -493,9 +558,9 @@ class Pipeline:
 
                 self.process_commits()
                 self.final_status = PipelineStatus.COMPLETED
-            except (SystemExit, KeyboardInterrupt) as e:
+            except (SystemExit, KeyboardInterrupt):
                 self.final_status = PipelineStatus.CANCELLED
-                logger.error("Caught error", exc_info=e)
+                logger.error("Caught error", exc_info=True)
                 raise
             except Exception as exc:
                 self.final_status = PipelineStatus.ERROR
@@ -503,8 +568,6 @@ class Pipeline:
             finally:
                 clear_global_warnings()
 
-                self._notify_reporters_on_ingestion_completion()
-
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:
         """
         Transforms the given sequence of records by passing the records through the transformers
@@ -578,15 +641,22 @@ class Pipeline:
         sink_failures = len(self.sink.get_report().failures)
         sink_warnings = len(self.sink.get_report().warnings)
         global_warnings = len(get_global_warnings())
+        source_aspects = self.source.get_report().get_aspects_dict()
+        source_aspects_by_subtype = (
+            self.source.get_report().get_aspects_by_subtypes_dict()
+        )
 
         telemetry_instance.ping(
             "ingest_stats",
             {
                 "source_type": self.source_type,
+                "source_aspects": source_aspects,
+                "source_aspects_by_subtype": source_aspects_by_subtype,
                 "sink_type": self.sink_type,
                 "transformer_types": [
                     transformer.type for transformer in self.config.transformers or []
                 ],
+                "extractor_type": self.config.source.extractor,
                 "records_written": stats.discretize(
                     self.sink.get_report().total_records_written
                 ),

datahub/ingestion/run/pipeline_config.py

@@ -6,7 +6,7 @@ from typing import Any, Dict, List, Optional
 
 from pydantic import Field, validator
 
-from datahub.configuration.common import ConfigModel, DynamicTypedConfig
+from datahub.configuration.common import ConfigModel, DynamicTypedConfig, HiddenFromDocs
 from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.ingestion.sink.file import FileSinkConfig
 
@@ -85,7 +85,7 @@ class PipelineConfig(ConfigModel):
     source: SourceConfig
     sink: Optional[DynamicTypedConfig] = None
     transformers: Optional[List[DynamicTypedConfig]] = None
-    flags: FlagsConfig = Field(default=FlagsConfig(), hidden_from_docs=True)
+    flags: HiddenFromDocs[FlagsConfig] = FlagsConfig()
     reporting: List[ReporterConfig] = []
     run_id: str = DEFAULT_RUN_ID
     datahub_api: Optional[DatahubClientConfig] = None

datahub/ingestion/sink/datahub_kafka.py

@@ -74,4 +74,5 @@ class DatahubKafkaSink(Sink[KafkaSinkConfig, SinkReport]):
             callback(err, f"Failed to write record: {err}")
 
     def close(self) -> None:
+        super().close()
         self.emitter.flush()

datahub/ingestion/sink/datahub_rest.py

@@ -3,7 +3,6 @@ import contextlib
 import dataclasses
 import functools
 import logging
-import os
 import threading
 import uuid
 from enum import auto
@@ -16,6 +15,10 @@ from datahub.configuration.common import (
     ConfigurationError,
     OperationalError,
 )
+from datahub.configuration.env_vars import (
+    get_rest_sink_default_max_threads,
+    get_rest_sink_default_mode,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
@@ -47,9 +50,7 @@ from datahub.utilities.server_config_util import set_gms_config
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_REST_SINK_MAX_THREADS = int(
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_MAX_THREADS", 15)
-)
+_DEFAULT_REST_SINK_MAX_THREADS = get_rest_sink_default_max_threads()
 
 
 class RestSinkMode(ConfigEnum):
@@ -63,13 +64,14 @@ class RestSinkMode(ConfigEnum):
 
 
 _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
-    RestSinkMode, os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE", RestSinkMode.ASYNC_BATCH)
+    RestSinkMode, get_rest_sink_default_mode() or RestSinkMode.ASYNC_BATCH
 )
 
 
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
     endpoint: RestSinkEndpoint = DEFAULT_REST_EMITTER_ENDPOINT
+    server_config_refresh_interval: Optional[int] = None
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -90,6 +92,7 @@ class DatahubRestSinkConfig(DatahubClientConfig):
 @dataclasses.dataclass
 class DataHubRestSinkReport(SinkReport):
     mode: Optional[RestSinkMode] = None
+    endpoint: Optional[RestSinkEndpoint] = None
     max_threads: Optional[int] = None
     gms_version: Optional[str] = None
     pending_requests: int = 0
@@ -140,6 +143,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
 
         self.report.gms_version = gms_config.service_version
         self.report.mode = self.config.mode
+        self.report.endpoint = self.config.endpoint
         self.report.max_threads = self.config.max_threads
         logger.debug("Setting env variables to override config")
         logger.debug("Setting gms config")
@@ -346,6 +350,10 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
         )
 
     def close(self):
+        # Execute pre-shutdown callbacks first (handled by parent class)
+        super().close()
+
+        # Then perform sink-specific shutdown
         with self.report.main_thread_blocking_timer:
             self.executor.shutdown()
 

datahub/ingestion/sink/file.py

@@ -79,6 +79,7 @@ class FileSink(Sink[FileSinkConfig, SinkReport]):
             write_callback.on_success(record_envelope, {})
 
     def close(self):
+        super().close()
         self.file.write("\n]")
         self.file.close()
 

datahub/ingestion/source/abs/config.py

@@ -151,7 +151,7 @@ class DataLakeSourceConfig(
             raise ValueError("platform must not be empty")
         return platform
 
-    @pydantic.root_validator()
+    @pydantic.root_validator(skip_on_failure=True)
     def ensure_profiling_pattern_is_passed_to_profiling(
         cls, values: Dict[str, Any]
     ) -> Dict[str, Any]:

datahub/ingestion/source/abs/datalake_profiler_config.py

@@ -72,7 +72,7 @@ class DataLakeProfilerConfig(ConfigModel):
         description="Whether to profile for the sample values for all columns.",
     )
 
-    @pydantic.root_validator()
+    @pydantic.root_validator(skip_on_failure=True)
     def ensure_field_level_settings_are_normalized(
         cls: "DataLakeProfilerConfig", values: Dict[str, Any]
     ) -> Dict[str, Any]:

datahub/ingestion/source/abs/source.py

@@ -44,7 +44,11 @@ from datahub.ingestion.source.azure.abs_utils import (
     get_key_prefix,
     strip_abs_prefix,
 )
-from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
+from datahub.ingestion.source.data_lake_common.data_lake_utils import (
+    ContainerWUCreator,
+    add_partition_columns_to_schema,
+)
 from datahub.ingestion.source.schema_inference import avro, csv_tsv, json, parquet
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -53,10 +57,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
-    SchemaField,
-    SchemaFieldDataType,
     SchemaMetadata,
-    StringTypeClass,
 )
 from datahub.metadata.schema_classes import (
     DataPlatformInstanceClass,
@@ -128,6 +129,14 @@ class TableData:
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.TAGS, "Can extract ABS object/container tags if enabled")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Extract ABS containers and folders",
+    subtype_modifier=[
+        SourceCapabilityModifier.FOLDER,
+        SourceCapabilityModifier.ABS_CONTAINER,
+    ],
+)
 class ABSSource(StatefulIngestionSourceBase):
     source_config: DataLakeSourceConfig
     report: DataLakeSourceReport
@@ -223,36 +232,12 @@ class ABSSource(StatefulIngestionSourceBase):
         fields = sorted(fields, key=lambda f: f.fieldPath)
 
         if self.source_config.add_partition_columns_to_schema:
-            self.add_partition_columns_to_schema(
+            add_partition_columns_to_schema(
                 fields=fields, path_spec=path_spec, full_path=table_data.full_path
             )
 
         return fields
 
-    def add_partition_columns_to_schema(
-        self, path_spec: PathSpec, full_path: str, fields: List[SchemaField]
-    ) -> None:
-        vars = path_spec.get_named_vars(full_path)
-        if vars is not None and "partition" in vars:
-            for partition in vars["partition"].values():
-                partition_arr = partition.split("=")
-                if len(partition_arr) != 2:
-                    logger.debug(
-                        f"Could not derive partition key from partition field {partition}"
-                    )
-                    continue
-                partition_key = partition_arr[0]
-                fields.append(
-                    SchemaField(
-                        fieldPath=f"{partition_key}",
-                        nativeDataType="string",
-                        type=SchemaFieldDataType(StringTypeClass()),
-                        isPartitioningKey=True,
-                        nullable=True,
-                        recursive=False,
-                    )
-                )
-
     def _create_table_operation_aspect(self, table_data: TableData) -> OperationClass:
         reported_time = int(time.time() * 1000)
 
@@ -533,7 +518,7 @@ class ABSSource(StatefulIngestionSourceBase):
             )
             path_spec.sample_files = False
             for obj in container_client.list_blobs(
-                prefix=f"{prefix}", results_per_page=PAGE_SIZE
+                name_starts_with=f"{prefix}", results_per_page=PAGE_SIZE
             ):
                 abs_path = self.create_abs_path(obj.name)
                 logger.debug(f"Path: {abs_path}")