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

datahub/ingestion/api/auto_work_units/auto_ensure_aspect_size.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 from typing import TYPE_CHECKING, Iterable, List
 
 from datahub.emitter.rest_emitter import INGEST_MAX_PAYLOAD_BYTES
@@ -7,15 +8,36 @@ from datahub.emitter.serialization_helper import pre_json_transform
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
     DatasetProfileClass,
+    QueryPropertiesClass,
+    QuerySubjectsClass,
     SchemaFieldClass,
     SchemaMetadataClass,
+    UpstreamLineageClass,
 )
 
 if TYPE_CHECKING:
     from datahub.ingestion.api.source import SourceReport
 
+
+# TODO: ordering
+# In the cases where we trim collections of data (e.g. fields in schema, upstream lineage, query subjects), given
+# those collections are typically unordered, we should consider sorting them by some criteria (e.g. size, alphabetically)
+# so that the trimming is deterministic and predictable and more importantly consistent across executions.
+# In the case of schemaMetadata, that's more relevant as currently we may be trimming fields while adding nested ones,
+# which may lead to poorly schema rendering in the UI.
+
 logger = logging.getLogger(__name__)
 
+DEFAULT_QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES = 5 * 1024 * 1024  # 5MB
+QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES = int(
+    os.environ.get(
+        "QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES",
+        DEFAULT_QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES,
+    )
+)
+
+QUERY_STATEMENT_TRUNCATION_BUFFER = 100
+
 
 class EnsureAspectSizeProcessor:
     def __init__(
@@ -81,6 +103,274 @@ class EnsureAspectSizeProcessor:
 
         schema.fields = accepted_fields
 
+    def ensure_query_subjects_size(
+        self, entity_urn: str, query_subjects: QuerySubjectsClass
+    ) -> None:
+        """
+        Ensure query subjects aspect does not exceed allowed size by removing column-level lineage first,
+        then table lineage if necessary.
+        """
+        if not query_subjects.subjects:
+            return
+
+        total_subjects_size = 0
+        accepted_table_level_subjects = []
+        accepted_column_level_subjects = []
+        column_level_subjects_with_sizes = []
+        table_level_subjects_with_sizes = []
+
+        # Separate column-level and table-level subjects
+        for subject in query_subjects.subjects:
+            subject_size = len(json.dumps(pre_json_transform(subject.to_obj())))
+
+            if subject.entity.startswith("urn:li:schemaField:"):
+                column_level_subjects_with_sizes.append((subject, subject_size))
+            else:
+                table_level_subjects_with_sizes.append((subject, subject_size))
+
+        # Once we find one that doesn't fit, stop everything else to prevent inconsistencies
+        first_skip_done = False
+
+        # First, try to include all table-level subjects
+        for subject, subject_size in table_level_subjects_with_sizes:
+            if total_subjects_size + subject_size < self.payload_constraint:
+                accepted_table_level_subjects.append(subject)
+                total_subjects_size += subject_size
+            else:
+                first_skip_done = True
+                break
+
+        # Then, add column-level subjects if there's remaining space
+        # Only process if we successfully included all table-level subjects
+        if not first_skip_done:
+            for subject, subject_size in column_level_subjects_with_sizes:
+                if total_subjects_size + subject_size < self.payload_constraint:
+                    accepted_column_level_subjects.append(subject)
+                    total_subjects_size += subject_size
+                else:
+                    first_skip_done = True
+                    break
+
+        if first_skip_done:
+            # Log aggregate warnings
+            table_level_skipped_count = len(table_level_subjects_with_sizes) - len(
+                accepted_table_level_subjects
+            )
+            column_level_skipped_count = len(column_level_subjects_with_sizes) - len(
+                accepted_column_level_subjects
+            )
+
+            self._maybe_warn_query_subjects(
+                entity_urn, table_level_skipped_count, "table-level lineage subjects"
+            )
+            self._maybe_warn_query_subjects(
+                entity_urn, column_level_skipped_count, "column-level lineage subjects"
+            )
+
+        query_subjects.subjects = (
+            accepted_table_level_subjects + accepted_column_level_subjects
+        )
+
+    def _maybe_warn_query_subjects(
+        self, entity_urn: str, skipped_count: int, item_type: str
+    ) -> None:
+        """Log warning for query subjects truncation if any items were skipped."""
+        if skipped_count > 0:
+            self.report.warning(
+                title="Query subjects truncated due to size constraint",
+                message="Query subjects contained too much data and would have caused ingestion to fail",
+                context=f"Skipped {skipped_count} {item_type} for {entity_urn} due to aspect size constraints",
+            )
+
+    def _maybe_warn_upstream_lineage(
+        self, entity_urn: str, skipped_count: int, item_type: str
+    ) -> None:
+        """Log warning for upstream lineage truncation if any items were skipped."""
+        if skipped_count > 0:
+            self.report.warning(
+                title="Upstream lineage truncated due to size constraint",
+                message="Upstream lineage contained too much data and would have caused ingestion to fail",
+                context=f"Skipped {skipped_count} {item_type} for {entity_urn} due to aspect size constraints",
+            )
+
+    def ensure_upstream_lineage_size(  # noqa: C901
+        self, entity_urn: str, upstream_lineage: UpstreamLineageClass
+    ) -> None:
+        """
+        Ensure upstream lineage aspect does not exceed allowed size by removing lineage in priority order:
+        first NONE fine-grained lineages (lowest priority), then FIELD_SET fine-grained lineages,
+        then DATASET fine-grained lineages, and finally upstreams (highest priority).
+        """
+        if not upstream_lineage.fineGrainedLineages and not upstream_lineage.upstreams:
+            return
+
+        total_lineage_size = 0
+        accepted_upstreams = []
+        accepted_dataset_fg_lineages = []
+        accepted_field_set_fg_lineages = []
+        accepted_none_fg_lineages = []
+        upstream_items_with_sizes = []
+        dataset_fg_items_with_sizes = []
+        field_set_fg_items_with_sizes = []
+        none_fg_items_with_sizes = []
+
+        # Add upstreams (highest priority)
+        if upstream_lineage.upstreams:
+            for upstream in upstream_lineage.upstreams:
+                upstream_size = len(json.dumps(pre_json_transform(upstream.to_obj())))
+                upstream_items_with_sizes.append((upstream, upstream_size))
+
+        # Separate fine-grained lineage items by upstreamType: DATASET > FIELD_SET > NONE
+        if upstream_lineage.fineGrainedLineages:
+            for fg_lineage in upstream_lineage.fineGrainedLineages:
+                fg_lineage_size = len(
+                    json.dumps(pre_json_transform(fg_lineage.to_obj()))
+                )
+
+                upstream_type_str = str(fg_lineage.upstreamType)
+                if upstream_type_str == "DATASET":
+                    dataset_fg_items_with_sizes.append((fg_lineage, fg_lineage_size))
+                elif upstream_type_str == "FIELD_SET":
+                    field_set_fg_items_with_sizes.append((fg_lineage, fg_lineage_size))
+                elif upstream_type_str == "NONE":
+                    none_fg_items_with_sizes.append((fg_lineage, fg_lineage_size))
+
+        # Once we find one that doesn't fit, stop everything else to prevent inconsistencies
+        first_skip_done = False
+
+        # First, include all upstreams (highest priority)
+        for item, item_size in upstream_items_with_sizes:
+            if total_lineage_size + item_size < self.payload_constraint:
+                accepted_upstreams.append(item)
+                total_lineage_size += item_size
+            else:
+                first_skip_done = True
+                break
+
+        # Second, include DATASET fine-grained lineages if no upstreams were skipped
+        if not first_skip_done:
+            for fg_lineage, fg_lineage_size in dataset_fg_items_with_sizes:
+                if total_lineage_size + fg_lineage_size < self.payload_constraint:
+                    accepted_dataset_fg_lineages.append(fg_lineage)
+                    total_lineage_size += fg_lineage_size
+                else:
+                    first_skip_done = True
+                    break
+
+        # Third, include FIELD_SET fine-grained lineages if no higher priority items were skipped
+        if not first_skip_done:
+            for fg_lineage, fg_lineage_size in field_set_fg_items_with_sizes:
+                if total_lineage_size + fg_lineage_size < self.payload_constraint:
+                    accepted_field_set_fg_lineages.append(fg_lineage)
+                    total_lineage_size += fg_lineage_size
+                else:
+                    first_skip_done = True
+                    break
+
+        # Finally, include NONE fine-grained lineages if no higher priority items were skipped
+        if not first_skip_done:
+            for fg_lineage, fg_lineage_size in none_fg_items_with_sizes:
+                if total_lineage_size + fg_lineage_size < self.payload_constraint:
+                    accepted_none_fg_lineages.append(fg_lineage)
+                    total_lineage_size += fg_lineage_size
+                else:
+                    first_skip_done = True
+                    break
+
+        # Log aggregate warnings instead of per-item warnings
+        if first_skip_done:
+            upstreams_skipped_count = len(upstream_items_with_sizes) - len(
+                accepted_upstreams
+            )
+            dataset_fg_skipped_count = len(dataset_fg_items_with_sizes) - len(
+                accepted_dataset_fg_lineages
+            )
+            field_set_fg_skipped_count = len(field_set_fg_items_with_sizes) - len(
+                accepted_field_set_fg_lineages
+            )
+            none_fg_skipped_count = len(none_fg_items_with_sizes) - len(
+                accepted_none_fg_lineages
+            )
+
+            self._maybe_warn_upstream_lineage(
+                entity_urn, upstreams_skipped_count, "upstream datasets"
+            )
+            self._maybe_warn_upstream_lineage(
+                entity_urn,
+                dataset_fg_skipped_count,
+                "dataset-level fine-grained lineages",
+            )
+            self._maybe_warn_upstream_lineage(
+                entity_urn,
+                field_set_fg_skipped_count,
+                "field-set-level fine-grained lineages",
+            )
+            self._maybe_warn_upstream_lineage(
+                entity_urn, none_fg_skipped_count, "none-level fine-grained lineages"
+            )
+
+        # Combine all accepted fine-grained lineages
+        accepted_fine_grained_lineages = (
+            accepted_dataset_fg_lineages
+            + accepted_field_set_fg_lineages
+            + accepted_none_fg_lineages
+        )
+
+        upstream_lineage.upstreams = accepted_upstreams
+        upstream_lineage.fineGrainedLineages = (
+            accepted_fine_grained_lineages if accepted_fine_grained_lineages else None
+        )
+
+    def ensure_query_properties_size(
+        self, entity_urn: str, query_properties: QueryPropertiesClass
+    ) -> None:
+        """
+        Ensure query properties aspect does not exceed allowed size by truncating the query statement value.
+        Uses a configurable max payload size that is the minimum between QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES
+        and INGEST_MAX_PAYLOAD_BYTES.
+
+        We have found surprisingly large query statements (e.g. 20MB+) that caused ingestion to fail;
+        that was INSERT INTO VALUES with huge list of values.
+        """
+        if not query_properties.statement or not query_properties.statement.value:
+            return
+
+        max_payload_size = min(
+            QUERY_PROPERTIES_STATEMENT_MAX_PAYLOAD_BYTES, self.payload_constraint
+        )
+
+        current_size = len(json.dumps(pre_json_transform(query_properties.to_obj())))
+
+        if current_size < max_payload_size:
+            return
+
+        reduction_needed = (
+            current_size - max_payload_size + QUERY_STATEMENT_TRUNCATION_BUFFER
+        )
+
+        statement_value_size = len(query_properties.statement.value)
+        original_statement_size = statement_value_size
+
+        # Only truncate if reduction is actually needed and possible
+        if statement_value_size > reduction_needed > 0:
+            new_statement_length = statement_value_size - reduction_needed
+            truncated_statement = query_properties.statement.value[
+                :new_statement_length
+            ]
+
+            truncation_message = f"... [original value was {original_statement_size} bytes and truncated to {new_statement_length} bytes]"
+            query_properties.statement.value = truncated_statement + truncation_message
+
+            self.report.warning(
+                title="Query properties truncated due to size constraint",
+                message="Query properties contained too much data and would have caused ingestion to fail",
+                context=f"Query statement was truncated from {original_statement_size} to {new_statement_length} characters for {entity_urn} due to aspect size constraints",
+            )
+        else:
+            logger.warning(
+                f"Cannot truncate query statement for {entity_urn} as it is smaller than or equal to the required reduction size {reduction_needed}. That means that 'ensure_query_properties_size' must be extended to trim other fields different than statement."
+            )
+
     def ensure_aspect_size(
         self,
         stream: Iterable[MetadataWorkUnit],
@@ -90,10 +380,16 @@ class EnsureAspectSizeProcessor:
         on GMS side and failure of the entire ingestion. This processor will attempt to trim suspected aspects.
         """
         for wu in stream:
-            logger.debug(f"Ensuring size of workunit: {wu.id}")
+            # logger.debug(f"Ensuring size of workunit: {wu.id}")
 
             if schema := wu.get_aspect_of_type(SchemaMetadataClass):
                 self.ensure_schema_metadata_size(wu.get_urn(), schema)
             elif profile := wu.get_aspect_of_type(DatasetProfileClass):
                 self.ensure_dataset_profile_size(wu.get_urn(), profile)
+            elif query_subjects := wu.get_aspect_of_type(QuerySubjectsClass):
+                self.ensure_query_subjects_size(wu.get_urn(), query_subjects)
+            elif upstream_lineage := wu.get_aspect_of_type(UpstreamLineageClass):
+                self.ensure_upstream_lineage_size(wu.get_urn(), upstream_lineage)
+            elif query_properties := wu.get_aspect_of_type(QueryPropertiesClass):
+                self.ensure_query_properties_size(wu.get_urn(), query_properties)
             yield wu

datahub/ingestion/api/auto_work_units/auto_validate_input_fields.py

@@ -0,0 +1,87 @@
+import logging
+from typing import TYPE_CHECKING, Iterable, List
+
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import InputFieldClass, InputFieldsClass
+
+if TYPE_CHECKING:
+    from datahub.ingestion.api.source import SourceReport
+
+logger = logging.getLogger(__name__)
+
+
+class ValidateInputFieldsProcessor:
+    def __init__(self, report: "SourceReport"):
+        self.report = report
+
+    def validate_input_fields(
+        self,
+        stream: Iterable[MetadataWorkUnit],
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Validate input fields and filter out invalid ones.
+
+        Invalid input fields have empty or missing fieldPath values, which would cause
+        URN generation to fail when sent to the server. This processor filters them out
+        and reports them as warnings.
+        """
+        for wu in stream:
+            input_fields_aspect = wu.get_aspect_of_type(InputFieldsClass)
+            if input_fields_aspect and input_fields_aspect.fields:
+                valid_fields: List[InputFieldClass] = []
+                invalid_count = 0
+
+                for input_field in input_fields_aspect.fields:
+                    if (
+                        input_field.schemaField
+                        and input_field.schemaField.fieldPath
+                        and input_field.schemaField.fieldPath.strip()
+                    ):
+                        valid_fields.append(input_field)
+                    else:
+                        invalid_count += 1
+
+                if invalid_count > 0:
+                    logger.debug(
+                        f"Filtered {invalid_count} invalid input field(s) with empty fieldPath for {wu.get_urn()}"
+                    )
+                    self.report.num_input_fields_filtered += invalid_count
+                    self.report.warning(
+                        title="Invalid input fields filtered",
+                        message="Input fields with empty fieldPath values were filtered out to prevent ingestion errors",
+                        context=f"Filtered {invalid_count} invalid input field(s) for {wu.get_urn()}",
+                    )
+
+                    # Update the aspect with only valid fields
+                    if valid_fields:
+                        input_fields_aspect.fields = valid_fields
+                    else:
+                        # If no valid fields remain, skip this workunit entirely
+                        logger.debug(
+                            f"All input fields were invalid for {wu.get_urn()}, skipping InputFieldsClass workunit"
+                        )
+                        # Don't yield this workunit
+                        continue
+
+            yield wu
+
+    def _remove_input_fields_aspect(self, wu: MetadataWorkUnit) -> MetadataWorkUnit:
+        """Remove InputFieldsClass aspect from a workunit."""
+        # For MCPs, we can simply not yield the aspect
+        # For MCEs, we need to remove it from the snapshot
+        if hasattr(wu.metadata, "aspect") and isinstance(
+            wu.metadata.aspect, InputFieldsClass
+        ):
+            # This is an MCP with InputFieldsClass, skip it
+            return wu
+
+        if hasattr(wu.metadata, "proposedSnapshot"):
+            snapshot = wu.metadata.proposedSnapshot
+            if hasattr(snapshot, "aspects"):
+                snapshot.aspects = [
+                    aspect
+                    for aspect in snapshot.aspects
+                    if not isinstance(aspect, InputFieldsClass)
+                ]
+
+        return wu

datahub/ingestion/api/decorators.py

@@ -1,12 +1,16 @@
+# So that SourceCapabilityModifier can be resolved at runtime
+from __future__ import annotations
+
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import Callable, Dict, Optional, Type
+from typing import Callable, Dict, List, Optional, Type
 
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.source import (
     Source,
     SourceCapability as SourceCapability,
 )
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 
 
 def config_class(config_cls: Type) -> Callable[[Type], Type]:
@@ -88,10 +92,14 @@ class CapabilitySetting:
     capability: SourceCapability
     description: str
     supported: bool
+    subtype_modifier: Optional[List[SourceCapabilityModifier]] = None
 
 
 def capability(
-    capability_name: SourceCapability, description: str, supported: bool = True
+    capability_name: SourceCapability,
+    description: str,
+    supported: bool = True,
+    subtype_modifier: Optional[List[SourceCapabilityModifier]] = None,
 ) -> Callable[[Type], Type]:
     """
     A decorator to mark a source as having a certain capability
@@ -104,6 +112,7 @@ def capability(
             for base in cls.__bases__
         ):
             cls.__capabilities = {}
+
             cls.get_capabilities = lambda: cls.__capabilities.values()
 
             # If the superclasses have capability annotations, copy those over.
@@ -113,7 +122,10 @@ def capability(
                     cls.__capabilities.update(base_caps)
 
         cls.__capabilities[capability_name] = CapabilitySetting(
-            capability=capability_name, description=description, supported=supported
+            capability=capability_name,
+            description=description,
+            supported=supported,
+            subtype_modifier=subtype_modifier,
         )
         return cls