acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/state/checkpoint.py

@@ -1,10 +1,8 @@
 import base64
 import bz2
-import contextlib
 import functools
 import json
 import logging
-import pickle
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import Callable, Generic, Optional, Type, TypeVar
@@ -70,7 +68,11 @@ class CheckpointStateBase(ConfigModel):
 
     @staticmethod
     def _to_bytes_utf8(model: ConfigModel) -> bytes:
-        return model.json(exclude={"version", "serde"}).encode("utf-8")
+        pydantic_json = model.model_dump_json(exclude={"version", "serde"})
+        # We decode and re-encode so that Python's default whitespace is included.
+        # This is purely to keep tests consistent as we migrate to pydantic v2,
+        # and can be removed once we're fully migrated.
+        return json.dumps(json.loads(pydantic_json)).encode("utf-8")
 
     @staticmethod
     def _to_bytes_base85_json(
@@ -117,10 +119,9 @@ class Checkpoint(Generic[StateType]):
                         checkpoint_aspect, state_class
                     )
                 elif checkpoint_aspect.state.serde == "base85":
-                    state_obj = Checkpoint._from_base85_bytes(
-                        checkpoint_aspect,
-                        functools.partial(bz2.decompress),
-                        state_class,
+                    raise ValueError(
+                        "The base85 encoding for stateful ingestion has been removed for security reasons. "
+                        "You may need to temporarily set `ignore_previous_checkpoint` to true to ignore the outdated checkpoint object."
                     )
                 elif checkpoint_aspect.state.serde == "base85-bz2-json":
                     state_obj = Checkpoint._from_base85_json_bytes(
@@ -164,28 +165,6 @@ class Checkpoint(Generic[StateType]):
         state_as_dict["serde"] = checkpoint_aspect.state.serde
         return state_class.parse_obj(state_as_dict)
 
-    @staticmethod
-    def _from_base85_bytes(
-        checkpoint_aspect: DatahubIngestionCheckpointClass,
-        decompressor: Callable[[bytes], bytes],
-        state_class: Type[StateType],
-    ) -> StateType:
-        state: StateType = pickle.loads(
-            decompressor(base64.b85decode(checkpoint_aspect.state.payload))  # type: ignore
-        )
-
-        with contextlib.suppress(Exception):
-            # When loading from pickle, the pydantic validators don't run.
-            # By re-serializing and re-parsing, we ensure that the state is valid.
-            # However, we also suppress any exceptions to make sure this doesn't blow up.
-            state = state_class.parse_obj(state.dict())
-
-        # Because the base85 method is deprecated in favor of base85-bz2-json,
-        # we will automatically switch the serde.
-        state.serde = "base85-bz2-json"
-
-        return state
-
     @staticmethod
     def _from_base85_json_bytes(
         checkpoint_aspect: DatahubIngestionCheckpointClass,

datahub/ingestion/source/state/entity_removal_state.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple, Type
 
 import pydantic
 
@@ -8,13 +8,16 @@ from datahub.utilities.checkpoint_state_util import CheckpointStateUtil
 from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.urns.urn import guess_entity_type
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 STATEFUL_INGESTION_IGNORED_ENTITY_TYPES = {
     "dataProcessInstance",
     "query",
 }
 
 
-def pydantic_state_migrator(mapping: Dict[str, str]) -> classmethod:
+def pydantic_state_migrator(mapping: Dict[str, str]) -> "V1RootValidator":
     # mapping would be something like:
     # {
     #    'encoded_view_urns': 'dataset',

datahub/ingestion/source/state/redundant_run_skip_handler.py

@@ -244,3 +244,24 @@ class RedundantUsageRunSkipHandler(RedundantRunSkipHandler):
             cur_state.begin_timestamp_millis = datetime_to_ts_millis(start_time)
             cur_state.end_timestamp_millis = datetime_to_ts_millis(end_time)
             cur_state.bucket_duration = bucket_duration
+
+
+class RedundantQueriesRunSkipHandler(RedundantRunSkipHandler):
+    """
+    Handler for stateful ingestion of queries v2 extraction.
+    Manages the time window for audit log extraction that combines
+    lineage, usage, operations, and queries.
+    """
+
+    def get_job_name_suffix(self):
+        return "_audit_window"
+
+    def update_state(
+        self, start_time: datetime, end_time: datetime, bucket_duration: BucketDuration
+    ) -> None:
+        cur_checkpoint = self.get_current_checkpoint()
+        if cur_checkpoint:
+            cur_state = cast(BaseTimeWindowCheckpointState, cur_checkpoint.state)
+            cur_state.begin_timestamp_millis = datetime_to_ts_millis(start_time)
+            cur_state.end_timestamp_millis = datetime_to_ts_millis(end_time)
+            cur_state.bucket_duration = bucket_duration

datahub/ingestion/source/state/stale_entity_removal_handler.py

@@ -45,7 +45,6 @@ class StatefulStaleMetadataRemovalConfig(StatefulIngestionConfig):
         description="Prevents large amount of soft deletes & the state from committing from accidental changes to the source configuration if the relative change percent in entities compared to the previous state is above the 'fail_safe_threshold'.",
         le=100.0,
         ge=0.0,
-        hidden_from_docs=True,
     )
 
 

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -10,6 +10,7 @@ from datahub.configuration.common import (
     ConfigModel,
     ConfigurationError,
     DynamicTypedConfig,
+    HiddenFromDocs,
 )
 from datahub.configuration.pydantic_migration_helpers import GenericModel
 from datahub.configuration.time_window_config import BaseTimeWindowConfig
@@ -55,25 +56,21 @@ class StatefulIngestionConfig(ConfigModel):
         description="Whether or not to enable stateful ingest. "
         "Default: True if a pipeline_name is set and either a datahub-rest sink or `datahub_api` is specified, otherwise False",
     )
-    max_checkpoint_state_size: pydantic.PositiveInt = Field(
+    max_checkpoint_state_size: HiddenFromDocs[pydantic.PositiveInt] = Field(
         default=2**24,  # 16 MB
         description="The maximum size of the checkpoint state in bytes. Default is 16MB",
-        hidden_from_docs=True,
     )
-    state_provider: Optional[DynamicTypedStateProviderConfig] = Field(
+    state_provider: HiddenFromDocs[Optional[DynamicTypedStateProviderConfig]] = Field(
         default=None,
         description="The ingestion state provider configuration.",
-        hidden_from_docs=True,
     )
-    ignore_old_state: bool = Field(
+    ignore_old_state: HiddenFromDocs[bool] = Field(
         default=False,
         description="If set to True, ignores the previous checkpoint state.",
-        hidden_from_docs=True,
     )
-    ignore_new_state: bool = Field(
+    ignore_new_state: HiddenFromDocs[bool] = Field(
         default=False,
         description="If set to True, ignores the current checkpoint state.",
-        hidden_from_docs=True,
     )
 
     @pydantic.root_validator(skip_on_failure=True)
@@ -104,7 +101,9 @@ class StatefulLineageConfigMixin(ConfigModel):
         default=True,
         description="Enable stateful lineage ingestion."
         " This will store lineage window timestamps after successful lineage ingestion. "
-        "and will not run lineage ingestion for same timestamps in subsequent run. ",
+        "and will not run lineage ingestion for same timestamps in subsequent run. "
+        "NOTE: This only works with use_queries_v2=False (legacy extraction path). "
+        "For queries v2, use enable_stateful_time_window instead.",
     )
 
     _store_last_lineage_extraction_timestamp = pydantic_renamed_field(
@@ -153,7 +152,9 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         default=True,
         description="Enable stateful lineage ingestion."
         " This will store usage window timestamps after successful usage ingestion. "
-        "and will not run usage ingestion for same timestamps in subsequent run. ",
+        "and will not run usage ingestion for same timestamps in subsequent run. "
+        "NOTE: This only works with use_queries_v2=False (legacy extraction path). "
+        "For queries v2, use enable_stateful_time_window instead.",
     )
 
     _store_last_usage_extraction_timestamp = pydantic_renamed_field(
@@ -172,6 +173,30 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         return values
 
 
+class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
+    enable_stateful_time_window: bool = Field(
+        default=False,
+        description="Enable stateful time window tracking."
+        " This will store the time window after successful extraction "
+        "and adjust the time window in subsequent runs to avoid reprocessing. "
+        "NOTE: This is ONLY applicable when using queries v2 (use_queries_v2=True). "
+        "This replaces enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion "
+        "for the queries v2 extraction path, since queries v2 extracts lineage, usage, operations, "
+        "and queries together from a single audit log and uses a unified time window.",
+    )
+
+    @root_validator(skip_on_failure=True)
+    def time_window_stateful_option_validator(cls, values: Dict) -> Dict:
+        sti = values.get("stateful_ingestion")
+        if not sti or not sti.enabled:
+            if values.get("enable_stateful_time_window"):
+                logger.warning(
+                    "Stateful ingestion is disabled, disabling enable_stateful_time_window config option as well"
+                )
+                values["enable_stateful_time_window"] = False
+        return values
+
+
 @dataclass
 class StatefulIngestionReport(SourceReport):
     pass
@@ -179,7 +204,7 @@ class StatefulIngestionReport(SourceReport):
 
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class StatefulIngestionSourceBase(Source):

datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py

@@ -10,7 +10,8 @@ from datahub.ingestion.api.ingestion_job_checkpointing_provider_base import (
     IngestionCheckpointingProviderConfig,
     JobId,
 )
-from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.metadata.schema_classes import DatahubIngestionCheckpointClass
 
 logger = logging.getLogger(__name__)