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

datahub/ingestion/source/sql/trino.py

@@ -18,6 +18,7 @@ from sqlalchemy.types import TypeEngine
 from trino.sqlalchemy import datatype
 from trino.sqlalchemy.dialect import TrinoDialect
 
+from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
@@ -36,6 +37,7 @@ from datahub.ingestion.api.decorators import (
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.extractor import schema_util
 from datahub.ingestion.source.common.data_reader import DataReader
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -221,7 +223,7 @@ class ConnectorDetail(PlatformInstanceConfigMixin, EnvConfigMixin):
 
 class TrinoConfig(BasicSQLAlchemyConfig):
     # defaults
-    scheme: str = Field(default="trino", description="", hidden_from_docs=True)
+    scheme: HiddenFromDocs[str] = Field(default="trino")
     database: str = Field(description="database (catalog)")
 
     catalog_to_connector_details: Dict[str, ConnectorDetail] = Field(
@@ -249,6 +251,14 @@ class TrinoConfig(BasicSQLAlchemyConfig):
 @support_status(SupportStatus.CERTIFIED)
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Extract table-level lineage",
+    subtype_modifier=[
+        SourceCapabilityModifier.TABLE,
+        SourceCapabilityModifier.VIEW,
+    ],
+)
 class TrinoSource(SQLAlchemySource):
     """
 

datahub/ingestion/source/sql/two_tier_sql_source.py

@@ -7,7 +7,7 @@ from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine import URL
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -27,11 +27,10 @@ class TwoTierSQLAlchemyConfig(BasicSQLAlchemyConfig):
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for databases to filter in ingestion.",
     )
-    schema_pattern: AllowDenyPattern = Field(
+    schema_pattern: HiddenFromDocs[AllowDenyPattern] = Field(
         # The superclass contains a `schema_pattern` field, so we need this here
         # to override the documentation.
         default=AllowDenyPattern.allow_all(),
-        hidden_from_docs=True,
         description="Deprecated in favour of database_pattern.",
     )
 

datahub/ingestion/source/sql/vertica.py

@@ -4,7 +4,8 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import pydantic
-from pydantic.class_validators import validator
+import pytest
+from pydantic import validator
 from vertica_sqlalchemy_dialect.base import VerticaInspector
 
 from datahub.configuration.common import AllowDenyPattern
@@ -25,6 +26,10 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_reader import DataReader
+from datahub.ingestion.source.common.subtypes import (
+    DatasetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -41,7 +46,6 @@ from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
-    ChangeTypeClass,
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
     SubTypesClass,
@@ -52,6 +56,8 @@ from datahub.utilities import config_clean
 
 if TYPE_CHECKING:
     from datahub.ingestion.source.ge_data_profiler import GEProfilerRequest
+
+pytestmark = pytest.mark.integration_batch_4
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -113,10 +119,14 @@ class VerticaConfig(BasicSQLAlchemyConfig):
 @capability(
     SourceCapability.LINEAGE_COARSE,
     "Enabled by default, can be disabled via configuration `include_view_lineage` and `include_projection_lineage`",
+    subtype_modifier=[
+        SourceCapabilityModifier.VIEW,
+        SourceCapabilityModifier.PROJECTIONS,
+    ],
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class VerticaSource(SQLAlchemySource):
@@ -493,11 +503,8 @@ class VerticaSource(SQLAlchemySource):
         if dpi_aspect:
             yield dpi_aspect
         yield MetadataChangeProposalWrapper(
-            entityType="dataset",
-            changeType=ChangeTypeClass.UPSERT,
             entityUrn=dataset_urn,
-            aspectName="subTypes",
-            aspect=SubTypesClass(typeNames=["Projections"]),
+            aspect=SubTypesClass(typeNames=[DatasetSubTypes.PROJECTIONS]),
         ).as_workunit()
 
         if self.config.domain:

datahub/ingestion/source/sql_queries.py

@@ -2,21 +2,22 @@ import json
 import logging
 import os
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime
 from functools import partial
-from typing import Iterable, List, Optional, Set
+from typing import ClassVar, Iterable, List, Optional, Union
 
-from pydantic import Field
+from pydantic import BaseModel, Field, validator
 
+from datahub.configuration.common import HiddenFromDocs
+from datahub.configuration.datetimes import parse_user_datetime
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
 )
 from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
-    make_user_urn,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -25,6 +26,10 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
+from datahub.ingestion.api.incremental_lineage_helper import (
+    IncrementalLineageConfigMixin,
+    auto_incremental_lineage,
+)
 from datahub.ingestion.api.source import (
     MetadataWorkUnitProcessor,
     Source,
@@ -35,13 +40,21 @@ from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
+from datahub.metadata.urns import CorpUserUrn, DatasetUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    KnownQueryLineageInfo,
+    ObservedQuery,
+    SqlAggregatorReport,
+    SqlParsingAggregator,
+)
 
 logger = logging.getLogger(__name__)
 
 
-class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
+class SqlQueriesSourceConfig(
+    PlatformInstanceConfigMixin, EnvConfigMixin, IncrementalLineageConfigMixin
+):
     query_file: str = Field(description="Path to file to ingest")
 
     platform: str = Field(
@@ -53,45 +66,34 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
         default=BaseUsageConfig(),
     )
 
-    use_schema_resolver: bool = Field(
+    use_schema_resolver: HiddenFromDocs[bool] = Field(
+        True,
         description="Read SchemaMetadata aspects from DataHub to aid in SQL parsing. Turn off only for testing.",
-        default=True,
-        hidden_from_docs=True,
     )
     default_db: Optional[str] = Field(
+        None,
         description="The default database to use for unqualified table names",
-        default=None,
     )
     default_schema: Optional[str] = Field(
+        None,
         description="The default schema to use for unqualified table names",
-        default=None,
     )
-    default_dialect: Optional[str] = Field(
+    override_dialect: Optional[str] = Field(
+        None,
         description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
-        default=None,
     )
 
 
+@dataclass
 class SqlQueriesSourceReport(SourceReport):
-    num_queries_parsed: int = 0
-    num_table_parse_failures: int = 0
-    num_column_parse_failures: int = 0
-
-    def compute_stats(self) -> None:
-        super().compute_stats()
-        self.table_failure_rate = (
-            f"{self.num_table_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
-        self.column_failure_rate = (
-            f"{self.num_column_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
+    num_entries_processed: int = 0
+    num_entries_failed: int = 0
+    num_queries_aggregator_failures: int = 0
 
+    sql_aggregator: Optional[SqlAggregatorReport] = None
 
-@platform_name("SQL Queries")
+
+@platform_name("SQL Queries", id="sql-queries")
 @config_class(SqlQueriesSourceConfig)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.LINEAGE_COARSE, "Parsed from SQL queries")
@@ -107,15 +109,25 @@ class SqlQueriesSource(Source):
     - user (optional): string - The user who ran the query.
     This user value will be directly converted into a DataHub user urn.
     - operation_type (optional): string - Platform-specific operation type, used if the operation type can't be parsed.
+    - session_id (optional): string - Session identifier for temporary table resolution across queries.
     - downstream_tables (optional): string[] - Fallback list of tables that the query writes to,
      used if the query can't be parsed.
     - upstream_tables (optional): string[] - Fallback list of tables the query reads from,
      used if the query can't be parsed.
+
+    ### Incremental Lineage
+    When `incremental_lineage` is enabled, this source will emit lineage as patches rather than full overwrites.
+    This allows you to add lineage edges without removing existing ones, which is useful for:
+    - Gradually building up lineage from multiple sources
+    - Preserving manually curated lineage
+    - Avoiding conflicts when multiple ingestion processes target the same datasets
+
+    Note: Incremental lineage only applies to UpstreamLineage aspects. Other aspects like queries and usage
+    statistics will still be emitted normally.
     """
 
-    urns: Optional[Set[str]]
-    schema_resolver: SchemaResolver
-    builder: SqlParsingBuilder
+    schema_resolver: Optional[SchemaResolver]
+    aggregator: SqlParsingAggregator
 
     def __init__(self, ctx: PipelineContext, config: SqlQueriesSourceConfig):
         if not ctx.graph:
@@ -128,22 +140,36 @@ class SqlQueriesSource(Source):
         self.config = config
         self.report = SqlQueriesSourceReport()
 
-        self.builder = SqlParsingBuilder(usage_config=self.config.usage)
-
         if self.config.use_schema_resolver:
+            # TODO: `initialize_schema_resolver_from_datahub` does a  bulk initialization by fetching all schemas
+            # for the given platform, platform instance, and env. Instead this should be configurable:
+            # bulk initialization vs lazy on-demand schema fetching.
             self.schema_resolver = self.graph.initialize_schema_resolver_from_datahub(
                 platform=self.config.platform,
                 platform_instance=self.config.platform_instance,
                 env=self.config.env,
             )
-            self.urns = self.schema_resolver.get_urns()
         else:
-            self.schema_resolver = self.graph._make_schema_resolver(
-                platform=self.config.platform,
-                platform_instance=self.config.platform_instance,
-                env=self.config.env,
-            )
-            self.urns = None
+            self.schema_resolver = None
+
+        self.aggregator = SqlParsingAggregator(
+            platform=self.config.platform,
+            platform_instance=self.config.platform_instance,
+            env=self.config.env,
+            schema_resolver=self.schema_resolver,
+            eager_graph_load=False,
+            generate_lineage=True,  # TODO: make this configurable
+            generate_queries=True,  # TODO: make this configurable
+            generate_query_subject_fields=True,  # TODO: make this configurable
+            generate_query_usage_statistics=True,  # This enables publishing SELECT query entities, otherwise only mutation queries are published
+            generate_usage_statistics=True,
+            generate_operations=True,  # TODO: make this configurable
+            usage_config=self.config.usage,
+            is_temp_table=None,
+            is_allowed_table=None,
+            format_queries=False,
+        )
+        self.report.sql_aggregator = self.aggregator.report
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SqlQueriesSource":
@@ -154,100 +180,172 @@ class SqlQueriesSource(Source):
         return self.report
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
-        return [partial(auto_workunit_reporter, self.get_report())]
+        return [
+            partial(auto_workunit_reporter, self.get_report()),
+            partial(
+                auto_incremental_lineage,
+                self.config.incremental_lineage,
+            ),
+        ]
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         logger.info(f"Parsing queries from {os.path.basename(self.config.query_file)}")
+
+        with self.report.new_stage("Collecting queries from file"):
+            queries = list(self._parse_query_file())
+            logger.info(f"Collected {len(queries)} queries for processing")
+
+        with self.report.new_stage("Processing queries through SQL parsing aggregator"):
+            for query_entry in queries:
+                self._add_query_to_aggregator(query_entry)
+
+        with self.report.new_stage("Generating metadata work units"):
+            logger.info("Generating workunits from SQL parsing aggregator")
+            yield from self.aggregator.gen_metadata()
+
+    def _parse_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse the query file and yield QueryEntry objects."""
         with open(self.config.query_file) as f:
             for line in f:
                 try:
                     query_dict = json.loads(line, strict=False)
                     entry = QueryEntry.create(query_dict, config=self.config)
-                    yield from self._process_query(entry)
+                    self.report.num_entries_processed += 1
+                    if self.report.num_entries_processed % 1000 == 0:
+                        logger.info(
+                            f"Processed {self.report.num_entries_processed} query entries"
+                        )
+                    yield entry
                 except Exception as e:
-                    logger.warning("Error processing query", exc_info=True)
-                    self.report.report_warning("process-query", str(e))
-
-        logger.info("Generating workunits")
-        yield from self.builder.gen_workunits()
-
-    def _process_query(self, entry: "QueryEntry") -> Iterable[MetadataWorkUnit]:
-        self.report.num_queries_parsed += 1
-        if self.report.num_queries_parsed % 1000 == 0:
-            logger.info(f"Parsed {self.report.num_queries_parsed} queries")
+                    self.report.num_entries_failed += 1
+                    self.report.warning(
+                        title="Error processing query",
+                        message="Query skipped due to parsing error",
+                        context=line.strip(),
+                        exc=e,
+                    )
 
-        result = sqlglot_lineage(
-            sql=entry.query,
-            schema_resolver=self.schema_resolver,
-            default_db=self.config.default_db,
-            default_schema=self.config.default_schema,
-            default_dialect=self.config.default_dialect,
-        )
-        if result.debug_info.table_error:
-            logger.info(f"Error parsing table lineage, {result.debug_info.table_error}")
-            self.report.num_table_parse_failures += 1
-            for downstream_urn in set(entry.downstream_tables):
-                self.builder.add_lineage(
-                    downstream_urn=downstream_urn,
-                    upstream_urns=entry.upstream_tables,
-                    timestamp=entry.timestamp,
-                    user=entry.user,
+    def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
+        """Add a query to the SQL parsing aggregator."""
+        try:
+            # If we have both upstream and downstream tables, use explicit lineage
+            if query_entry.upstream_tables and query_entry.downstream_tables:
+                logger.debug("Using explicit lineage from query file")
+                for downstream_table in query_entry.downstream_tables:
+                    known_lineage = KnownQueryLineageInfo(
+                        query_text=query_entry.query,
+                        downstream=str(downstream_table),
+                        upstreams=[str(urn) for urn in query_entry.upstream_tables],
+                        timestamp=query_entry.timestamp,
+                        session_id=query_entry.session_id,
+                    )
+                    self.aggregator.add_known_query_lineage(known_lineage)
+            else:
+                # Warn if only partial lineage information is provided
+                # XOR: true if exactly one of upstream_tables or downstream_tables is provided
+                if bool(query_entry.upstream_tables) ^ bool(
+                    query_entry.downstream_tables
+                ):
+                    query_preview = (
+                        query_entry.query[:150] + "..."
+                        if len(query_entry.query) > 150
+                        else query_entry.query
+                    )
+                    missing_upstream = (
+                        "Missing upstream. " if not query_entry.upstream_tables else ""
+                    )
+                    missing_downstream = (
+                        "Missing downstream. "
+                        if not query_entry.downstream_tables
+                        else ""
+                    )
+                    logger.info(
+                        f"Only partial lineage information provided, falling back to SQL parsing for complete lineage detection. {missing_upstream}{missing_downstream}Query: {query_preview}"
+                    )
+                # No explicit lineage, rely on parsing
+                observed_query = ObservedQuery(
+                    query=query_entry.query,
+                    timestamp=query_entry.timestamp,
+                    user=query_entry.user,
+                    session_id=query_entry.session_id,
+                    default_db=self.config.default_db,
+                    default_schema=self.config.default_schema,
+                    override_dialect=self.config.override_dialect,
                 )
-            return
-        elif result.debug_info.column_error:
-            logger.debug(
-                f"Error parsing column lineage, {result.debug_info.column_error}"
+                self.aggregator.add_observed_query(observed_query)
+
+        except Exception as e:
+            self.report.num_queries_aggregator_failures += 1
+            self.report.warning(
+                title="Error adding query to aggregator",
+                message="Query skipped due to failure when adding query to SQL parsing aggregator",
+                context=query_entry.query,
+                exc=e,
             )
-            self.report.num_column_parse_failures += 1
-
-        yield from self.builder.process_sql_parsing_result(
-            result,
-            query=entry.query,
-            query_timestamp=entry.timestamp,
-            user=entry.user,
-            custom_operation_type=entry.operation_type,
-            include_urns=self.urns,
-        )
 
 
-@dataclass
-class QueryEntry:
+class QueryEntry(BaseModel):
     query: str
-    timestamp: Optional[datetime]
-    user: Optional[str]
-    operation_type: Optional[str]
-    downstream_tables: List[str]
-    upstream_tables: List[str]
+    timestamp: Optional[datetime] = None
+    user: Optional[CorpUserUrn] = None
+    operation_type: Optional[str] = None
+    downstream_tables: List[DatasetUrn] = Field(default_factory=list)
+    upstream_tables: List[DatasetUrn] = Field(default_factory=list)
+    session_id: Optional[str] = None
+
+    # Validation context for URN creation
+    _validation_context: ClassVar[Optional[SqlQueriesSourceConfig]] = None
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @validator("timestamp", pre=True)
+    def parse_timestamp(cls, v):
+        return None if v is None else parse_user_datetime(str(v))
+
+    @validator("user", pre=True)
+    def parse_user(cls, v):
+        if v is None:
+            return None
+
+        return v if isinstance(v, CorpUserUrn) else CorpUserUrn(v)
+
+    @validator("downstream_tables", "upstream_tables", pre=True)
+    def parse_tables(cls, v):
+        if not v:
+            return []
+
+        result = []
+        for item in v:
+            if isinstance(item, DatasetUrn):
+                result.append(item)
+            elif isinstance(item, str):
+                # Skip empty/whitespace-only strings
+                if item and item.strip():
+                    # Convert to URN using validation context
+                    assert cls._validation_context, (
+                        "Validation context must be set for URN creation"
+                    )
+                    urn_string = make_dataset_urn_with_platform_instance(
+                        name=item,
+                        platform=cls._validation_context.platform,
+                        platform_instance=cls._validation_context.platform_instance,
+                        env=cls._validation_context.env,
+                    )
+                    result.append(DatasetUrn.from_string(urn_string))
+
+        return result
 
     @classmethod
     def create(
         cls, entry_dict: dict, *, config: SqlQueriesSourceConfig
     ) -> "QueryEntry":
-        return cls(
-            query=entry_dict["query"],
-            timestamp=(
-                datetime.fromtimestamp(entry_dict["timestamp"], tz=timezone.utc)
-                if "timestamp" in entry_dict
-                else None
-            ),
-            user=make_user_urn(entry_dict["user"]) if "user" in entry_dict else None,
-            operation_type=entry_dict.get("operation_type"),
-            downstream_tables=[
-                make_dataset_urn_with_platform_instance(
-                    name=table,
-                    platform=config.platform,
-                    platform_instance=config.platform_instance,
-                    env=config.env,
-                )
-                for table in entry_dict.get("downstream_tables", [])
-            ],
-            upstream_tables=[
-                make_dataset_urn_with_platform_instance(
-                    name=table,
-                    platform=config.platform,
-                    platform_instance=config.platform_instance,
-                    env=config.env,
-                )
-                for table in entry_dict.get("upstream_tables", [])
-            ],
-        )
+        """Create QueryEntry from dict with config context."""
+        # Set validation context for URN creation
+        cls._validation_context = config
+        try:
+            return cls.parse_obj(entry_dict)
+        finally:
+            cls._validation_context = None

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',