acryl-datahub 1.3.0.1rc9__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

datahub/ingestion/source/sql_queries.py

@@ -1,12 +1,14 @@
 import json
 import logging
 import os
-from dataclasses import dataclass
+import re
+from dataclasses import dataclass, field
 from datetime import datetime
 from functools import partial
-from typing import ClassVar, Iterable, List, Optional, Union
+from typing import Any, ClassVar, Iterable, List, Optional, Union, cast
 
-from pydantic import BaseModel, Field, validator
+import smart_open
+from pydantic import BaseModel, Field, field_validator
 
 from datahub.configuration.common import HiddenFromDocs
 from datahub.configuration.datetimes import parse_user_datetime
@@ -36,12 +38,13 @@ from datahub.ingestion.api.source import (
     SourceCapability,
     SourceReport,
 )
-from datahub.ingestion.api.source_helpers import auto_workunit_reporter
+from datahub.ingestion.api.source_helpers import auto_workunit, auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.source.aws.aws_common import AwsConnectionConfig
 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.schema_resolver import SchemaResolver, SchemaResolverReport
 from datahub.sql_parsing.sql_parsing_aggregator import (
     KnownQueryLineageInfo,
     ObservedQuery,
@@ -82,6 +85,24 @@ class SqlQueriesSourceConfig(
         None,
         description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
     )
+    temp_table_patterns: List[str] = Field(
+        description="Regex patterns for temporary tables to filter in lineage ingestion. "
+        "Specify regex to match the entire table name. This is useful for platforms like Athena "
+        "that don't have native temp tables but use naming patterns for fake temp tables.",
+        default=[],
+    )
+
+    enable_lazy_schema_loading: bool = Field(
+        default=True,
+        description="Enable lazy schema loading for better performance. When enabled, schemas are fetched on-demand "
+        "instead of bulk loading all schemas upfront, reducing startup time and memory usage.",
+    )
+
+    # AWS/S3 configuration
+    aws_config: Optional[AwsConnectionConfig] = Field(
+        default=None,
+        description="AWS configuration for S3 access. Required when query_file is an S3 URI (s3://).",
+    )
 
 
 @dataclass
@@ -89,8 +110,13 @@ class SqlQueriesSourceReport(SourceReport):
     num_entries_processed: int = 0
     num_entries_failed: int = 0
     num_queries_aggregator_failures: int = 0
+    num_queries_processed_sequential: int = 0
+    num_temp_tables_detected: int = 0
+    temp_table_patterns_used: List[str] = field(default_factory=list)
+    peak_memory_usage_mb: float = 0.0
 
     sql_aggregator: Optional[SqlAggregatorReport] = None
+    schema_resolver_report: Optional[SchemaResolverReport] = None
 
 
 @platform_name("SQL Queries", id="sql-queries")
@@ -115,6 +141,18 @@ class SqlQueriesSource(Source):
     - upstream_tables (optional): string[] - Fallback list of tables the query reads from,
      used if the query can't be parsed.
 
+    **Lazy Schema Loading**:
+    - Fetches schemas on-demand during query parsing instead of bulk loading all schemas upfront
+    - Caches fetched schemas for future lookups to avoid repeated network requests
+    - Reduces initial startup time and memory usage significantly
+    - Automatically handles large platforms efficiently without memory issues
+
+    **Query Processing**:
+    - Loads the entire query file into memory at once
+    - Processes all queries sequentially before generating metadata work units
+    - Preserves temp table mappings and lineage relationships to ensure consistent lineage tracking
+    - Query deduplication is handled automatically by the SQL parsing aggregator
+
     ### 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:
@@ -124,6 +162,12 @@ class SqlQueriesSource(Source):
 
     Note: Incremental lineage only applies to UpstreamLineage aspects. Other aspects like queries and usage
     statistics will still be emitted normally.
+
+    ### Temporary Table Support
+    For platforms like Athena that don't have native temporary tables, you can use the `temp_table_patterns`
+    configuration to specify regex patterns that identify fake temporary tables. This allows the source to
+    process these tables like other sources that support native temp tables, enabling proper lineage tracking
+    across temporary table operations.
     """
 
     schema_resolver: Optional[SchemaResolver]
@@ -141,13 +185,19 @@ class SqlQueriesSource(Source):
         self.report = SqlQueriesSourceReport()
 
         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(
+            # Create schema resolver report for tracking
+            self.report.schema_resolver_report = SchemaResolverReport()
+
+            # Use lazy loading - schemas will be fetched on-demand and cached
+            logger.info(
+                "Using lazy schema loading - schemas will be fetched on-demand and cached"
+            )
+            self.schema_resolver = SchemaResolver(
                 platform=self.config.platform,
                 platform_instance=self.config.platform_instance,
                 env=self.config.env,
+                graph=self.graph,
+                report=self.report.schema_resolver_report,
             )
         else:
             self.schema_resolver = None
@@ -156,7 +206,9 @@ class SqlQueriesSource(Source):
             platform=self.config.platform,
             platform_instance=self.config.platform_instance,
             env=self.config.env,
-            schema_resolver=self.schema_resolver,
+            schema_resolver=cast(SchemaResolver, self.schema_resolver)
+            if self.schema_resolver
+            else None,
             eager_graph_load=False,
             generate_lineage=True,  # TODO: make this configurable
             generate_queries=True,  # TODO: make this configurable
@@ -165,7 +217,9 @@ class SqlQueriesSource(Source):
             generate_usage_statistics=True,
             generate_operations=True,  # TODO: make this configurable
             usage_config=self.config.usage,
-            is_temp_table=None,
+            is_temp_table=self.is_temp_table
+            if self.config.temp_table_patterns
+            else None,
             is_allowed_table=None,
             format_queries=False,
         )
@@ -193,20 +247,73 @@ class SqlQueriesSource(Source):
     ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         logger.info(f"Parsing queries from {os.path.basename(self.config.query_file)}")
 
+        logger.info("Processing all queries in batch mode")
+        yield from self._process_queries_batch()
+
+    def _process_queries_batch(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
+        """Process all queries in memory (original behavior)."""
         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)
+            logger.info("Using sequential processing")
+            self._process_queries_sequential(queries)
 
         with self.report.new_stage("Generating metadata work units"):
             logger.info("Generating workunits from SQL parsing aggregator")
-            yield from self.aggregator.gen_metadata()
+            yield from auto_workunit(self.aggregator.gen_metadata())
 
-    def _parse_query_file(self) -> Iterable["QueryEntry"]:
-        """Parse the query file and yield QueryEntry objects."""
+    def _is_s3_uri(self, path: str) -> bool:
+        """Check if the path is an S3 URI."""
+        return path.startswith("s3://")
+
+    def _parse_s3_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse query file from S3 using smart_open."""
+        if not self.config.aws_config:
+            raise ValueError("AWS configuration required for S3 file access")
+
+        logger.info(f"Reading query file from S3: {self.config.query_file}")
+
+        try:
+            # Use smart_open for efficient S3 streaming, similar to S3FileSystem
+            s3_client = self.config.aws_config.get_s3_client()
+
+            with smart_open.open(
+                self.config.query_file, mode="r", transport_params={"client": s3_client}
+            ) as file_stream:
+                for line in file_stream:
+                    if line.strip():
+                        try:
+                            query_dict = json.loads(line, strict=False)
+                            entry = QueryEntry.create(query_dict, config=self.config)
+                            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 from S3"
+                                )
+                            yield entry
+                        except Exception as e:
+                            self.report.num_entries_failed += 1
+                            self.report.warning(
+                                title="Error processing query from S3",
+                                message="Query skipped due to parsing error",
+                                context=line.strip(),
+                                exc=e,
+                            )
+        except Exception as e:
+            self.report.warning(
+                title="Error reading S3 file",
+                message="Failed to read S3 file",
+                context=self.config.query_file,
+                exc=e,
+            )
+            raise
+
+    def _parse_local_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse local query file (existing logic)."""
         with open(self.config.query_file) as f:
             for line in f:
                 try:
@@ -227,6 +334,30 @@ class SqlQueriesSource(Source):
                         exc=e,
                     )
 
+    def _parse_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse the query file and yield QueryEntry objects."""
+        if self._is_s3_uri(self.config.query_file):
+            yield from self._parse_s3_query_file()
+        else:
+            yield from self._parse_local_query_file()
+
+    def _process_queries_sequential(self, queries: List["QueryEntry"]) -> None:
+        """Process queries sequentially."""
+        total_queries = len(queries)
+        logger.info(f"Processing {total_queries} queries sequentially")
+
+        # Process each query sequentially
+        for i, query_entry in enumerate(queries):
+            self._add_query_to_aggregator(query_entry)
+            self.report.num_queries_processed_sequential += 1
+
+            # Simple progress reporting every 1000 queries
+            if (i + 1) % 1000 == 0:
+                progress_pct = ((i + 1) / total_queries) * 100
+                logger.info(
+                    f"Processed {i + 1}/{total_queries} queries ({progress_pct:.1f}%)"
+                )
+
     def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
         """Add a query to the SQL parsing aggregator."""
         try:
@@ -285,6 +416,24 @@ class SqlQueriesSource(Source):
                 exc=e,
             )
 
+    def is_temp_table(self, name: str) -> bool:
+        """Check if a table name matches any of the configured temp table patterns."""
+        if not self.config.temp_table_patterns:
+            return False
+
+        try:
+            for pattern in self.config.temp_table_patterns:
+                if re.match(pattern, name, flags=re.IGNORECASE):
+                    logger.debug(
+                        f"Table '{name}' matched temp table pattern: {pattern}"
+                    )
+                    self.report.num_temp_tables_detected += 1
+                    return True
+        except re.error as e:
+            logger.warning(f"Invalid regex pattern '{pattern}': {e}")
+
+        return False
+
 
 class QueryEntry(BaseModel):
     query: str
@@ -301,19 +450,22 @@ class QueryEntry(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    @validator("timestamp", pre=True)
-    def parse_timestamp(cls, v):
+    @field_validator("timestamp", mode="before")
+    @classmethod
+    def parse_timestamp(cls, v: Any) -> Any:
         return None if v is None else parse_user_datetime(str(v))
 
-    @validator("user", pre=True)
-    def parse_user(cls, v):
+    @field_validator("user", mode="before")
+    @classmethod
+    def parse_user(cls, v: Any) -> Any:
         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):
+    @field_validator("downstream_tables", "upstream_tables", mode="before")
+    @classmethod
+    def parse_tables(cls, v: Any) -> Any:
         if not v:
             return []
 

datahub/ingestion/source/state/checkpoint.py

@@ -163,7 +163,7 @@ class Checkpoint(Generic[StateType]):
         )
         state_as_dict["version"] = checkpoint_aspect.state.formatVersion
         state_as_dict["serde"] = checkpoint_aspect.state.serde
-        return state_class.parse_obj(state_as_dict)
+        return state_class.model_validate(state_as_dict)
 
     @staticmethod
     def _from_base85_json_bytes(
@@ -179,7 +179,7 @@ class Checkpoint(Generic[StateType]):
         state_as_dict = json.loads(state_uncompressed.decode("utf-8"))
         state_as_dict["version"] = checkpoint_aspect.state.formatVersion
         state_as_dict["serde"] = checkpoint_aspect.state.serde
-        return state_class.parse_obj(state_as_dict)
+        return state_class.model_validate(state_as_dict)
 
     def to_checkpoint_aspect(
         self, max_allowed_state_size: int

datahub/ingestion/source/state/entity_removal_state.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple, Type
 
 import pydantic
+from pydantic import model_validator
 
 from datahub.emitter.mce_builder import make_assertion_urn, make_container_urn
 from datahub.ingestion.source.state.checkpoint import CheckpointStateBase
@@ -59,7 +60,7 @@ def pydantic_state_migrator(mapping: Dict[str, str]) -> "V1RootValidator":
 
         return values
 
-    return pydantic.root_validator(pre=True, allow_reuse=True)(_validate_field_rename)
+    return model_validator(mode="before")(_validate_field_rename)
 
 
 class GenericCheckpointState(CheckpointStateBase):

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Generic, Optional, Type, TypeVar
 
 import pydantic
-from pydantic import root_validator
+from pydantic import model_validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import (
@@ -73,14 +73,14 @@ class StatefulIngestionConfig(ConfigModel):
         description="If set to True, ignores the current checkpoint state.",
     )
 
-    @pydantic.root_validator(skip_on_failure=True)
-    def validate_config(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("enabled"):
-            if values.get("state_provider") is None:
-                values["state_provider"] = DynamicTypedStateProviderConfig(
+    @model_validator(mode="after")
+    def validate_config(self) -> "StatefulIngestionConfig":
+        if self.enabled:
+            if self.state_provider is None:
+                self.state_provider = DynamicTypedStateProviderConfig(
                     type="datahub", config={}
                 )
-        return values
+        return self
 
 
 CustomConfig = TypeVar("CustomConfig", bound=StatefulIngestionConfig)
@@ -110,17 +110,19 @@ class StatefulLineageConfigMixin(ConfigModel):
         "store_last_lineage_extraction_timestamp", "enable_stateful_lineage_ingestion"
     )
 
-    @root_validator(skip_on_failure=True)
-    def lineage_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_lineage_ingestion"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_lineage_ingestion config option as well"
-                )
-                values["enable_stateful_lineage_ingestion"] = False
-
-        return values
+    @model_validator(mode="after")
+    def lineage_stateful_option_validator(self) -> "StatefulLineageConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_lineage_ingestion", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_lineage_ingestion config option as well"
+                    )
+                    self.enable_stateful_lineage_ingestion = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful lineage validation due to: {e}")
+        return self
 
 
 class StatefulProfilingConfigMixin(ConfigModel):
@@ -135,16 +137,19 @@ class StatefulProfilingConfigMixin(ConfigModel):
         "store_last_profiling_timestamps", "enable_stateful_profiling"
     )
 
-    @root_validator(skip_on_failure=True)
-    def profiling_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_profiling"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_profiling config option as well"
-                )
-                values["enable_stateful_profiling"] = False
-        return values
+    @model_validator(mode="after")
+    def profiling_stateful_option_validator(self) -> "StatefulProfilingConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_profiling", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_profiling config option as well"
+                    )
+                    self.enable_stateful_profiling = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful profiling validation due to: {e}")
+        return self
 
 
 class StatefulUsageConfigMixin(BaseTimeWindowConfig):
@@ -161,16 +166,21 @@ class StatefulUsageConfigMixin(BaseTimeWindowConfig):
         "store_last_usage_extraction_timestamp", "enable_stateful_usage_ingestion"
     )
 
-    @root_validator(skip_on_failure=True)
-    def last_usage_extraction_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_usage_ingestion"):
-                logger.warning(
-                    "Stateful ingestion is disabled, disabling enable_stateful_usage_ingestion config option as well"
-                )
-                values["enable_stateful_usage_ingestion"] = False
-        return values
+    @model_validator(mode="after")
+    def last_usage_extraction_stateful_option_validator(
+        self,
+    ) -> "StatefulUsageConfigMixin":
+        try:
+            sti = getattr(self, "stateful_ingestion", None)
+            if not sti or not getattr(sti, "enabled", False):
+                if getattr(self, "enable_stateful_usage_ingestion", False):
+                    logger.warning(
+                        "Stateful ingestion is disabled, disabling enable_stateful_usage_ingestion config option as well"
+                    )
+                    self.enable_stateful_usage_ingestion = False
+        except (AttributeError, RecursionError) as e:
+            logger.debug(f"Skipping stateful usage validation due to: {e}")
+        return self
 
 
 class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
@@ -185,16 +195,16 @@ class StatefulTimeWindowConfigMixin(BaseTimeWindowConfig):
         "and queries together from a single audit log and uses a unified time window.",
     )
 
-    @root_validator(skip_on_failure=True)
-    def time_window_stateful_option_validator(cls, values: Dict) -> Dict:
-        sti = values.get("stateful_ingestion")
-        if not sti or not sti.enabled:
-            if values.get("enable_stateful_time_window"):
+    @model_validator(mode="after")
+    def time_window_stateful_option_validator(self) -> "StatefulTimeWindowConfigMixin":
+        sti = getattr(self, "stateful_ingestion", None)
+        if not sti or not getattr(sti, "enabled", False):
+            if getattr(self, "enable_stateful_time_window", False):
                 logger.warning(
                     "Stateful ingestion is disabled, disabling enable_stateful_time_window config option as well"
                 )
-                values["enable_stateful_time_window"] = False
-        return values
+                self.enable_stateful_time_window = False
+        return self
 
 
 @dataclass

datahub/ingestion/source/state_provider/datahub_ingestion_checkpointing_provider.py

@@ -40,7 +40,7 @@ class DatahubIngestionCheckpointingProvider(IngestionCheckpointingProviderBase):
     def create(
         cls, config_dict: Dict[str, Any], ctx: PipelineContext
     ) -> "DatahubIngestionCheckpointingProvider":
-        config = DatahubIngestionStateProviderConfig.parse_obj(config_dict)
+        config = DatahubIngestionStateProviderConfig.model_validate(config_dict)
         if config.datahub_api is not None:
             return cls(DataHubGraph(config.datahub_api))
         elif ctx.graph:

datahub/ingestion/source/state_provider/file_ingestion_checkpointing_provider.py

@@ -32,7 +32,7 @@ class FileIngestionCheckpointingProvider(IngestionCheckpointingProviderBase):
     def create(
         cls, config_dict: Dict[str, Any], ctx: PipelineContext
     ) -> "FileIngestionCheckpointingProvider":
-        config = FileIngestionStateProviderConfig.parse_obj(config_dict)
+        config = FileIngestionStateProviderConfig.model_validate(config_dict)
         return cls(config)
 
     def get_latest_checkpoint(

datahub/ingestion/source/superset.py

@@ -9,7 +9,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 import dateutil.parser as dp
 import requests
 import sqlglot
-from pydantic import BaseModel, root_validator, validator
+from pydantic import BaseModel, field_validator, model_validator
 from pydantic.fields import Field
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
@@ -246,16 +246,16 @@ class SupersetConfig(
         # This is required to allow preset configs to get parsed
         extra = "allow"
 
-    @validator("connect_uri", "display_uri")
-    def remove_trailing_slash(cls, v):
+    @field_validator("connect_uri", "display_uri", mode="after")
+    @classmethod
+    def remove_trailing_slash(cls, v: str) -> str:
         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) -> "SupersetConfig":
+        if self.display_uri is None:
+            self.display_uri = self.connect_uri
+        return self
 
 
 def get_metric_name(metric):

datahub/ingestion/source/tableau/tableau.py

@@ -25,7 +25,7 @@ from urllib.parse import quote, urlparse
 
 import dateutil.parser as dp
 import tableauserverclient as TSC
-from pydantic import root_validator, validator
+from pydantic import field_validator, model_validator
 from pydantic.fields import Field
 from requests.adapters import HTTPAdapter
 from tableauserverclient import (
@@ -257,8 +257,9 @@ class TableauConnectionConfig(ConfigModel):
         description="When enabled, extracts column-level lineage from Tableau Datasources",
     )
 
-    @validator("connect_uri")
-    def remove_trailing_slash(cls, v):
+    @field_validator("connect_uri", mode="after")
+    @classmethod
+    def remove_trailing_slash(cls, v: str) -> str:
         return config_clean.remove_trailing_slashes(v)
 
     def get_tableau_auth(
@@ -652,8 +653,9 @@ class TableauConfig(
         "fetch_size",
     )
 
-    # pre = True because we want to take some decision before pydantic initialize the configuration to default values
-    @root_validator(pre=True)
+    # mode = "before" because we want to take some decision before pydantic initialize the configuration to default values
+    @model_validator(mode="before")
+    @classmethod
     def projects_backward_compatibility(cls, values: Dict) -> Dict:
         # In-place update of the input dict would cause state contamination. This was discovered through test failures
         # in test_hex.py where the same dict is reused.
@@ -683,27 +685,23 @@ class TableauConfig(
 
         return values
 
-    @root_validator(skip_on_failure=True)
-    def validate_config_values(cls, values: Dict) -> Dict:
-        tags_for_hidden_assets = values.get("tags_for_hidden_assets")
-        ingest_tags = values.get("ingest_tags")
+    @model_validator(mode="after")
+    def validate_config_values(self) -> "TableauConfig":
         if (
-            not ingest_tags
-            and tags_for_hidden_assets
-            and len(tags_for_hidden_assets) > 0
+            not self.ingest_tags
+            and self.tags_for_hidden_assets
+            and len(self.tags_for_hidden_assets) > 0
         ):
             raise ValueError(
                 "tags_for_hidden_assets is only allowed with ingest_tags enabled. Be aware that this will overwrite tags entered from the UI."
             )
 
-        use_email_as_username = values.get("use_email_as_username")
-        ingest_owner = values.get("ingest_owner")
-        if use_email_as_username and not ingest_owner:
+        if self.use_email_as_username and not self.ingest_owner:
             raise ValueError(
                 "use_email_as_username requires ingest_owner to be enabled."
             )
 
-        return values
+        return self
 
 
 class WorkbookKey(ContainerKey):

datahub/ingestion/source/unity/azure_auth_config.py

@@ -0,0 +1,15 @@
+from pydantic import Field, SecretStr
+
+from datahub.configuration import ConfigModel
+
+
+class AzureAuthConfig(ConfigModel):
+    client_secret: SecretStr = Field(
+        description="Azure application client secret used for authentication. This is a confidential credential that should be kept secure."
+    )
+    client_id: str = Field(
+        description="Azure application (client) ID. This is the unique identifier for the registered Azure AD application.",
+    )
+    tenant_id: str = Field(
+        description="Azure tenant (directory) ID. This identifies the Azure AD tenant where the application is registered.",
+    )