acryl-datahub 1.2.0.6rc1__py3-none-any.whl → 1.2.0.7__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_schema.py

@@ -3,7 +3,7 @@ import os
 from collections import defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, Callable, Dict, Iterable, List, MutableMapping, Optional
+from typing import Any, Callable, Dict, Iterable, List, MutableMapping, Optional, Tuple
 
 from datahub.ingestion.api.report import SupportsAsObj
 from datahub.ingestion.source.common.subtypes import DatasetSubTypes
@@ -239,12 +239,16 @@ class _SnowflakeTagCache:
 
 class SnowflakeDataDictionary(SupportsAsObj):
     def __init__(
-        self, connection: SnowflakeConnection, report: SnowflakeV2Report
+        self,
+        connection: SnowflakeConnection,
+        report: SnowflakeV2Report,
+        fetch_views_from_information_schema: bool = False,
     ) -> None:
         self.connection = connection
         self.report = report
+        self._fetch_views_from_information_schema = fetch_views_from_information_schema
 
-    def as_obj(self) -> Dict[str, Dict[str, int]]:
+    def as_obj(self) -> Dict[str, Any]:
         # TODO: Move this into a proper report type that gets computed.
 
         # Reports how many times we reset in-memory `functools.lru_cache` caches of data,
@@ -260,7 +264,9 @@ class SnowflakeDataDictionary(SupportsAsObj):
             self.get_fk_constraints_for_schema,
         ]
 
-        report = {}
+        report: Dict[str, Any] = {
+            "fetch_views_from_information_schema": self._fetch_views_from_information_schema,
+        }
         for func in lru_cache_functions:
             report[func.__name__] = func.cache_info()._asdict()  # type: ignore
         return report
@@ -430,7 +436,17 @@ class SnowflakeDataDictionary(SupportsAsObj):
         return tables
 
     @serialized_lru_cache(maxsize=1)
-    def get_views_for_database(self, db_name: str) -> Dict[str, List[SnowflakeView]]:
+    def get_views_for_database(
+        self, db_name: str
+    ) -> Optional[Dict[str, List[SnowflakeView]]]:
+        if self._fetch_views_from_information_schema:
+            return self._get_views_for_database_using_information_schema(db_name)
+        else:
+            return self._get_views_for_database_using_show(db_name)
+
+    def _get_views_for_database_using_show(
+        self, db_name: str
+    ) -> Dict[str, List[SnowflakeView]]:
         page_limit = SHOW_COMMAND_MAX_PAGE_SIZE
 
         views: Dict[str, List[SnowflakeView]] = {}
@@ -461,10 +477,9 @@ class SnowflakeDataDictionary(SupportsAsObj):
                     SnowflakeView(
                         name=view_name,
                         created=view["created_on"],
-                        # last_altered=table["last_altered"],
                         comment=view["comment"],
                         view_definition=view["text"],
-                        last_altered=view["created_on"],
+                        last_altered=view["created_on"],  # TODO: This is not correct.
                         materialized=(
                             view.get("is_materialized", "false").lower() == "true"
                         ),
@@ -479,6 +494,163 @@ class SnowflakeDataDictionary(SupportsAsObj):
                 )
                 view_pagination_marker = view_name
 
+        # Because this is in a cached function, this will only log once per database.
+        view_counts = {schema_name: len(views[schema_name]) for schema_name in views}
+        logger.info(
+            f"Finished fetching views in {db_name}; counts by schema {view_counts}"
+        )
+        return views
+
+    def _map_view(self, db_name: str, row: Dict[str, Any]) -> Tuple[str, SnowflakeView]:
+        schema_name = row["VIEW_SCHEMA"]
+        view_definition = row.get("VIEW_DEFINITION")
+        fragment_view_definition = (
+            view_definition[:50].strip() if view_definition else None
+        )
+        logger.info(
+            f"Mapping view {db_name}.{schema_name}.{row['VIEW_NAME']} with view definition: {fragment_view_definition}..."
+        )
+
+        return schema_name, SnowflakeView(
+            name=row["VIEW_NAME"],
+            created=row["CREATED"],
+            comment=row["COMMENT"],
+            view_definition=view_definition,
+            last_altered=row["LAST_ALTERED"],
+            is_secure=(row.get("IS_SECURE", "false").lower() == "true"),
+            # TODO: This doesn't work for materialized views.
+            materialized=False,
+        )
+
+    def _maybe_populate_empty_view_definitions(
+        self,
+        db_name: str,
+        schema_name: str,
+        views_with_empty_definition: List[SnowflakeView],
+    ) -> List[SnowflakeView]:
+        if not views_with_empty_definition:
+            return []
+
+        view_names = [view.name for view in views_with_empty_definition]
+        batches = [
+            batch[0]
+            for batch in build_prefix_batches(
+                view_names, max_batch_size=1000, max_groups_in_batch=1
+            )
+            if batch
+            # Skip empty batch if so, also max_groups_in_batch=1 makes it safe to access batch[0]
+        ]
+
+        view_map: Dict[str, SnowflakeView] = {
+            view.name: view for view in views_with_empty_definition
+        }
+        views_found_count = 0
+
+        logger.info(
+            f"Fetching definitions for {len(view_map)} views in {db_name}.{schema_name} "
+            f"using batched 'SHOW VIEWS ... LIKE ...' queries. Found {len(batches)} batch(es)."
+        )
+
+        for batch_index, prefix_group in enumerate(batches):
+            query = f'SHOW VIEWS LIKE \'{prefix_group.prefix}%\' IN SCHEMA "{db_name}"."{schema_name}"'
+            logger.info(f"Processing batch {batch_index + 1}/{len(batches)}: {query}")
+
+            try:
+                cur = self.connection.query(query)
+                for row in cur:
+                    view_name = row["name"]
+                    if view_name in view_map:
+                        view_definition = row.get("text")
+                        if view_definition:  # Ensure definition is not None or empty
+                            view_map[view_name].view_definition = view_definition
+                            views_found_count += 1
+                            logger.debug(
+                                f"Fetched view definition for {db_name}.{schema_name}.{view_name}"
+                            )
+                            # If all targeted views are found, we could theoretically break early,
+                            # but SHOW VIEWS doesn't guarantee order, so we must process all results.
+                        else:
+                            logger.warning(
+                                f"'text' field missing or empty in SHOW VIEWS result for {db_name}.{schema_name}.{view_name}"
+                            )
+
+            except Exception as e:
+                logger.error(
+                    f"Failed to execute query for batch {batch_index + 1} ('{query}') for {db_name}.{schema_name} or process its results.",
+                    exc_info=e,
+                )
+                # Returning the original list; some views might still be missing definitions.
+                # This also means subsequent batches for this schema (in this call) are skipped.
+                return views_with_empty_definition
+
+        logger.info(
+            f"Finished processing 'SHOW VIEWS' batches for {db_name}.{schema_name}. "
+            f"Fetched definitions for {views_found_count} out of {len(view_map)} targeted views."
+        )
+
+        if views_found_count < len(view_map):
+            missing_count = len(view_map) - views_found_count
+            logger.warning(
+                f"Could not fetch definitions for {missing_count} views in {db_name}.{schema_name} after processing all batches."
+            )
+        # The SnowflakeView objects in the original list were modified in place via view_map
+        return views_with_empty_definition
+
+    def _get_views_for_database_using_information_schema(
+        self, db_name: str
+    ) -> Optional[Dict[str, List[SnowflakeView]]]:
+        try:
+            cur = self.connection.query(
+                SnowflakeQuery.get_views_for_database(db_name),
+            )
+        except Exception as e:
+            logger.debug(f"Failed to get all views for database {db_name}", exc_info=e)
+            # Error - Information schema query returned too much data. Please repeat query with more selective predicates.
+            return None
+
+        views: Dict[str, List[SnowflakeView]] = {}
+        views_with_empty_definition: Dict[str, List[SnowflakeView]] = {}
+
+        for row in cur:
+            schema_name, view = self._map_view(db_name, row)
+            if view.view_definition is None or view.view_definition == "":
+                views_with_empty_definition.setdefault(schema_name, []).append(view)
+            else:
+                views.setdefault(schema_name, []).append(view)
+
+        for schema_name, empty_views in views_with_empty_definition.items():
+            updated_views = self._maybe_populate_empty_view_definitions(
+                db_name, schema_name, empty_views
+            )
+            views.setdefault(schema_name, []).extend(updated_views)
+
+        return views
+
+    def get_views_for_schema_using_information_schema(
+        self, *, schema_name: str, db_name: str
+    ) -> List[SnowflakeView]:
+        cur = self.connection.query(
+            SnowflakeQuery.get_views_for_schema(
+                db_name=db_name, schema_name=schema_name
+            ),
+        )
+
+        views: List[SnowflakeView] = []
+        views_with_empty_definition: List[SnowflakeView] = []
+
+        for row in cur:
+            schema_name, view = self._map_view(db_name, row)
+            if view.view_definition is None or view.view_definition == "":
+                views_with_empty_definition.append(view)
+            else:
+                views.append(view)
+
+        if views_with_empty_definition:
+            updated_empty_views = self._maybe_populate_empty_view_definitions(
+                db_name, schema_name, views_with_empty_definition
+            )
+            views.extend(updated_empty_views)
+
         return views
 
     @serialized_lru_cache(maxsize=SCHEMA_PARALLELISM)

datahub/ingestion/source/snowflake/snowflake_schema_gen.py

@@ -166,8 +166,8 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
 
     def __init__(
         self,
-        config: SnowflakeV2Config,
-        report: SnowflakeV2Report,
+        config: SnowflakeV2Config,  # FIXME: SnowflakeSummary is passing here SnowflakeSummaryConfig
+        report: SnowflakeV2Report,  # FIXME: SnowflakeSummary is passing here SnowflakeSummaryReport
         connection: SnowflakeConnection,
         filters: SnowflakeFilter,
         identifiers: SnowflakeIdentifierBuilder,
@@ -175,6 +175,7 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         profiler: Optional[SnowflakeProfiler],
         aggregator: Optional[SqlParsingAggregator],
         snowsight_url_builder: Optional[SnowsightUrlBuilder],
+        fetch_views_from_information_schema: bool = False,
     ) -> None:
         self.config: SnowflakeV2Config = config
         self.report: SnowflakeV2Report = report
@@ -183,7 +184,9 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         self.identifiers: SnowflakeIdentifierBuilder = identifiers
 
         self.data_dictionary: SnowflakeDataDictionary = SnowflakeDataDictionary(
-            connection=self.connection, report=self.report
+            connection=self.connection,
+            report=self.report,
+            fetch_views_from_information_schema=fetch_views_from_information_schema,
         )
         self.report.data_dictionary_cache = self.data_dictionary
 
@@ -438,13 +441,16 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
             tables = self.fetch_tables_for_schema(
                 snowflake_schema, db_name, schema_name
             )
+        if self.config.include_views:
+            views = self.fetch_views_for_schema(snowflake_schema, db_name, schema_name)
+
+        if self.config.include_tables:
             db_tables[schema_name] = tables
             yield from self._process_tables(
                 tables, snowflake_schema, db_name, schema_name
             )
 
         if self.config.include_views:
-            views = self.fetch_views_for_schema(snowflake_schema, db_name, schema_name)
             yield from self._process_views(
                 views, snowflake_schema, db_name, schema_name
             )
@@ -1241,7 +1247,10 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
         # falling back to get tables for schema
         if tables is None:
             self.report.num_get_tables_for_schema_queries += 1
-            return self.data_dictionary.get_tables_for_schema(schema_name, db_name)
+            return self.data_dictionary.get_tables_for_schema(
+                db_name=db_name,
+                schema_name=schema_name,
+            )
 
         # Some schema may not have any table
         return tables.get(schema_name, [])
@@ -1251,8 +1260,17 @@ class SnowflakeSchemaGenerator(SnowflakeStructuredReportMixin):
     ) -> List[SnowflakeView]:
         views = self.data_dictionary.get_views_for_database(db_name)
 
-        # Some schema may not have any table
-        return views.get(schema_name, [])
+        if views is not None:
+            # Some schemas may not have any views
+            return views.get(schema_name, [])
+
+        # Usually this fails when there are too many views in the schema.
+        # Fall back to per-schema queries.
+        self.report.num_get_views_for_schema_queries += 1
+        return self.data_dictionary.get_views_for_schema_using_information_schema(
+            db_name=db_name,
+            schema_name=schema_name,
+        )
 
     def get_columns_for_table(
         self, table_name: str, snowflake_schema: SnowflakeSchema, db_name: str

datahub/ingestion/source/snowflake/snowflake_summary.py

@@ -86,6 +86,7 @@ class SnowflakeSummarySource(Source):
                 filter_config=self.config,
                 structured_reporter=self.report,
             ),
+            fetch_views_from_information_schema=False,  # we haven't enabled this config for SnowflakeSummarySource
         )
 
         # Databases.

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -9,6 +9,7 @@ from datahub.emitter.mce_builder import (
 from datahub.emitter.mcp_builder import DatabaseKey, SchemaKey
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.snowflake.constants import (
+    DEFAULT_SNOWFLAKE_DOMAIN,
     SNOWFLAKE_REGION_CLOUD_REGION_MAPPING,
     SnowflakeCloudProvider,
     SnowflakeObjectDomain,
@@ -34,16 +35,21 @@ class SnowsightUrlBuilder:
         "us-east-1",
         "eu-west-1",
         "eu-central-1",
-        "ap-southeast-1",
         "ap-southeast-2",
     ]
 
     snowsight_base_url: str
 
-    def __init__(self, account_locator: str, region: str, privatelink: bool = False):
+    def __init__(
+        self,
+        account_locator: str,
+        region: str,
+        privatelink: bool = False,
+        snowflake_domain: str = DEFAULT_SNOWFLAKE_DOMAIN,
+    ):
         cloud, cloud_region_id = self.get_cloud_region_from_snowflake_region_id(region)
         self.snowsight_base_url = self.create_snowsight_base_url(
-            account_locator, cloud_region_id, cloud, privatelink
+            account_locator, cloud_region_id, cloud, privatelink, snowflake_domain
         )
 
     @staticmethod
@@ -52,6 +58,7 @@ class SnowsightUrlBuilder:
         cloud_region_id: str,
         cloud: str,
         privatelink: bool = False,
+        snowflake_domain: str = DEFAULT_SNOWFLAKE_DOMAIN,
     ) -> str:
         if cloud:
             url_cloud_provider_suffix = f".{cloud}"
@@ -67,9 +74,15 @@ class SnowsightUrlBuilder:
             else:
                 url_cloud_provider_suffix = f".{cloud}"
         if privatelink:
-            url = f"https://app.{account_locator}.{cloud_region_id}.privatelink.snowflakecomputing.com/"
+            url = f"https://app.{account_locator}.{cloud_region_id}.privatelink.{snowflake_domain}/"
         else:
-            url = f"https://app.snowflake.com/{cloud_region_id}{url_cloud_provider_suffix}/{account_locator}/"
+            # Standard Snowsight URL format - works for most regions
+            # China region may use app.snowflake.cn instead of app.snowflake.com. This is not documented, just
+            # guessing Based on existence of snowflake.cn domain (https://domainindex.com/domains/snowflake.cn)
+            if snowflake_domain == "snowflakecomputing.cn":
+                url = f"https://app.snowflake.cn/{cloud_region_id}{url_cloud_provider_suffix}/{account_locator}/"
+            else:
+                url = f"https://app.snowflake.com/{cloud_region_id}{url_cloud_provider_suffix}/{account_locator}/"
         return url
 
     @staticmethod

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -172,7 +172,9 @@ class SnowflakeV2Source(
 
         # For database, schema, tables, views, etc
         self.data_dictionary = SnowflakeDataDictionary(
-            connection=self.connection, report=self.report
+            connection=self.connection,
+            report=self.report,
+            fetch_views_from_information_schema=self.config.fetch_views_from_information_schema,
         )
         self.lineage_extractor: Optional[SnowflakeLineageExtractor] = None
 
@@ -197,6 +199,7 @@ class SnowflakeV2Source(
                 ),
                 generate_usage_statistics=False,
                 generate_operations=False,
+                generate_queries=self.config.include_queries,
                 format_queries=self.config.format_sql_queries,
                 is_temp_table=self._is_temp_table,
                 is_allowed_table=self._is_allowed_table,
@@ -528,6 +531,7 @@ class SnowflakeV2Source(
             snowsight_url_builder=snowsight_url_builder,
             filters=self.filters,
             identifiers=self.identifiers,
+            fetch_views_from_information_schema=self.config.fetch_views_from_information_schema,
         )
 
         with self.report.new_stage(f"*: {METADATA_EXTRACTION}"):
@@ -747,6 +751,7 @@ class SnowflakeV2Source(
                 # For privatelink, account identifier ends with .privatelink
                 # See https://docs.snowflake.com/en/user-guide/organizations-connect.html#private-connectivity-urls
                 privatelink=self.config.account_id.endswith(".privatelink"),
+                snowflake_domain=self.config.snowflake_domain,
             )
 
         except Exception as e:

datahub/ingestion/source/sql/hive_metastore.py

@@ -174,6 +174,7 @@ class HiveMetastore(BasicSQLAlchemyConfig):
     "Enabled by default",
     subtype_modifier=[
         SourceCapabilityModifier.CATALOG,
+        SourceCapabilityModifier.SCHEMA,
     ],
 )
 class HiveMetastoreSource(SQLAlchemySource):

datahub/ingestion/source/sql/mssql/job_models.py

@@ -134,7 +134,9 @@ class StoredProcedure:
 
     @property
     def escape_full_name(self) -> str:
-        return f"[{self.db}].[{self.schema}].[{self.formatted_name}]"
+        return f"[{self.db}].[{self.schema}].[{self.formatted_name}]".replace(
+            "'", r"''"
+        )
 
     def to_base_procedure(self) -> BaseProcedure:
         return BaseProcedure(

datahub/ingestion/source/sql/mssql/source.py

@@ -10,6 +10,7 @@ from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.base import Connection
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.exc import ProgrammingError, ResourceClosedError
+from sqlalchemy.sql import quoted_name
 
 import datahub.metadata.schema_classes as models
 from datahub.configuration.common import AllowDenyPattern
@@ -130,10 +131,14 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         "match the entire table name in database.schema.table format. Defaults are to set in such a way "
         "to ignore the temporary staging tables created by known ETL tools.",
     )
+    quote_schemas: bool = Field(
+        default=False,
+        description="Represent a schema identifiers combined with quoting preferences. See [sqlalchemy quoted_name docs](https://docs.sqlalchemy.org/en/20/core/sqlelement.html#sqlalchemy.sql.expression.quoted_name).",
+    )
 
     @pydantic.validator("uri_args")
     def passwords_match(cls, v, values, **kwargs):
-        if values["use_odbc"] and "driver" not in v:
+        if values["use_odbc"] and not values["sqlalchemy_uri"] and "driver" not in v:
             raise ValueError("uri_args must contain a 'driver' option")
         elif not values["use_odbc"] and v:
             raise ValueError("uri_args is not supported when ODBC is disabled")
@@ -159,7 +164,15 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
             uri_opts=uri_opts,
         )
         if self.use_odbc:
-            uri = f"{uri}?{urllib.parse.urlencode(self.uri_args)}"
+            final_uri_args = self.uri_args.copy()
+            if final_uri_args and current_db:
+                final_uri_args.update({"database": current_db})
+
+            uri = (
+                f"{uri}?{urllib.parse.urlencode(final_uri_args)}"
+                if final_uri_args
+                else uri
+            )
         return uri
 
     @property
@@ -923,7 +936,11 @@ class SQLServerSource(SQLAlchemySource):
         logger.debug(f"sql_alchemy_url={url}")
         engine = create_engine(url, **self.config.options)
 
-        if self.config.database and self.config.database != "":
+        if (
+            self.config.database
+            and self.config.database != ""
+            or (self.config.sqlalchemy_uri and self.config.sqlalchemy_uri != "")
+        ):
             inspector = inspect(engine)
             yield inspector
         else:
@@ -1020,3 +1037,45 @@ class SQLServerSource(SQLAlchemySource):
             if self.config.convert_urns_to_lowercase
             else table_ref_str
         )
+
+    def get_allowed_schemas(self, inspector: Inspector, db_name: str) -> Iterable[str]:
+        for schema in super().get_allowed_schemas(inspector, db_name):
+            if self.config.quote_schemas:
+                yield quoted_name(schema, True)
+            else:
+                yield schema
+
+    def get_db_name(self, inspector: Inspector) -> str:
+        engine = inspector.engine
+
+        try:
+            if (
+                engine
+                and hasattr(engine, "url")
+                and hasattr(engine.url, "database")
+                and engine.url.database
+            ):
+                return str(engine.url.database).strip('"')
+
+            if (
+                engine
+                and hasattr(engine, "url")
+                and hasattr(engine.url, "query")
+                and "odbc_connect" in engine.url.query
+            ):
+                # According to the ODBC connection keywords: https://learn.microsoft.com/en-us/sql/connect/odbc/dsn-connection-string-attribute?view=sql-server-ver17#supported-dsnconnection-string-keywords-and-connection-attributes
+                database = re.search(
+                    r"DATABASE=([^;]*);",
+                    urllib.parse.unquote_plus(str(engine.url.query["odbc_connect"])),
+                    flags=re.IGNORECASE,
+                )
+
+                if database and database.group(1):
+                    return database.group(1)
+
+            return ""
+
+        except Exception as e:
+            raise RuntimeError(
+                "Unable to get database name from Sqlalchemy inspector"
+            ) from e

datahub/ingestion/source/sql_queries.py

@@ -25,6 +25,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,
@@ -48,7 +52,9 @@ from datahub.sql_parsing.sql_parsing_aggregator import (
 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(
@@ -109,6 +115,16 @@ class SqlQueriesSource(Source):
      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.
     """
 
     schema_resolver: Optional[SchemaResolver]
@@ -165,7 +181,13 @@ 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,

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
@@ -117,10 +115,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 +161,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,