acryl-datahub 1.0.0rc16__py3-none-any.whl → 1.0.0rc18__py3-none-any.whl

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -403,6 +403,7 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
                     res["session_id"],
                     res["query_start_time"],
                     object_modified_by_ddl,
+                    res["query_type"],
                 )
             if known_ddl_entry:
                 return known_ddl_entry
@@ -537,40 +538,42 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
         session_id: str,
         timestamp: datetime,
         object_modified_by_ddl: dict,
+        query_type: str,
     ) -> Optional[Union[TableRename, TableSwap]]:
         timestamp = timestamp.astimezone(timezone.utc)
-        if object_modified_by_ddl[
-            "operationType"
-        ] == "ALTER" and object_modified_by_ddl["properties"].get("swapTargetName"):
-            urn1 = self.identifiers.gen_dataset_urn(
+        if (
+            object_modified_by_ddl["operationType"] == "ALTER"
+            and query_type == "RENAME_TABLE"
+            and object_modified_by_ddl["properties"].get("objectName")
+        ):
+            original_un = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
                     object_modified_by_ddl["objectName"]
                 )
             )
 
-            urn2 = self.identifiers.gen_dataset_urn(
+            new_urn = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
-                    object_modified_by_ddl["properties"]["swapTargetName"]["value"]
+                    object_modified_by_ddl["properties"]["objectName"]["value"]
                 )
             )
-
-            return TableSwap(urn1, urn2, query, session_id, timestamp)
+            return TableRename(original_un, new_urn, query, session_id, timestamp)
         elif object_modified_by_ddl[
             "operationType"
-        ] == "RENAME_TABLE" and object_modified_by_ddl["properties"].get("objectName"):
-            original_un = self.identifiers.gen_dataset_urn(
+        ] == "ALTER" and object_modified_by_ddl["properties"].get("swapTargetName"):
+            urn1 = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
                     object_modified_by_ddl["objectName"]
                 )
             )
 
-            new_urn = self.identifiers.gen_dataset_urn(
+            urn2 = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
-                    object_modified_by_ddl["properties"]["objectName"]["value"]
+                    object_modified_by_ddl["properties"]["swapTargetName"]["value"]
                 )
             )
 
-            return TableRename(original_un, new_urn, query, session_id, timestamp)
+            return TableSwap(urn1, urn2, query, session_id, timestamp)
         else:
             self.report.num_ddl_queries_dropped += 1
             return None

datahub/ingestion/source/sql/hive.py

@@ -777,6 +777,7 @@ class HiveSource(TwoTierSQLAlchemySource):
             column,
             inspector,
             pk_constraints,
+            partition_keys=partition_keys,
         )
 
         if self._COMPLEX_TYPE.match(fields[0].nativeDataType) and isinstance(
@@ -849,3 +850,15 @@ class HiveSource(TwoTierSQLAlchemySource):
                 default_db=default_db,
                 default_schema=default_schema,
             )
+
+    def get_partitions(
+        self, inspector: Inspector, schema: str, table: str
+    ) -> Optional[List[str]]:
+        partition_columns: List[dict] = inspector.get_indexes(
+            table_name=table, schema=schema
+        )
+        for partition_column in partition_columns:
+            if partition_column.get("column_names"):
+                return partition_column.get("column_names")
+
+        return []