acryl-datahub-cloud 0.3.14__tar.gz → 0.3.16.1rc5__tar.gz

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/PKG-INFO

@@ -1,101 +1,98 @@
 Metadata-Version: 2.4
 Name: acryl-datahub-cloud
-Version: 0.3.14
+Version: 0.3.16.1rc5
 Requires-Python: >=3.10
 Requires-Dist: avro-gen3==0.7.16
 Requires-Dist: acryl-datahub
+Requires-Dist: acryl-datahub~=1.3.1.5
 Requires-Dist: croniter
 Requires-Dist: pytz
 Requires-Dist: types-croniter
 Requires-Dist: tzlocal
 Provides-Extra: datahub-lineage-features
 Requires-Dist: duckdb; extra == "datahub-lineage-features"
+Requires-Dist: pandas; extra == "datahub-lineage-features"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-lineage-features"
 Requires-Dist: tenacity; extra == "datahub-lineage-features"
-Requires-Dist: pydantic<2; extra == "datahub-lineage-features"
 Requires-Dist: pyarrow; extra == "datahub-lineage-features"
-Requires-Dist: pandas; extra == "datahub-lineage-features"
 Provides-Extra: datahub-reporting-forms
 Requires-Dist: duckdb; extra == "datahub-reporting-forms"
-Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Requires-Dist: termcolor==2.5.0; extra == "datahub-reporting-forms"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-forms"
-Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
+Requires-Dist: boto3; extra == "datahub-reporting-forms"
 Requires-Dist: pandas; extra == "datahub-reporting-forms"
+Requires-Dist: pyarrow; extra == "datahub-reporting-forms"
 Provides-Extra: datahub-reporting-extract-graph
 Requires-Dist: duckdb; extra == "datahub-reporting-extract-graph"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-graph"
+Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-graph"
 Requires-Dist: pyarrow; extra == "datahub-reporting-extract-graph"
-Requires-Dist: pandas; extra == "datahub-reporting-extract-graph"
 Provides-Extra: datahub-reporting-extract-sql
 Requires-Dist: duckdb; extra == "datahub-reporting-extract-sql"
 Requires-Dist: boto3; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pydantic<2; extra == "datahub-reporting-extract-sql"
-Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
 Requires-Dist: pandas; extra == "datahub-reporting-extract-sql"
+Requires-Dist: pyarrow; extra == "datahub-reporting-extract-sql"
 Provides-Extra: datahub-usage-reporting
-Requires-Dist: polars==1.30.0; extra == "datahub-usage-reporting"
-Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
-Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
+Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
 Requires-Dist: duckdb; extra == "datahub-usage-reporting"
-Requires-Dist: boto3; extra == "datahub-usage-reporting"
-Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
+Requires-Dist: pandas; extra == "datahub-usage-reporting"
 Requires-Dist: opensearch-py==2.4.2; extra == "datahub-usage-reporting"
-Requires-Dist: pydantic<2; extra == "datahub-usage-reporting"
-Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
-Requires-Dist: numpy<2; extra == "datahub-usage-reporting"
+Requires-Dist: elasticsearch==7.13.4; extra == "datahub-usage-reporting"
 Requires-Dist: pyarrow<=18.0.0; extra == "datahub-usage-reporting"
-Requires-Dist: pandas; extra == "datahub-usage-reporting"
+Requires-Dist: pyarrow; extra == "datahub-usage-reporting"
+Requires-Dist: termcolor==2.5.0; extra == "datahub-usage-reporting"
+Requires-Dist: boto3; extra == "datahub-usage-reporting"
+Requires-Dist: polars==1.34.0; extra == "datahub-usage-reporting"
+Requires-Dist: scipy<=1.14.1; extra == "datahub-usage-reporting"
 Provides-Extra: datahub-metadata-sharing
 Requires-Dist: tenacity; extra == "datahub-metadata-sharing"
 Provides-Extra: datahub-action-request-owner
 Requires-Dist: tenacity; extra == "datahub-action-request-owner"
 Provides-Extra: acryl-cs-issues
+Requires-Dist: zenpy; extra == "acryl-cs-issues"
 Requires-Dist: jinja2; extra == "acryl-cs-issues"
 Requires-Dist: openai; extra == "acryl-cs-issues"
-Requires-Dist: zenpy; extra == "acryl-cs-issues"
 Requires-Dist: slack-sdk; extra == "acryl-cs-issues"
 Provides-Extra: datahub-forms-notifications
 Requires-Dist: tenacity; extra == "datahub-forms-notifications"
 Provides-Extra: all
-Requires-Dist: polars==1.30.0; extra == "all"
-Requires-Dist: termcolor==2.5.0; extra == "all"
-Requires-Dist: tenacity; extra == "all"
-Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: numpy<2; extra == "all"
 Requires-Dist: duckdb; extra == "all"
-Requires-Dist: boto3; extra == "all"
+Requires-Dist: pandas; extra == "all"
+Requires-Dist: slack-sdk; extra == "all"
 Requires-Dist: opensearch-py==2.4.2; extra == "all"
-Requires-Dist: numpy<2; extra == "all"
-Requires-Dist: jinja2; extra == "all"
+Requires-Dist: elasticsearch==7.13.4; extra == "all"
 Requires-Dist: pyarrow<=18.0.0; extra == "all"
-Requires-Dist: zenpy; extra == "all"
+Requires-Dist: tenacity; extra == "all"
+Requires-Dist: pyarrow; extra == "all"
+Requires-Dist: termcolor==2.5.0; extra == "all"
+Requires-Dist: boto3; extra == "all"
+Requires-Dist: polars==1.34.0; extra == "all"
 Requires-Dist: openai; extra == "all"
-Requires-Dist: slack-sdk; extra == "all"
-Requires-Dist: elasticsearch==7.13.4; extra == "all"
+Requires-Dist: zenpy; extra == "all"
 Requires-Dist: scipy<=1.14.1; extra == "all"
-Requires-Dist: pydantic<2; extra == "all"
-Requires-Dist: pandas; extra == "all"
+Requires-Dist: jinja2; extra == "all"
 Provides-Extra: dev
-Requires-Dist: polars==1.30.0; extra == "dev"
-Requires-Dist: acryl-datahub[dev]; extra == "dev"
-Requires-Dist: openai; extra == "dev"
-Requires-Dist: termcolor==2.5.0; extra == "dev"
+Requires-Dist: duckdb; extra == "dev"
 Requires-Dist: tenacity; extra == "dev"
+Requires-Dist: polars==1.34.0; extra == "dev"
 Requires-Dist: slack-sdk; extra == "dev"
-Requires-Dist: pyarrow; extra == "dev"
-Requires-Dist: duckdb; extra == "dev"
-Requires-Dist: boto3; extra == "dev"
-Requires-Dist: opensearch-py==2.4.2; extra == "dev"
-Requires-Dist: elasticsearch==7.13.4; extra == "dev"
-Requires-Dist: pydantic<2; extra == "dev"
-Requires-Dist: scipy<=1.14.1; extra == "dev"
 Requires-Dist: numpy<2; extra == "dev"
-Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: opensearch-py==2.4.2; extra == "dev"
+Requires-Dist: acryl-datahub[dev]; extra == "dev"
 Requires-Dist: pyarrow<=18.0.0; extra == "dev"
+Requires-Dist: openai; extra == "dev"
+Requires-Dist: boto3; extra == "dev"
 Requires-Dist: zenpy; extra == "dev"
+Requires-Dist: jinja2; extra == "dev"
+Requires-Dist: pandas-stubs; extra == "dev"
+Requires-Dist: pyarrow-stubs; extra == "dev"
+Requires-Dist: termcolor==2.5.0; extra == "dev"
+Requires-Dist: scipy-stubs; extra == "dev"
 Requires-Dist: pandas; extra == "dev"
+Requires-Dist: elasticsearch==7.13.4; extra == "dev"
+Requires-Dist: pyarrow; extra == "dev"
+Requires-Dist: scipy<=1.14.1; extra == "dev"
 Dynamic: provides-extra
 Dynamic: requires-dist
 Dynamic: requires-python

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/setup.py

@@ -6,17 +6,22 @@ from setuptools import setup
 _codegen_config_file = pathlib.Path("./src/acryl_datahub_cloud/_codegen_config.json")
 _codegen_config: dict = json.loads(_codegen_config_file.read_text())
 
+acryl_datahub = [
+    # 1.3.0: Pydantic v2 was introduced in 1.2.0.10 and later RCs fixed in OSS some missing code from /metadata-ingestion SaaS
+    "acryl-datahub~=1.3.1.5"
+]
+
 # Note: We are using the croniter library for cron parsing which is different from executor, which uses apscheduler, so there is a risk of mismatch here.
 # croniter is now maintained at: https://github.com/pallets-eco/croniter
 base_requirements = [
+    *acryl_datahub,
     "croniter",
     "pytz",
     "types-croniter",
     "tzlocal",
 ]
 
-# Adding pydantic<2 since we use pydantic models to map to pyarrow models and that is only compatible in pydantic v1
-stats_common = {"pandas", "pyarrow", "duckdb", "pydantic<2"}
+stats_common = {"pandas", "pyarrow", "duckdb"}
 aws_common = {"boto3"}
 open_search_common = {"opensearch-py==2.4.2"}
 
@@ -33,7 +38,7 @@ plugins = {
     | aws_common
     | {
         "opensearch-py==2.4.2",
-        "polars==1.30.0",
+        "polars==1.34.0",
         "elasticsearch==7.13.4",
         "numpy<2",
         "scipy<=1.14.1",
@@ -49,6 +54,10 @@ plugins = {
 dev_requirements = {
     # acryl-datahub[dev] pulls in more things than are strictly necessary, but it's fine.
     "acryl-datahub[dev]",
+    # Type stubs for external libraries
+    "pyarrow-stubs",
+    "scipy-stubs",
+    "pandas-stubs",
     *list(
         dependency
         for plugin in [

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.14",
+  "version": "0.3.16.1rc5",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/acryl_cs_issues/models.py

@@ -4,7 +4,7 @@ from enum import Enum
 from typing import Any, Iterable, List, Optional
 
 import pydantic
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 
 class DataHubProductGlossary:
@@ -308,7 +308,8 @@ class Project(BaseModel):
         platform: str
         id: str
 
-        @validator("id", pre=True)
+        @field_validator("id", mode="before")
+        @classmethod
         def id_must_be_coerced_to_string(cls, v: Any) -> str:
             if not v:
                 raise ValueError("id must be present")
@@ -328,7 +329,8 @@ class Issue(BaseModel):
         platform: str
         id: str
 
-        @validator("id", pre=True)
+        @field_validator("id", mode="before")
+        @classmethod
         def id_must_be_coerced_to_string(cls, v: Any) -> str:
             if not v:
                 raise ValueError("id must be present")

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/action_request/action_request_owner_source.py

@@ -98,7 +98,7 @@ class ActionRequestOwnerSource(Source):
         action_type = action_request.get("type")
         action_request_entity = action_request.get("entity")
         if action_request_entity is None:
-            self.report.failure(
+            self.report.warning(
                 title="Action request entity not found",
                 message="Action request entity not found",
                 context=str(action_request_urn),

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/datahub_forms_notifications/forms_notifications_source.py

@@ -1,4 +1,3 @@
-import json
 import logging
 import time
 from dataclasses import dataclass
@@ -18,6 +17,7 @@ from acryl_datahub_cloud.datahub_forms_notifications.query import (
     GRAPHQL_SCROLL_FORMS_FOR_NOTIFICATIONS,
     GRAPHQL_SEND_FORM_NOTIFICATION_REQUEST,
 )
+from acryl_datahub_cloud.graphql_utils import parse_extra_properties_for_model
 from acryl_datahub_cloud.notifications.notification_recipient_builder import (
     NotificationRecipientBuilder,
 )
@@ -391,7 +391,7 @@ class DataHubFormsNotificationsSource(Source):
         user_urns = []
         group_urns = []
 
-        extra_fields = [f for f in DataHubDatasetSearchRow.__fields__]
+        extra_fields = [f for f in DataHubDatasetSearchRow.model_fields]
         results = self.graph.get_results_by_filter(
             extra_or_filters=self._get_incomplete_assets_for_form(form_urn, form.type),
             extra_source_fields=extra_fields,
@@ -399,9 +399,9 @@ class DataHubFormsNotificationsSource(Source):
         )
         for result in results:
             extra_properties = result["extraProperties"]
-            extra_properties_map = {
-                x["name"]: json.loads(x["value"]) for x in extra_properties
-            }
+            extra_properties_map = parse_extra_properties_for_model(
+                extra_properties, DataHubDatasetSearchRow
+            )
             search_row = DataHubDatasetSearchRow(**extra_properties_map)
             for owner in search_row.owners:
                 if owner.startswith(USER_URN_PREFIX):

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/datahub_reporting/datahub_dataset.py

@@ -5,14 +5,14 @@ import pathlib
 import tempfile
 import time
 from enum import Enum
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union, cast
 
 import boto3
 import duckdb
 import pandas
 import pyarrow as pa
 import pyarrow.parquet as pq
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow, SchemaField
 from datahub.configuration.common import ConfigModel
@@ -73,7 +73,9 @@ class FileStoreBackedDatasetConfig(ConfigModel):
     store_platform: str = "s3"
     file_name: str = "data"
     file_extension: str = "parquet"
-    file_compression: str = "snappy"
+    file_compression: Literal[
+        "gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"
+    ] = "snappy"
     file_overwrite_existing: bool = True
     snapshot_partitioning_strategy: str = PartitioningStrategy.DATE
     generate_presigned_url: bool = True
@@ -85,7 +87,8 @@ class FileStoreBackedDatasetConfig(ConfigModel):
 
     datahub_platform: str = "acryl"
 
-    @validator("snapshot_partitioning_strategy")
+    @field_validator("snapshot_partitioning_strategy")
+    @classmethod
     def validate_partitioning_strategy(cls, v):
         if v not in PartitioningStrategy._value2member_map_:
             raise ValueError(f"Unsupported partitioning strategy: {v}")
@@ -119,9 +122,14 @@ class DataHubBasedS3Dataset:
         self.local_file_path: str = (
             config.file if config.file else self._initialize_local_file()
         )
-        self.file_writer = None
+        self.file_writer: Optional[pq.ParquetWriter] = None
         self.schema = (
-            pa.schema([(x.name, x.type) for x in self.dataset_metadata.schemaFields])
+            pa.schema(
+                [
+                    pa.field(x.name, BaseModelRow.string_to_pyarrow_type(x.type))
+                    for x in self.dataset_metadata.schemaFields
+                ]
+            )
             if self.dataset_metadata.schemaFields
             else None
         )
@@ -163,18 +171,32 @@ class DataHubBasedS3Dataset:
                     self.schema = row.arrow_schema()
                 else:
                     # hail mary: infer schema from the first row and cast everything to string
-                    self.schema = pa.schema([(key, pa.string()) for key in row])
+                    self.schema = pa.schema([pa.field(key, pa.string()) for key in row])
                     self.stringify_row = True
 
             self._initialize_local_file()
+            # Map compression names to PyArrow format (most are direct mappings)
+            compression_map = {
+                "gzip": "gzip",
+                "bz2": "brotli",  # PyArrow doesn't support bz2, use brotli
+                "brotli": "brotli",
+                "lz4": "lz4",
+                "zstd": "zstd",
+                "snappy": "snappy",
+                "none": "none",
+            }
+            compression = cast(
+                Literal["gzip", "bz2", "brotli", "lz4", "zstd", "snappy", "none"],
+                compression_map.get(self.config.file_compression, "snappy"),
+            )
             self.file_writer = pq.ParquetWriter(
                 self.local_file_path,
                 self.schema,
-                compression=self.config.file_compression,
+                compression=compression,
             )
         if isinstance(row, (BaseModel, BaseModelRow)):
             # for anything extending BaseModel, we want to use the dict representation
-            write_row: Dict[str, Any] = row.dict()
+            write_row: Dict[str, Any] = row.model_dump()
         elif isinstance(row, dict):
             write_row = row
         else:
@@ -271,7 +293,7 @@ class DataHubBasedS3Dataset:
             )
 
     def _generate_schema_metadata(
-        self, duckdb_columns: List[Tuple[str, str]]
+        self, duckdb_columns: List[Tuple[str, Any]]
     ) -> SchemaMetadataClass:
         def get_type_from_dtype(dtype: str) -> SchemaFieldDataTypeClass:
             if "int" in dtype or "float" in dtype or "number" in dtype:
@@ -302,7 +324,7 @@ class DataHubBasedS3Dataset:
         )
         for column in duckdb_columns:
             # generate data type
-            data_type = column[1].lower()
+            data_type = str(column[1]).lower()
             schema_metadata.fields.append(
                 SchemaFieldClass(
                     fieldPath=column[0],
@@ -341,7 +363,7 @@ class DataHubBasedS3Dataset:
                 # generate min, max, avg, distinct count, null count
                 column_name = column[0]
                 logger.info(f"Generating field profile for {column_name}")
-                data_type = column[1].lower()
+                data_type = str(column[1]).lower()
                 if "int" in data_type or "float" in data_type:
                     query = (
                         f"SELECT COUNT(DISTINCT {column_name}), COUNT(*) - COUNT({column_name}), MIN({column_name}), MAX({column_name}), AVG({column_name})"
@@ -396,7 +418,9 @@ class DataHubBasedS3Dataset:
                 assert dataset_profiles.fieldProfiles is not None
                 dataset_profiles.fieldProfiles.append(field_profile)
             logger.info("Generated dataset profile")
-            schema_metadata = self._generate_schema_metadata(columns)
+            schema_metadata = self._generate_schema_metadata(
+                [(col[0], col[1]) for col in columns]
+            )
         return dataset_profiles, schema_metadata
 
     def register_dataset(

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/datahub_reporting/datahub_form_reporting.py

@@ -1,13 +1,13 @@
-import json
 import logging
 from datetime import date, datetime, timezone
 from enum import Enum
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 
 import pandas as pd
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 from acryl_datahub_cloud.elasticsearch.graph_service import BaseModelRow
+from acryl_datahub_cloud.graphql_utils import parse_extra_properties_for_model
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.graph.filters import RawSearchFilter
@@ -130,6 +130,22 @@ class DataHubFormReportingData(FormData):
         platformInstance: Optional[str] = None
         domains: List[str] = []
 
+        @field_validator(
+            "completedFormsIncompletePromptResponseTimes",
+            "completedFormsCompletedPromptResponseTimes",
+            "incompleteFormsIncompletePromptResponseTimes",
+            "incompleteFormsCompletedPromptResponseTimes",
+            mode="before",
+        )
+        @classmethod
+        def convert_timestamps_to_strings(
+            cls, v: Union[List[int], List[str]]
+        ) -> List[str]:
+            """Convert timestamp integers to strings for compatibility with GMS data."""
+            if not isinstance(v, list):
+                return v
+            return [str(item) for item in v]
+
     def __init__(self, graph: DataHubGraph, allowed_forms: Optional[List[str]] = None):
         self.graph: DataHubGraph = graph
         self.form_registry = FormRegistry(graph)
@@ -143,7 +159,7 @@ class DataHubFormReportingData(FormData):
         on_form_scanned: Callable[[str], Any],
     ) -> pd.DataFrame:
         return pd.DataFrame(
-            x.dict()
+            x.model_dump()
             for x in self.get_data(
                 on_asset_scanned=on_asset_scanned, on_form_scanned=on_form_scanned
             )
@@ -306,7 +322,7 @@ class DataHubFormReportingData(FormData):
         on_asset_scanned: Optional[Callable[[str], Any]] = None,
         on_form_scanned: Optional[Callable[[str], Any]] = None,
     ) -> Iterable[FormReportingRow]:
-        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__]
+        extra_fields = [f for f in self.DataHubDatasetSearchRow.model_fields]
         # TODO: Replace with the new search/filter SDK.
         result = self.graph.get_results_by_filter(
             extra_or_filters=self.get_form_existence_or_filters(),
@@ -320,10 +336,9 @@ class DataHubFormReportingData(FormData):
             if row_index % 100 == 0:
                 logger.info(f"Scanned {row_index} assets")
             extra_properties = row["extraProperties"]
-
-            extra_properties_map = {
-                x["name"]: json.loads(x["value"]) for x in extra_properties
-            }
+            extra_properties_map = parse_extra_properties_for_model(
+                extra_properties, self.DataHubDatasetSearchRow
+            )
             search_row = self.DataHubDatasetSearchRow(**extra_properties_map)
             if on_asset_scanned:
                 on_asset_scanned(search_row.urn)
@@ -430,7 +445,7 @@ class DataHubFormReportingData(FormData):
                             question_status=QuestionStatus.COMPLETED,
                             question_completed_date=datetime.fromtimestamp(
                                 float(prompt_response_time) / 1000, tz=timezone.utc
-                            ),
+                            ).date(),
                             snapshot_date=self.snapshot_date,
                         )
             complete_forms = (
@@ -532,7 +547,7 @@ class DataHubFormReportingData(FormData):
                             question_status=QuestionStatus.COMPLETED,
                             question_completed_date=datetime.fromtimestamp(
                                 float(prompt_response_time) / 1000, tz=timezone.utc
-                            ),
+                            ).date(),
                             snapshot_date=self.snapshot_date,
                         )
 

{acryl_datahub_cloud-0.3.14 → acryl_datahub_cloud-0.3.16.1rc5}/src/acryl_datahub_cloud/datahub_reporting/extract_graph.py

@@ -6,7 +6,7 @@ from typing import List, Optional
 
 import boto3
 from opensearchpy import OpenSearch
-from pydantic import validator
+from pydantic import field_validator
 
 from acryl_datahub_cloud.datahub_reporting.datahub_dataset import (
     DataHubBasedS3Dataset,
@@ -43,8 +43,9 @@ class DataHubReportingExtractGraphSourceConfig(ConfigModel):
     query_timeout: int = 30
     extract_batch_size: int = 2000
 
-    @validator("extract_graph_store", pre=True, always=True)
-    def set_default_extract_soft_delete_flag(cls, v, values):
+    @field_validator("extract_graph_store", mode="before")
+    @classmethod
+    def set_default_extract_soft_delete_flag(cls, v):
         if v is not None:
             if "dataset_registration_spec" not in v:
                 v["dataset_registration_spec"] = DatasetRegistrationSpec(