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

datahub/ingestion/source/snowflake/constants.py

@@ -9,6 +9,8 @@ class SnowflakeCloudProvider(StrEnum):
 
 SNOWFLAKE_DEFAULT_CLOUD = SnowflakeCloudProvider.AWS
 
+DEFAULT_SNOWFLAKE_DOMAIN = "snowflakecomputing.com"
+
 
 class SnowflakeEdition(StrEnum):
     STANDARD = "Standard"
@@ -55,6 +57,7 @@ class SnowflakeObjectDomain(StrEnum):
     ICEBERG_TABLE = "iceberg table"
     STREAM = "stream"
     PROCEDURE = "procedure"
+    DYNAMIC_TABLE = "dynamic table"
 
 
 GENERIC_PERMISSION_ERROR_KEY = "permission-error"

datahub/ingestion/source/snowflake/snowflake_config.py

@@ -1,12 +1,13 @@
 import logging
 from collections import defaultdict
 from dataclasses import dataclass
+from enum import Enum
 from typing import Dict, List, Optional, Set
 
 import pydantic
 from pydantic import Field, root_validator, validator
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
 from datahub.configuration.pattern_utils import UUID_REGEX
 from datahub.configuration.source_common import (
     EnvConfigMixin,
@@ -30,6 +31,7 @@ from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterCo
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulLineageConfigMixin,
     StatefulProfilingConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulUsageConfigMixin,
 )
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
@@ -49,9 +51,15 @@ DEFAULT_TEMP_TABLES_PATTERNS = [
     rf".*\.SEGMENT_{UUID_REGEX}",  # segment
     rf".*\.STAGING_.*_{UUID_REGEX}",  # stitch
     r".*\.(GE_TMP_|GE_TEMP_|GX_TEMP_)[0-9A-F]{8}",  # great expectations
+    r".*\.SNOWPARK_TEMP_TABLE_.+",  # snowpark
 ]
 
 
+class QueryDedupStrategyType(Enum):
+    STANDARD = "STANDARD"
+    NONE = "NONE"
+
+
 class TagOption(StrEnum):
     with_lineage = "with_lineage"
     without_lineage = "without_lineage"
@@ -60,13 +68,10 @@ class TagOption(StrEnum):
 
 @dataclass(frozen=True)
 class DatabaseId:
-    database: str = Field(
-        description="Database created from share in consumer account."
-    )
-    platform_instance: Optional[str] = Field(
-        default=None,
-        description="Platform instance of consumer snowflake account.",
-    )
+    # Database created from share in consumer account
+    database: str
+    # Platform instance of consumer snowflake account
+    platform_instance: Optional[str] = None
 
 
 class SnowflakeShareConfig(ConfigModel):
@@ -154,14 +159,11 @@ class SnowflakeIdentifierConfig(
 
     email_domain: Optional[str] = pydantic.Field(
         default=None,
-        description="Email domain of your organization so users can be displayed on UI appropriately.",
+        description="Email domain of your organization so users can be displayed on UI appropriately. This is used only if we cannot infer email ID.",
     )
 
-    email_as_user_identifier: bool = Field(
-        default=True,
-        description="Format user urns as an email, if the snowflake user's email is set. If `email_domain` is "
-        "provided, generates email addresses for snowflake users with unset emails, based on their "
-        "username.",
+    _email_as_user_identifier = pydantic_removed_field(
+        "email_as_user_identifier",
     )
 
 
@@ -198,6 +200,7 @@ class SnowflakeV2Config(
     SnowflakeUsageConfig,
     StatefulLineageConfigMixin,
     StatefulUsageConfigMixin,
+    StatefulTimeWindowConfigMixin,
     StatefulProfilingConfigMixin,
     ClassificationSourceConfigMixin,
     IncrementalPropertiesConfigMixin,
@@ -212,6 +215,16 @@ class SnowflakeV2Config(
         description="If enabled, populates the ingested views' definitions.",
     )
 
+    fetch_views_from_information_schema: bool = Field(
+        default=False,
+        description="If enabled, uses information_schema.views to fetch view definitions instead of SHOW VIEWS command. "
+        "This alternative method can be more reliable for databases with large numbers of views (> 10K views), as the "
+        "SHOW VIEWS approach has proven unreliable and can lead to missing views in such scenarios. However, this method "
+        "requires OWNERSHIP privileges on views to retrieve their definitions. For views without ownership permissions "
+        "(where VIEW_DEFINITION is null/empty), the system will automatically fall back to using batched SHOW VIEWS queries "
+        "to populate the missing definitions.",
+    )
+
     include_technical_schema: bool = Field(
         default=True,
         description="If enabled, populates the snowflake technical schema and descriptions.",
@@ -232,7 +245,7 @@ class SnowflakeV2Config(
     )
 
     use_queries_v2: bool = Field(
-        default=False,
+        default=True,
         description="If enabled, uses the new queries extractor to extract queries from snowflake.",
     )
     include_queries: bool = Field(
@@ -250,6 +263,11 @@ class SnowflakeV2Config(
         "This is useful if you have a large number of schemas and want to avoid bulk fetching the schema for each table/view.",
     )
 
+    query_dedup_strategy: QueryDedupStrategyType = Field(
+        default=QueryDedupStrategyType.STANDARD,
+        description=f"Experimental: Choose the strategy for query deduplication (default value is appropriate for most use-cases; make sure you understand performance implications before changing it). Allowed values are: {', '.join([s.name for s in QueryDedupStrategyType])}",
+    )
+
     _check_role_grants_removed = pydantic_removed_field("check_role_grants")
     _provision_role_removed = pydantic_removed_field("provision_role")
 
@@ -263,10 +281,11 @@ class SnowflakeV2Config(
         description="If enabled along with `extract_tags`, extracts snowflake's key-value tags as DataHub structured properties instead of DataHub tags.",
     )
 
-    structured_properties_template_cache_invalidation_interval: int = Field(
-        hidden_from_docs=True,
-        default=60,
-        description="Interval in seconds to invalidate the structured properties template cache.",
+    structured_properties_template_cache_invalidation_interval: HiddenFromDocs[int] = (
+        Field(
+            default=60,
+            description="Interval in seconds to invalidate the structured properties template cache.",
+        )
     )
 
     include_external_url: bool = Field(
@@ -315,7 +334,7 @@ class SnowflakeV2Config(
         "to ignore the temporary staging tables created by known ETL tools.",
     )
 
-    rename_upstreams_deny_pattern_to_temporary_table_pattern = pydantic_renamed_field(
+    rename_upstreams_deny_pattern_to_temporary_table_pattern = pydantic_renamed_field(  # type: ignore[pydantic-field]
         "upstreams_deny_pattern", "temporary_tables_pattern"
     )
 
@@ -333,8 +352,7 @@ class SnowflakeV2Config(
     )
 
     # Allows empty containers to be ingested before datasets are added, avoiding permission errors
-    warn_no_datasets: bool = Field(
-        hidden_from_docs=True,
+    warn_no_datasets: HiddenFromDocs[bool] = Field(
         default=False,
         description="If True, warns when no datasets are found during ingestion. If False, ingestion fails when no datasets are found.",
     )
@@ -347,11 +365,32 @@ class SnowflakeV2Config(
 
     pushdown_deny_usernames: List[str] = Field(
         default=[],
-        description="List of snowflake usernames which will not be considered for lineage/usage/queries extraction. "
+        description="List of snowflake usernames (SQL LIKE patterns, e.g., 'SERVICE_%', '%_PROD', 'TEST_USER') which will NOT be considered for lineage/usage/queries extraction. "
         "This is primarily useful for improving performance by filtering out users with extremely high query volumes. "
         "Only applicable if `use_queries_v2` is enabled.",
     )
 
+    pushdown_allow_usernames: List[str] = Field(
+        default=[],
+        description="List of snowflake usernames (SQL LIKE patterns, e.g., 'ANALYST_%', '%_USER', 'MAIN_ACCOUNT') which WILL be considered for lineage/usage/queries extraction. "
+        "This is primarily useful for improving performance by filtering in only specific users. "
+        "Only applicable if `use_queries_v2` is enabled. If not specified, all users not in deny list are included.",
+    )
+
+    push_down_database_pattern_access_history: bool = Field(
+        default=False,
+        description="If enabled, pushes down database pattern filtering to the access_history table for improved performance. "
+        "This filters on the accessed objects in access_history.",
+    )
+
+    additional_database_names_allowlist: List[str] = Field(
+        default=[],
+        description="Additional database names (no pattern matching) to be included in the access_history filter. "
+        "Only applies if push_down_database_pattern_access_history=True. "
+        "These databases will be included in the filter being pushed down regardless of database_pattern settings."
+        "This may be required in the case of _eg_ temporary tables being created in a different database than the ones in the database_name patterns.",
+    )
+
     @validator("convert_urns_to_lowercase")
     def validate_convert_urns_to_lowercase(cls, v):
         if not v:
@@ -440,6 +479,20 @@ class SnowflakeV2Config(
 
         return shares
 
+    @root_validator(pre=False, skip_on_failure=True)
+    def validate_queries_v2_stateful_ingestion(cls, values: Dict) -> Dict:
+        if values.get("use_queries_v2"):
+            if values.get("enable_stateful_lineage_ingestion") or values.get(
+                "enable_stateful_usage_ingestion"
+            ):
+                logger.warning(
+                    "enable_stateful_lineage_ingestion and enable_stateful_usage_ingestion are deprecated "
+                    "when using use_queries_v2=True. These configs only work with the legacy (non-queries v2) extraction path. "
+                    "For queries v2, use enable_stateful_time_window instead to enable stateful ingestion "
+                    "for the unified time window extraction (lineage + usage + operations + queries)."
+                )
+        return values
+
     def outbounds(self) -> Dict[str, Set[DatabaseId]]:
         """
         Returns mapping of

datahub/ingestion/source/snowflake/snowflake_connection.py

@@ -15,13 +15,19 @@ from snowflake.connector.network import (
     OAUTH_AUTHENTICATOR,
 )
 
-from datahub.configuration.common import ConfigModel, ConfigurationError, MetaError
+from datahub.configuration.common import (
+    ConfigModel,
+    ConfigurationError,
+    HiddenFromDocs,
+    MetaError,
+)
 from datahub.configuration.connection_resolver import auto_connection_resolver
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.source.snowflake.constants import (
     CLIENT_PREFETCH_THREADS,
     CLIENT_SESSION_KEEP_ALIVE,
+    DEFAULT_SNOWFLAKE_DOMAIN,
 )
 from datahub.ingestion.source.snowflake.oauth_config import (
     OAuthConfiguration,
@@ -47,8 +53,6 @@ _VALID_AUTH_TYPES: Dict[str, str] = {
     "OAUTH_AUTHENTICATOR_TOKEN": OAUTH_AUTHENTICATOR,
 }
 
-_SNOWFLAKE_HOST_SUFFIX = ".snowflakecomputing.com"
-
 
 class SnowflakePermissionError(MetaError):
     """A permission error has happened"""
@@ -64,7 +68,7 @@ class SnowflakeConnectionConfig(ConfigModel):
         description="Any options specified here will be passed to [SQLAlchemy.create_engine](https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine) as kwargs.",
     )
 
-    scheme: str = "snowflake"
+    scheme: HiddenFromDocs[str] = "snowflake"
     username: Optional[str] = pydantic.Field(
         default=None, description="Snowflake username."
     )
@@ -110,18 +114,25 @@ class SnowflakeConnectionConfig(ConfigModel):
         default=None,
         description="OAuth token from external identity provider. Not recommended for most use cases because it will not be able to refresh once expired.",
     )
+    snowflake_domain: str = pydantic.Field(
+        default=DEFAULT_SNOWFLAKE_DOMAIN,
+        description="Snowflake domain. Use 'snowflakecomputing.com' for most regions or 'snowflakecomputing.cn' for China (cn-northwest-1) region.",
+    )
 
     def get_account(self) -> str:
         assert self.account_id
         return self.account_id
 
-    rename_host_port_to_account_id = pydantic_renamed_field("host_port", "account_id")
+    rename_host_port_to_account_id = pydantic_renamed_field("host_port", "account_id")  # type: ignore[pydantic-field]
 
     @pydantic.validator("account_id")
-    def validate_account_id(cls, account_id: str) -> str:
+    def validate_account_id(cls, account_id: str, values: Dict) -> str:
         account_id = remove_protocol(account_id)
         account_id = remove_trailing_slashes(account_id)
-        account_id = remove_suffix(account_id, _SNOWFLAKE_HOST_SUFFIX)
+        # Get the domain from config, fallback to default
+        domain = values.get("snowflake_domain", DEFAULT_SNOWFLAKE_DOMAIN)
+        snowflake_host_suffix = f".{domain}"
+        account_id = remove_suffix(account_id, snowflake_host_suffix)
         return account_id
 
     @pydantic.validator("authentication_type", always=True)
@@ -311,6 +322,7 @@ class SnowflakeConnectionConfig(ConfigModel):
             warehouse=self.warehouse,
             authenticator=_VALID_AUTH_TYPES.get(self.authentication_type),
             application=_APPLICATION_NAME,
+            host=f"{self.account_id}.{self.snowflake_domain}",
             **connect_args,
         )
 
@@ -324,6 +336,7 @@ class SnowflakeConnectionConfig(ConfigModel):
             role=self.role,
             authenticator=_VALID_AUTH_TYPES.get(self.authentication_type),
             application=_APPLICATION_NAME,
+            host=f"{self.account_id}.{self.snowflake_domain}",
             **connect_args,
         )
 
@@ -337,6 +350,7 @@ class SnowflakeConnectionConfig(ConfigModel):
                 warehouse=self.warehouse,
                 role=self.role,
                 application=_APPLICATION_NAME,
+                host=f"{self.account_id}.{self.snowflake_domain}",
                 **connect_args,
             )
         elif self.authentication_type == "OAUTH_AUTHENTICATOR_TOKEN":
@@ -348,6 +362,7 @@ class SnowflakeConnectionConfig(ConfigModel):
                 warehouse=self.warehouse,
                 role=self.role,
                 application=_APPLICATION_NAME,
+                host=f"{self.account_id}.{self.snowflake_domain}",
                 **connect_args,
             )
         elif self.authentication_type == "OAUTH_AUTHENTICATOR":
@@ -363,6 +378,7 @@ class SnowflakeConnectionConfig(ConfigModel):
                 role=self.role,
                 authenticator=_VALID_AUTH_TYPES.get(self.authentication_type),
                 application=_APPLICATION_NAME,
+                host=f"{self.account_id}.{self.snowflake_domain}",
                 **connect_args,
             )
         else:
@@ -408,7 +424,7 @@ class SnowflakeConnection(Closeable):
             # We often run multiple queries in parallel across multiple threads,
             # so we need to number them to help with log readability.
             query_num = self.get_query_no()
-            logger.info(f"Query #{query_num}: {query}", stacklevel=2)
+            logger.info(f"Query #{query_num}: {query.rstrip()}", stacklevel=2)
             resp = self._connection.cursor(DictCursor).execute(query)
             if resp is not None and resp.rowcount is not None:
                 logger.info(

datahub/ingestion/source/snowflake/snowflake_lineage_v2.py

@@ -2,7 +2,17 @@ import json
 import logging
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Collection, Iterable, List, Optional, Set, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Collection,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+)
 
 from pydantic import BaseModel, Field, validator
 
@@ -44,6 +54,9 @@ from datahub.sql_parsing.sqlglot_utils import get_query_fingerprint
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.time import ts_millis_to_datetime
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1Validator
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 EXTERNAL_LINEAGE = "external_lineage"
@@ -51,7 +64,7 @@ TABLE_LINEAGE = "table_lineage"
 VIEW_LINEAGE = "view_lineage"
 
 
-def pydantic_parse_json(field: str) -> classmethod:
+def pydantic_parse_json(field: str) -> "V1Validator":
     def _parse_from_json(cls: Type, v: Any) -> dict:
         if isinstance(v, str):
             return json.loads(v)
@@ -72,7 +85,7 @@ class ColumnUpstreamJob(BaseModel):
 
 
 class ColumnUpstreamLineage(BaseModel):
-    column_name: Optional[str]
+    column_name: Optional[str] = None
     upstreams: List[ColumnUpstreamJob] = Field(default_factory=list)
 
 
@@ -91,9 +104,9 @@ class Query(BaseModel):
 class UpstreamLineageEdge(BaseModel):
     DOWNSTREAM_TABLE_NAME: str
     DOWNSTREAM_TABLE_DOMAIN: str
-    UPSTREAM_TABLES: Optional[List[UpstreamTableNode]]
-    UPSTREAM_COLUMNS: Optional[List[ColumnUpstreamLineage]]
-    QUERIES: Optional[List[Query]]
+    UPSTREAM_TABLES: Optional[List[UpstreamTableNode]] = None
+    UPSTREAM_COLUMNS: Optional[List[ColumnUpstreamLineage]] = None
+    QUERIES: Optional[List[Query]] = None
 
     _json_upstream_tables = pydantic_parse_json("UPSTREAM_TABLES")
     _json_upstream_columns = pydantic_parse_json("UPSTREAM_COLUMNS")