acryl-datahub 1.1.0.4rc3__py3-none-any.whl → 1.1.0.5rc2__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -63,7 +63,10 @@ from datahub.sql_parsing.sqlglot_lineage import (
     DownstreamColumnRef,
 )
 from datahub.sql_parsing.sqlglot_utils import get_query_fingerprint
-from datahub.utilities.file_backed_collections import ConnectionWrapper, FileBackedList
+from datahub.utilities.file_backed_collections import (
+    ConnectionWrapper,
+    FileBackedList,
+)
 from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
@@ -243,6 +246,12 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         audit_log_file = self.local_temp_path / "audit_log.sqlite"
         use_cached_audit_log = audit_log_file.exists()
 
+        if self.config.local_temp_path is None:
+            self._exit_stack.callback(lambda: audit_log_file.unlink(missing_ok=True))
+
+        shared_connection = self._exit_stack.enter_context(
+            ConnectionWrapper(audit_log_file)
+        )
         queries: FileBackedList[
             Union[
                 KnownLineageMapping,
@@ -251,27 +260,16 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 TableSwap,
                 ObservedQuery,
             ]
-        ]
+        ] = self._exit_stack.enter_context(FileBackedList(shared_connection))
+
         if use_cached_audit_log:
-            logger.info("Using cached audit log")
-            shared_connection = ConnectionWrapper(audit_log_file)
-            queries = FileBackedList(shared_connection)
+            logger.info(f"Using cached audit log at {audit_log_file}")
         else:
-            audit_log_file.unlink(missing_ok=True)
-
-            shared_connection = ConnectionWrapper(audit_log_file)
-            queries = FileBackedList(shared_connection)
-            entry: Union[
-                KnownLineageMapping,
-                PreparsedQuery,
-                TableRename,
-                TableSwap,
-                ObservedQuery,
-            ]
+            logger.info(f"Fetching audit log into {audit_log_file}")
 
             with self.report.copy_history_fetch_timer:
-                for entry in self.fetch_copy_history():
-                    queries.append(entry)
+                for copy_entry in self.fetch_copy_history():
+                    queries.append(copy_entry)
 
             with self.report.query_log_fetch_timer:
                 for entry in self.fetch_query_log(users):
@@ -281,13 +279,10 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
             for i, query in enumerate(queries):
                 if i % 1000 == 0:
                     logger.info(f"Added {i} query log entries to SQL aggregator")
+
                 self.aggregator.add(query)
 
         yield from auto_workunit(self.aggregator.gen_metadata())
-        if not use_cached_audit_log:
-            queries.close()
-            shared_connection.close()
-            audit_log_file.unlink(missing_ok=True)
 
     def fetch_users(self) -> UsersMapping:
         users: UsersMapping = dict()
@@ -403,8 +398,9 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
 
         # TODO need to map snowflake query types to ours
         query_text: str = res["query_text"]
+        snowflake_query_type: str = res["query_type"]
         query_type: QueryType = SNOWFLAKE_QUERY_TYPE_MAPPING.get(
-            res["query_type"], QueryType.UNKNOWN
+            snowflake_query_type, QueryType.UNKNOWN
         )
 
         direct_objects_accessed = res["direct_objects_accessed"]
@@ -421,7 +417,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                     res["session_id"],
                     timestamp,
                     object_modified_by_ddl,
-                    res["query_type"],
+                    snowflake_query_type,
                 )
             if known_ddl_entry:
                 return known_ddl_entry
@@ -436,6 +432,16 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 res["user_name"], users.get(res["user_name"])
             )
         )
+        extra_info = {
+            "snowflake_query_id": res["query_id"],
+            "snowflake_root_query_id": res["root_query_id"],
+            "snowflake_query_type": res["query_type"],
+            "snowflake_role_name": res["role_name"],
+            "query_duration": res["query_duration"],
+            "rows_inserted": res["rows_inserted"],
+            "rows_updated": res["rows_updated"],
+            "rows_deleted": res["rows_deleted"],
+        }
 
         # There are a couple cases when we'd want to prefer our own SQL parsing
         # over Snowflake's metadata.
@@ -470,6 +476,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                 query_hash=get_query_fingerprint(
                     query_text, self.identifiers.platform, fast=True
                 ),
+                extra_info=extra_info,
             )
 
         upstreams = []
@@ -556,6 +563,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
             timestamp=timestamp,
             session_id=res["session_id"],
             query_type=query_type,
+            extra_info=extra_info,
         )
         return entry
 
@@ -667,7 +675,7 @@ def _build_enriched_query_log_query(
     start_time_millis = int(start_time.timestamp() * 1000)
     end_time_millis = int(end_time.timestamp() * 1000)
 
-    users_filter = ""
+    users_filter = "TRUE"
     if deny_usernames:
         user_not_in = ",".join(f"'{user.upper()}'" for user in deny_usernames)
         users_filter = f"user_name NOT IN ({user_not_in})"
@@ -694,10 +702,10 @@ fingerprinted_queries as (
     FROM
         snowflake.account_usage.query_history
     WHERE
-        query_history.start_time >= to_timestamp_ltz({start_time_millis}, 3)
-        AND query_history.start_time < to_timestamp_ltz({end_time_millis}, 3)
+        query_history.start_time >= to_timestamp_ltz({start_time_millis}, 3) -- {start_time.isoformat()}
+        AND query_history.start_time < to_timestamp_ltz({end_time_millis}, 3) -- {end_time.isoformat()}
         AND execution_status = 'SUCCESS'
-        AND {users_filter or "TRUE"}
+        AND {users_filter}
 )
 , deduplicated_queries as (
     SELECT
@@ -715,6 +723,7 @@ fingerprinted_queries as (
 , raw_access_history AS (
     SELECT
         query_id,
+        root_query_id,
         query_start_time,
         user_name,
         direct_objects_accessed,
@@ -723,9 +732,9 @@ fingerprinted_queries as (
     FROM
         snowflake.account_usage.access_history
     WHERE
-        query_start_time >= to_timestamp_ltz({start_time_millis}, 3)
-        AND query_start_time < to_timestamp_ltz({end_time_millis}, 3)
-        AND {users_filter or "TRUE"}
+        query_start_time >= to_timestamp_ltz({start_time_millis}, 3) -- {start_time.isoformat()}
+        AND query_start_time < to_timestamp_ltz({end_time_millis}, 3) -- {end_time.isoformat()}
+        AND {users_filter}
         AND query_id IN (
             SELECT query_id FROM deduplicated_queries
         )
@@ -734,6 +743,7 @@ fingerprinted_queries as (
     -- TODO: Add table filter clause.
     SELECT
         query_id,
+        root_query_id,
         query_start_time,
         ARRAY_SLICE(
             FILTER(direct_objects_accessed, o -> o:objectDomain IN {SnowflakeQuery.ACCESS_HISTORY_TABLE_VIEW_DOMAINS_FILTER}),
@@ -764,6 +774,7 @@ fingerprinted_queries as (
         q.rows_deleted AS "ROWS_DELETED",
         q.user_name AS "USER_NAME",
         q.role_name AS "ROLE_NAME",
+        a.root_query_id,
         a.direct_objects_accessed,
         a.objects_modified,
         a.object_modified_by_ddl

datahub/ingestion/source/snowflake/snowflake_v2.py

@@ -118,7 +118,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 @capability(

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

@@ -936,25 +936,25 @@ class SQLServerSource(SQLAlchemySource):
         url = self.config.get_sql_alchemy_url()
         logger.debug(f"sql_alchemy_url={url}")
         engine = create_engine(url, **self.config.options)
-        with engine.connect() as conn:
-            if self.config.database and self.config.database != "":
-                inspector = inspect(conn)
-                yield inspector
-            else:
+
+        if self.config.database and self.config.database != "":
+            inspector = inspect(engine)
+            yield inspector
+        else:
+            with engine.begin() as conn:
                 databases = conn.execute(
                     "SELECT name FROM master.sys.databases WHERE name NOT IN \
                   ('master', 'model', 'msdb', 'tempdb', 'Resource', \
                        'distribution' , 'reportserver', 'reportservertempdb'); "
-                )
-                for db in databases:
-                    if self.config.database_pattern.allowed(db["name"]):
-                        url = self.config.get_sql_alchemy_url(current_db=db["name"])
-                        with create_engine(
-                            url, **self.config.options
-                        ).connect() as conn:
-                            inspector = inspect(conn)
-                            self.current_database = db["name"]
-                            yield inspector
+                ).fetchall()
+
+            for db in databases:
+                if self.config.database_pattern.allowed(db["name"]):
+                    url = self.config.get_sql_alchemy_url(current_db=db["name"])
+                    engine = create_engine(url, **self.config.options)
+                    inspector = inspect(engine)
+                    self.current_database = db["name"]
+                    yield inspector
 
     def get_identifier(
         self, *, schema: str, entity: str, inspector: Inspector, **kwargs: Any

datahub/ingestion/source/sql/vertica.py

@@ -116,7 +116,7 @@ class VerticaConfig(BasicSQLAlchemyConfig):
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class VerticaSource(SQLAlchemySource):

datahub/ingestion/source/state/stateful_ingestion_base.py

@@ -179,7 +179,7 @@ class StatefulIngestionReport(SourceReport):
 
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 class StatefulIngestionSourceBase(Source):

datahub/ingestion/source/superset.py

@@ -272,7 +272,7 @@ def get_filter_name(filter_obj):
 @config_class(SupersetConfig)
 @support_status(SupportStatus.CERTIFIED)
 @capability(
-    SourceCapability.DELETION_DETECTION, "Optionally enabled via stateful_ingestion"
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 @capability(SourceCapability.DOMAINS, "Enabled by `domain` config to assign domain_key")
 @capability(SourceCapability.LINEAGE_COARSE, "Supported by default")

datahub/ingestion/source/unity/source.py

@@ -159,7 +159,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
 @capability(SourceCapability.TEST_CONNECTION, "Enabled by default")

datahub/metadata/_internal_schema_classes.py

@@ -4599,6 +4599,9 @@ class FabricTypeClass(object):
     SIT = "SIT"
     """System Integration Testing"""
     
+    SBX = "SBX"
+    """Alternative spelling for sandbox"""
+    
     SANDBOX = "SANDBOX"
     """Designates sandbox fabrics"""
     

datahub/metadata/schema.avsc

@@ -9508,6 +9508,7 @@
               "QA": "Designates quality assurance fabrics",
               "RVW": "Designates review fabrics",
               "SANDBOX": "Designates sandbox fabrics",
+              "SBX": "Alternative spelling for sandbox",
               "SIT": "System Integration Testing",
               "STG": "Designates staging fabrics",
               "TEST": "Designates testing fabrics",
@@ -9531,6 +9532,7 @@
               "PRD",
               "TST",
               "SIT",
+              "SBX",
               "SANDBOX"
             ],
             "doc": "Fabric group type"

datahub/metadata/schemas/ContainerProperties.avsc

@@ -99,6 +99,7 @@
             "QA": "Designates quality assurance fabrics",
             "RVW": "Designates review fabrics",
             "SANDBOX": "Designates sandbox fabrics",
+            "SBX": "Alternative spelling for sandbox",
             "SIT": "System Integration Testing",
             "STG": "Designates staging fabrics",
             "TEST": "Designates testing fabrics",
@@ -122,6 +123,7 @@
             "PRD",
             "TST",
             "SIT",
+            "SBX",
             "SANDBOX"
           ],
           "doc": "Fabric group type"

datahub/metadata/schemas/DataFlowInfo.avsc

@@ -153,6 +153,7 @@
             "QA": "Designates quality assurance fabrics",
             "RVW": "Designates review fabrics",
             "SANDBOX": "Designates sandbox fabrics",
+            "SBX": "Alternative spelling for sandbox",
             "SIT": "System Integration Testing",
             "STG": "Designates staging fabrics",
             "TEST": "Designates testing fabrics",
@@ -176,6 +177,7 @@
             "PRD",
             "TST",
             "SIT",
+            "SBX",
             "SANDBOX"
           ],
           "doc": "Fabric group type"

datahub/metadata/schemas/DataJobInfo.avsc

@@ -219,6 +219,7 @@
             "QA": "Designates quality assurance fabrics",
             "RVW": "Designates review fabrics",
             "SANDBOX": "Designates sandbox fabrics",
+            "SBX": "Alternative spelling for sandbox",
             "SIT": "System Integration Testing",
             "STG": "Designates staging fabrics",
             "TEST": "Designates testing fabrics",
@@ -242,6 +243,7 @@
             "PRD",
             "TST",
             "SIT",
+            "SBX",
             "SANDBOX"
           ],
           "doc": "Fabric group type"

datahub/metadata/schemas/DataProcessKey.avsc

@@ -52,6 +52,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -75,6 +76,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/DatasetKey.avsc

@@ -89,6 +89,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -112,6 +113,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/IcebergWarehouseInfo.avsc

@@ -64,6 +64,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -87,6 +88,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/MLModelDeploymentKey.avsc

@@ -60,6 +60,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -83,6 +84,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/MLModelGroupKey.avsc

@@ -67,6 +67,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -90,6 +91,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/MLModelKey.avsc

@@ -81,6 +81,7 @@
           "QA": "Designates quality assurance fabrics",
           "RVW": "Designates review fabrics",
           "SANDBOX": "Designates sandbox fabrics",
+          "SBX": "Alternative spelling for sandbox",
           "SIT": "System Integration Testing",
           "STG": "Designates staging fabrics",
           "TEST": "Designates testing fabrics",
@@ -104,6 +105,7 @@
           "PRD",
           "TST",
           "SIT",
+          "SBX",
           "SANDBOX"
         ],
         "doc": "Fabric group type"

datahub/metadata/schemas/MetadataChangeEvent.avsc

@@ -2430,6 +2430,7 @@
                               "QA": "Designates quality assurance fabrics",
                               "RVW": "Designates review fabrics",
                               "SANDBOX": "Designates sandbox fabrics",
+                              "SBX": "Alternative spelling for sandbox",
                               "SIT": "System Integration Testing",
                               "STG": "Designates staging fabrics",
                               "TEST": "Designates testing fabrics",
@@ -2453,6 +2454,7 @@
                               "PRD",
                               "TST",
                               "SIT",
+                              "SBX",
                               "SANDBOX"
                             ],
                             "doc": "Fabric group type"

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -58,6 +58,7 @@ from datahub.sql_parsing.tool_meta_extractor import (
     ToolMetaExtractorReport,
 )
 from datahub.utilities.cooperative_timeout import CooperativeTimeoutError
+from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.file_backed_collections import (
     ConnectionWrapper,
     FileBackedDict,
@@ -140,6 +141,7 @@ class QueryMetadata:
 
     used_temp_tables: bool = True
 
+    extra_info: Optional[dict] = None
     origin: Optional[Urn] = None
 
     def make_created_audit_stamp(self) -> models.AuditStampClass:
@@ -263,7 +265,7 @@ class PreparsedQuery:
     query_type_props: QueryTypeProps = dataclasses.field(
         default_factory=lambda: QueryTypeProps()
     )
-    # Use this to store addtitional key-value information about query for debugging
+    # Use this to store additional key-value information about the query for debugging.
     extra_info: Optional[dict] = None
     origin: Optional[Urn] = None
 
@@ -948,6 +950,7 @@ class SqlParsingAggregator(Closeable):
                 column_usage=parsed.column_usage or {},
                 confidence_score=parsed.confidence_score,
                 used_temp_tables=session_has_temp_tables,
+                extra_info=parsed.extra_info,
                 origin=parsed.origin,
             )
         )
@@ -1706,7 +1709,7 @@ class SqlParsingAggregator(Closeable):
         )
 
         merged_query_text = ";\n\n".join(
-            [q.formatted_query_string for q in ordered_queries]
+            deduplicate_list([q.formatted_query_string for q in ordered_queries])
         )
 
         resolved_query = dataclasses.replace(