acryl-datahub 0.15.0.1rc8__py3-none-any.whl → 0.15.0.1rc10__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -66,6 +66,11 @@ from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
 
+# Define a type alias
+UserName = str
+UserEmail = str
+UsersMapping = Dict[UserName, UserEmail]
+
 
 class SnowflakeQueriesExtractorConfig(ConfigModel):
     # TODO: Support stateful ingestion for the time windows.
@@ -114,11 +119,13 @@ class SnowflakeQueriesSourceConfig(
 class SnowflakeQueriesExtractorReport(Report):
     copy_history_fetch_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
     query_log_fetch_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    users_fetch_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
 
     audit_log_load_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
     sql_aggregator: Optional[SqlAggregatorReport] = None
 
     num_ddl_queries_dropped: int = 0
+    num_users: int = 0
 
 
 @dataclass
@@ -225,6 +232,9 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
     def get_workunits_internal(
         self,
     ) -> Iterable[MetadataWorkUnit]:
+        with self.report.users_fetch_timer:
+            users = self.fetch_users()
+
         # TODO: Add some logic to check if the cached audit log is stale or not.
         audit_log_file = self.local_temp_path / "audit_log.sqlite"
         use_cached_audit_log = audit_log_file.exists()
@@ -248,7 +258,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                     queries.append(entry)
 
             with self.report.query_log_fetch_timer:
-                for entry in self.fetch_query_log():
+                for entry in self.fetch_query_log(users):
                     queries.append(entry)
 
         with self.report.audit_log_load_timer:
@@ -263,6 +273,25 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
             shared_connection.close()
             audit_log_file.unlink(missing_ok=True)
 
+    def fetch_users(self) -> UsersMapping:
+        users: UsersMapping = dict()
+        with self.structured_reporter.report_exc("Error fetching users from Snowflake"):
+            logger.info("Fetching users from Snowflake")
+            query = SnowflakeQuery.get_all_users()
+            resp = self.connection.query(query)
+
+            for row in resp:
+                try:
+                    users[row["NAME"]] = row["EMAIL"]
+                    self.report.num_users += 1
+                except Exception as e:
+                    self.structured_reporter.warning(
+                        "Error parsing user row",
+                        context=f"{row}",
+                        exc=e,
+                    )
+        return users
+
     def fetch_copy_history(self) -> Iterable[KnownLineageMapping]:
         # Derived from _populate_external_lineage_from_copy_history.
 
@@ -298,7 +327,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                         yield result
 
     def fetch_query_log(
-        self,
+        self, users: UsersMapping
     ) -> Iterable[Union[PreparsedQuery, TableRename, TableSwap]]:
         query_log_query = _build_enriched_query_log_query(
             start_time=self.config.window.start_time,
@@ -319,7 +348,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
 
                 assert isinstance(row, dict)
                 try:
-                    entry = self._parse_audit_log_row(row)
+                    entry = self._parse_audit_log_row(row, users)
                 except Exception as e:
                     self.structured_reporter.warning(
                         "Error parsing query log row",
@@ -331,7 +360,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                         yield entry
 
     def _parse_audit_log_row(
-        self, row: Dict[str, Any]
+        self, row: Dict[str, Any], users: UsersMapping
     ) -> Optional[Union[TableRename, TableSwap, PreparsedQuery]]:
         json_fields = {
             "DIRECT_OBJECTS_ACCESSED",
@@ -430,9 +459,11 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                     )
                 )
 
-        # TODO: Fetch email addresses from Snowflake to map user -> email
-        # TODO: Support email_domain fallback for generating user urns.
-        user = CorpUserUrn(self.identifiers.snowflake_identifier(res["user_name"]))
+        user = CorpUserUrn(
+            self.identifiers.get_user_identifier(
+                res["user_name"], users.get(res["user_name"])
+            )
+        )
 
         timestamp: datetime = res["query_start_time"]
         timestamp = timestamp.astimezone(timezone.utc)

datahub/ingestion/source/snowflake/snowflake_query.py

@@ -947,4 +947,8 @@ WHERE table_schema='{schema_name}' AND {extra_clause}"""
                 AND METRIC_NAME ilike '{pattern}' escape '{escape_pattern}'
                 ORDER BY MEASUREMENT_TIME ASC;
 
-"""
+            """
+
+    @staticmethod
+    def get_all_users() -> str:
+        return """SELECT name as "NAME", email as "EMAIL" FROM SNOWFLAKE.ACCOUNT_USAGE.USERS"""

datahub/ingestion/source/snowflake/snowflake_usage_v2.py

@@ -342,10 +342,9 @@ class SnowflakeUsageExtractor(SnowflakeCommonMixin, Closeable):
                 filtered_user_counts.append(
                     DatasetUserUsageCounts(
                         user=make_user_urn(
-                            self.get_user_identifier(
+                            self.identifiers.get_user_identifier(
                                 user_count["user_name"],
                                 user_email,
-                                self.config.email_as_user_identifier,
                             )
                         ),
                         count=user_count["total"],
@@ -453,9 +452,7 @@ class SnowflakeUsageExtractor(SnowflakeCommonMixin, Closeable):
             reported_time: int = int(time.time() * 1000)
             last_updated_timestamp: int = int(start_time.timestamp() * 1000)
             user_urn = make_user_urn(
-                self.get_user_identifier(
-                    user_name, user_email, self.config.email_as_user_identifier
-                )
+                self.identifiers.get_user_identifier(user_name, user_email)
             )
 
             # NOTE: In earlier `snowflake-usage` connector this was base_objects_accessed, which is incorrect

datahub/ingestion/source/snowflake/snowflake_utils.py

@@ -300,6 +300,28 @@ class SnowflakeIdentifierBuilder:
     def get_quoted_identifier_for_table(db_name, schema_name, table_name):
         return f'"{db_name}"."{schema_name}"."{table_name}"'
 
+    # Note - decide how to construct user urns.
+    # Historically urns were created using part before @ from user's email.
+    # Users without email were skipped from both user entries as well as aggregates.
+    # However email is not mandatory field in snowflake user, user_name is always present.
+    def get_user_identifier(
+        self,
+        user_name: str,
+        user_email: Optional[str],
+    ) -> str:
+        if user_email:
+            return self.snowflake_identifier(
+                user_email
+                if self.identifier_config.email_as_user_identifier is True
+                else user_email.split("@")[0]
+            )
+        return self.snowflake_identifier(
+            f"{user_name}@{self.identifier_config.email_domain}"
+            if self.identifier_config.email_as_user_identifier is True
+            and self.identifier_config.email_domain is not None
+            else user_name
+        )
+
 
 class SnowflakeCommonMixin(SnowflakeStructuredReportMixin):
     platform = "snowflake"
@@ -315,24 +337,6 @@ class SnowflakeCommonMixin(SnowflakeStructuredReportMixin):
     def identifiers(self) -> SnowflakeIdentifierBuilder:
         return SnowflakeIdentifierBuilder(self.config, self.report)
 
-    # Note - decide how to construct user urns.
-    # Historically urns were created using part before @ from user's email.
-    # Users without email were skipped from both user entries as well as aggregates.
-    # However email is not mandatory field in snowflake user, user_name is always present.
-    def get_user_identifier(
-        self,
-        user_name: str,
-        user_email: Optional[str],
-        email_as_user_identifier: bool,
-    ) -> str:
-        if user_email:
-            return self.identifiers.snowflake_identifier(
-                user_email
-                if email_as_user_identifier is True
-                else user_email.split("@")[0]
-            )
-        return self.identifiers.snowflake_identifier(user_name)
-
     # TODO: Revisit this after stateful ingestion can commit checkpoint
     # for failures that do not affect the checkpoint
     # TODO: Add additional parameters to match the signature of the .warning and .failure methods

datahub/ingestion/source_report/ingestion_stage.py

@@ -42,4 +42,5 @@ class IngestionStageReport:
             self._timer = PerfTimer()
 
         self.ingestion_stage = f"{stage} at {datetime.now(timezone.utc)}"
+        logger.info(f"Stage started: {self.ingestion_stage}")
         self._timer.start()