acryl-datahub 1.3.1__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

datahub/ingestion/source/powerbi/config.py

@@ -4,7 +4,7 @@ from enum import Enum
 from typing import Dict, List, Literal, Optional, Union
 
 import pydantic
-from pydantic import root_validator, validator
+from pydantic import field_validator, model_validator
 
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
@@ -540,8 +540,8 @@ class PowerBiDashboardSourceConfig(
         description="timeout in seconds for Metadata Rest Api.",
     )
 
-    @root_validator(skip_on_failure=True)
-    def validate_extract_column_level_lineage(cls, values: Dict) -> Dict:
+    @model_validator(mode="after")
+    def validate_extract_column_level_lineage(self) -> "PowerBiDashboardSourceConfig":
         flags = [
             "native_query_parsing",
             "enable_advance_lineage_sql_construct",
@@ -549,26 +549,23 @@ class PowerBiDashboardSourceConfig(
             "extract_dataset_schema",
         ]
 
-        if (
-            "extract_column_level_lineage" in values
-            and values["extract_column_level_lineage"] is False
-        ):
+        if self.extract_column_level_lineage is False:
             # Flag is not set. skip validation
-            return values
+            return self
 
         logger.debug(f"Validating additional flags: {flags}")
 
         is_flag_enabled: bool = True
         for flag in flags:
-            if flag not in values or values[flag] is False:
+            if not getattr(self, flag, True):
                 is_flag_enabled = False
 
         if not is_flag_enabled:
             raise ValueError(f"Enable all these flags in recipe: {flags} ")
 
-        return values
+        return self
 
-    @validator("dataset_type_mapping")
+    @field_validator("dataset_type_mapping", mode="after")
     @classmethod
     def map_data_platform(cls, value):
         # For backward compatibility convert input PostgreSql to PostgreSQL
@@ -580,28 +577,32 @@ class PowerBiDashboardSourceConfig(
 
         return value
 
-    @root_validator(skip_on_failure=True)
-    def workspace_id_backward_compatibility(cls, values: Dict) -> Dict:
-        workspace_id = values.get("workspace_id")
-        workspace_id_pattern = values.get("workspace_id_pattern")
-
-        if workspace_id_pattern == AllowDenyPattern.allow_all() and workspace_id:
+    @model_validator(mode="after")
+    def workspace_id_backward_compatibility(self) -> "PowerBiDashboardSourceConfig":
+        if (
+            self.workspace_id_pattern == AllowDenyPattern.allow_all()
+            and self.workspace_id
+        ):
             logger.warning(
                 "workspace_id_pattern is not set but workspace_id is set, setting workspace_id as "
                 "workspace_id_pattern. workspace_id will be deprecated, please use workspace_id_pattern instead."
             )
-            values["workspace_id_pattern"] = AllowDenyPattern(
-                allow=[f"^{workspace_id}$"]
+            self.workspace_id_pattern = AllowDenyPattern(
+                allow=[f"^{self.workspace_id}$"]
             )
-        elif workspace_id_pattern != AllowDenyPattern.allow_all() and workspace_id:
+        elif (
+            self.workspace_id_pattern != AllowDenyPattern.allow_all()
+            and self.workspace_id
+        ):
             logger.warning(
                 "workspace_id will be ignored in favour of workspace_id_pattern. workspace_id will be deprecated, "
                 "please use workspace_id_pattern only."
             )
-            values.pop("workspace_id")
-        return values
+            self.workspace_id = None
+        return self
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def raise_error_for_dataset_type_mapping(cls, values: Dict) -> Dict:
         if (
             values.get("dataset_type_mapping") is not None
@@ -613,18 +614,18 @@ class PowerBiDashboardSourceConfig(
 
         return values
 
-    @root_validator(skip_on_failure=True)
-    def validate_extract_dataset_schema(cls, values: Dict) -> Dict:
-        if values.get("extract_dataset_schema") is False:
+    @model_validator(mode="after")
+    def validate_extract_dataset_schema(self) -> "PowerBiDashboardSourceConfig":
+        if self.extract_dataset_schema is False:
             add_global_warning(
                 "Please use `extract_dataset_schema: true`, otherwise dataset schema extraction will be skipped."
             )
-        return values
+        return self
 
-    @root_validator(skip_on_failure=True)
-    def validate_dsn_to_database_schema(cls, values: Dict) -> Dict:
-        if values.get("dsn_to_database_schema") is not None:
-            dsn_mapping = values.get("dsn_to_database_schema")
+    @model_validator(mode="after")
+    def validate_dsn_to_database_schema(self) -> "PowerBiDashboardSourceConfig":
+        if self.dsn_to_database_schema is not None:
+            dsn_mapping = self.dsn_to_database_schema
             if not isinstance(dsn_mapping, dict):
                 raise ValueError("dsn_to_database_schema must contain key-value pairs")
 
@@ -639,4 +640,4 @@ class PowerBiDashboardSourceConfig(
                         f"dsn_to_database_schema invalid mapping value: {value}"
                     )
 
-        return values
+        return self

datahub/ingestion/source/powerbi/dataplatform_instance_resolver.py

@@ -41,7 +41,7 @@ class ResolvePlatformInstanceFromDatasetTypeMapping(
         if isinstance(platform, PlatformDetail):
             return platform
 
-        return PlatformDetail.parse_obj({})
+        return PlatformDetail.model_validate({})
 
 
 class ResolvePlatformInstanceFromServerToPlatformInstance(
@@ -56,7 +56,7 @@ class ResolvePlatformInstanceFromServerToPlatformInstance(
             ]
             if data_platform_detail.data_platform_server
             in self.config.server_to_platform_instance
-            else PlatformDetail.parse_obj({})
+            else PlatformDetail.model_validate({})
         )
 
 

datahub/ingestion/source/powerbi/powerbi.py

@@ -1316,7 +1316,7 @@ class PowerBiDashboardSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = PowerBiDashboardSourceConfig.parse_obj(config_dict)
+        config = PowerBiDashboardSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_allowed_workspaces(self) -> List[powerbi_data_classes.Workspace]:

datahub/ingestion/source/powerbi_report_server/report_server.py

@@ -213,7 +213,7 @@ class PowerBiReportServerAPI:
 
             if response_dict.get("value"):
                 reports.extend(
-                    report_types_mapping[report_type].parse_obj(report)
+                    report_types_mapping[report_type].model_validate(report)
                     for report in response_dict.get("value")
                 )
 
@@ -517,7 +517,7 @@ class PowerBiReportServerDashboardSource(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = PowerBiReportServerDashboardSourceConfig.parse_obj(config_dict)
+        config = PowerBiReportServerDashboardSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:

datahub/ingestion/source/powerbi_report_server/report_server_domain.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, Field, validator
+from pydantic import BaseModel, Field, model_validator
 
 from datahub.ingestion.source.powerbi_report_server.constants import (
     RelationshipDirection,
@@ -30,11 +30,13 @@ class CatalogItem(BaseModel):
     has_data_sources: bool = Field(False, alias="HasDataSources")
     data_sources: Optional[List["DataSource"]] = Field(None, alias="DataSources")
 
-    @validator("display_name", always=True)
-    def validate_diplay_name(cls, value, values):
-        if values["created_by"]:
-            return values["created_by"].split("\\")[-1]
-        return ""
+    @model_validator(mode="after")
+    def validate_diplay_name(self):
+        if self.created_by:
+            self.display_name = self.created_by.split("\\")[-1]
+        else:
+            self.display_name = ""
+        return self
 
     def get_urn_part(self):
         return f"reports.{self.id}"

datahub/ingestion/source/preset.py

@@ -2,7 +2,7 @@ import logging
 from typing import Dict, Optional
 
 import requests
-from pydantic import root_validator, validator
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 
 from datahub.emitter.mce_builder import DEFAULT_ENV
@@ -55,16 +55,16 @@ class PresetConfig(SupersetConfig):
         description="Can be used to change mapping for database names in superset to what you have in datahub",
     )
 
-    @validator("connect_uri", "display_uri")
+    @field_validator("connect_uri", "display_uri", mode="after")
+    @classmethod
     def remove_trailing_slash(cls, v):
         return config_clean.remove_trailing_slashes(v)
 
-    @root_validator(skip_on_failure=True)
-    def default_display_uri_to_connect_uri(cls, values):
-        base = values.get("display_uri")
-        if base is None:
-            values["display_uri"] = values.get("connect_uri")
-        return values
+    @model_validator(mode="after")
+    def default_display_uri_to_connect_uri(self) -> "PresetConfig":
+        if self.display_uri is None:
+            self.display_uri = self.connect_uri
+        return self
 
 
 @platform_name("Preset")

datahub/ingestion/source/pulsar.py

@@ -235,7 +235,7 @@ class PulsarSource(StatefulIngestionSourceBase):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = PulsarSourceConfig.parse_obj(config_dict)
+        config = PulsarSourceConfig.model_validate(config_dict)
 
         # Do not include each individual partition for partitioned topics,
         if config.exclude_individual_partitions:

datahub/ingestion/source/qlik_sense/data_classes.py

@@ -3,7 +3,7 @@ from datetime import datetime
 from enum import Enum
 from typing import Dict, List, Optional, Type, Union
 
-from pydantic import BaseModel, ConfigDict, Field, root_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.source.qlik_sense.config import QLIK_DATETIME_FORMAT, Constant
@@ -92,7 +92,8 @@ class Space(_QlikBaseModel):
     updatedAt: datetime
     ownerId: Optional[str] = None
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -121,7 +122,8 @@ class SchemaField(_QlikBaseModel):
     primaryKey: Optional[bool] = None
     nullable: Optional[bool] = None
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -138,7 +140,8 @@ class QlikDataset(Item):
     itemId: str
     datasetSchema: List[SchemaField]
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -174,7 +177,8 @@ class Chart(_QlikBaseModel):
     qDimension: List[AxisProperty]
     qMeasure: List[AxisProperty]
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -193,7 +197,8 @@ class Sheet(_QlikBaseModel):
     updatedAt: datetime
     charts: List[Chart] = []
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -220,7 +225,8 @@ class QlikTable(_QlikBaseModel):
     databaseName: Optional[str] = None
     schemaName: Optional[str] = None
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -239,7 +245,8 @@ class App(Item):
     sheets: List[Sheet] = []
     tables: List[QlikTable] = []
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def update_values(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)

datahub/ingestion/source/qlik_sense/qlik_api.py

@@ -56,7 +56,7 @@ class QlikAPI:
                 response.raise_for_status()
                 response_dict = response.json()
                 for space_dict in response_dict[Constant.DATA]:
-                    space = Space.parse_obj(space_dict)
+                    space = Space.model_validate(space_dict)
                     spaces.append(space)
                     self.spaces[space.id] = space.name
                 if Constant.NEXT in response_dict[Constant.LINKS]:
@@ -64,7 +64,7 @@ class QlikAPI:
                 else:
                     break
             # Add personal space entity
-            spaces.append(Space.parse_obj(PERSONAL_SPACE_DICT))
+            spaces.append(Space.model_validate(PERSONAL_SPACE_DICT))
             self.spaces[PERSONAL_SPACE_DICT[Constant.ID]] = PERSONAL_SPACE_DICT[
                 Constant.NAME
             ]
@@ -78,7 +78,7 @@ class QlikAPI:
             response.raise_for_status()
             response_dict = response.json()
             response_dict[Constant.ITEMID] = item_id
-            return QlikDataset.parse_obj(response_dict)
+            return QlikDataset.model_validate(response_dict)
         except Exception as e:
             self._log_http_error(
                 message=f"Unable to fetch dataset with id {dataset_id}. Exception: {e}"
@@ -119,7 +119,7 @@ class QlikAPI:
                     f"Chart with id {chart_id} of sheet {sheet_id} does not have hypercube. q_layout: {q_layout}"
                 )
                 return None
-            return Chart.parse_obj(q_layout)
+            return Chart.model_validate(q_layout)
         except Exception as e:
             self._log_http_error(
                 message=f"Unable to fetch chart {chart_id} of sheet {sheet_id}. Exception: {e}"
@@ -140,7 +140,7 @@ class QlikAPI:
             if Constant.OWNERID not in sheet_dict[Constant.QMETA]:
                 # That means sheet is private sheet
                 return None
-            sheet = Sheet.parse_obj(sheet_dict[Constant.QMETA])
+            sheet = Sheet.model_validate(sheet_dict[Constant.QMETA])
             if Constant.QCHILDLIST not in sheet_dict:
                 logger.warning(
                     f"Sheet {sheet.title} with id {sheet_id} does not have any charts. sheet_dict: {sheet_dict}"
@@ -222,7 +222,7 @@ class QlikAPI:
                 return []
             response = websocket_connection.websocket_send_request(method="GetLayout")
             for table_dict in response[Constant.QLAYOUT][Constant.TABLES]:
-                tables.append(QlikTable.parse_obj(table_dict))
+                tables.append(QlikTable.model_validate(table_dict))
             websocket_connection.handle.pop()
             self._add_qri_of_tables(tables, app_id)
         except Exception as e:
@@ -270,7 +270,7 @@ class QlikAPI:
             response = websocket_connection.websocket_send_request(
                 method="GetAppLayout"
             )
-            app = App.parse_obj(response[Constant.QLAYOUT])
+            app = App.model_validate(response[Constant.QLAYOUT])
             app.sheets = self._get_app_sheets(websocket_connection, app_id)
             app.tables = self._get_app_used_tables(websocket_connection, app_id)
             websocket_connection.close_websocket()

datahub/ingestion/source/qlik_sense/qlik_sense.py

@@ -148,7 +148,7 @@ class QlikSenseSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = QlikSourceConfig.parse_obj(config_dict)
+        config = QlikSourceConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     def _gen_space_key(self, space_id: str) -> SpaceKey:

datahub/ingestion/source/redshift/config.py

@@ -3,7 +3,7 @@ from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from pydantic import root_validator
+from pydantic import model_validator
 from pydantic.fields import Field
 
 from datahub.configuration import ConfigModel
@@ -182,7 +182,8 @@ class RedshiftConfig(
         description="Whether to skip EXTERNAL tables.",
     )
 
-    @root_validator(pre=True)
+    @model_validator(mode="before")
+    @classmethod
     def check_email_is_set_on_usage(cls, values):
         if values.get("include_usage_statistics"):
             assert "email_domain" in values and values["email_domain"], (
@@ -190,31 +191,28 @@ class RedshiftConfig(
             )
         return values
 
-    @root_validator(skip_on_failure=True)
-    def check_database_is_set(cls, values):
-        assert values.get("database"), "database must be set"
-        return values
-
-    @root_validator(skip_on_failure=True)
-    def backward_compatibility_configs_set(cls, values: Dict) -> Dict:
-        match_fully_qualified_names = values.get("match_fully_qualified_names")
-
-        schema_pattern: Optional[AllowDenyPattern] = values.get("schema_pattern")
+    @model_validator(mode="after")
+    def check_database_is_set(self) -> "RedshiftConfig":
+        assert self.database, "database must be set"
+        return self
 
+    @model_validator(mode="after")
+    def backward_compatibility_configs_set(self) -> "RedshiftConfig":
         if (
-            schema_pattern is not None
-            and schema_pattern != AllowDenyPattern.allow_all()
-            and match_fully_qualified_names is not None
-            and not match_fully_qualified_names
+            self.schema_pattern is not None
+            and self.schema_pattern != AllowDenyPattern.allow_all()
+            and self.match_fully_qualified_names is not None
+            and not self.match_fully_qualified_names
         ):
             logger.warning(
                 "Please update `schema_pattern` to match against fully qualified schema name `<database_name>.<schema_name>` and set config `match_fully_qualified_names : True`."
                 "Current default `match_fully_qualified_names: False` is only to maintain backward compatibility. "
                 "The config option `match_fully_qualified_names` will be deprecated in future and the default behavior will assume `match_fully_qualified_names: True`."
             )
-        return values
+        return self
 
-    @root_validator(skip_on_failure=True)
+    @model_validator(mode="before")
+    @classmethod
     def connection_config_compatibility_set(cls, values: Dict) -> Dict:
         # Create a copy to avoid modifying the input dictionary, preventing state contamination in tests
         values = deepcopy(values)
@@ -231,8 +229,8 @@ class RedshiftConfig(
         if "options" in values and "connect_args" in values["options"]:
             values["extra_client_options"] = values["options"]["connect_args"]
 
-        if values["extra_client_options"]:
-            if values["options"]:
+        if values.get("extra_client_options"):
+            if values.get("options"):
                 values["options"]["connect_args"] = values["extra_client_options"]
             else:
                 values["options"] = {"connect_args": values["extra_client_options"]}

datahub/ingestion/source/redshift/redshift.py

@@ -236,7 +236,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
             RedshiftConfig.Config.extra = (
                 pydantic.Extra.allow
             )  # we are okay with extra fields during this stage
-            config = RedshiftConfig.parse_obj(config_dict)
+            config = RedshiftConfig.model_validate(config_dict)
             # source = RedshiftSource(config, report)
             connection: redshift_connector.Connection = (
                 RedshiftSource.get_redshift_connection(config)
@@ -316,7 +316,7 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
 
     @classmethod
     def create(cls, config_dict, ctx):
-        config = RedshiftConfig.parse_obj(config_dict)
+        config = RedshiftConfig.model_validate(config_dict)
         return cls(config, ctx)
 
     @staticmethod

datahub/ingestion/source/redshift/usage.py

@@ -1,12 +1,12 @@
 import collections
 import logging
 import time
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
 
 import cachetools
-import pydantic.error_wrappers
 import redshift_connector
+from pydantic import ValidationError, field_validator
 from pydantic.fields import Field
 from pydantic.main import BaseModel
 
@@ -64,6 +64,26 @@ class RedshiftAccessEvent(BaseModel):
     starttime: datetime
     endtime: datetime
 
+    @field_validator("starttime", "endtime", mode="before")
+    @classmethod
+    def ensure_utc_datetime(cls, v):
+        """Ensure datetime fields are treated as UTC for consistency with Pydantic V1 behavior.
+
+        Pydantic V2 assumes local timezone for naive datetime strings, whereas Pydantic V1 assumed UTC.
+        This validator restores V1 behavior to maintain timestamp consistency.
+        """
+        if isinstance(v, str):
+            # Parse as naive datetime, then assume UTC (matching V1 behavior)
+            dt = datetime.fromisoformat(v)
+            if dt.tzinfo is None:
+                # Treat naive datetime as UTC (this was the V1 behavior)
+                dt = dt.replace(tzinfo=timezone.utc)
+            return dt
+        elif isinstance(v, datetime) and v.tzinfo is None:
+            # If we get a naive datetime object, assume UTC
+            return v.replace(tzinfo=timezone.utc)
+        return v
+
 
 class RedshiftUsageExtractor:
     """
@@ -291,7 +311,7 @@ class RedshiftUsageExtractor:
                             else None
                         ),
                     )
-                except pydantic.error_wrappers.ValidationError as e:
+                except ValidationError as e:
                     logging.warning(
                         f"Validation error on access event creation from row {row}. The error was: {e} Skipping ...."
                     )