updated pytests and golden files

datahub-project · Jan 16, 2025 · 5f25f3c · 5f25f3c
1 parent 2968f6c
commit 5f25f3c
Show file tree

Hide file tree

Showing 5 changed files with 1,332 additions and 7 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py b/metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_query.py
@@ -970,8 +970,4 @@ def streams_for_database(
         from_clause = (
             f"""FROM '{stream_pagination_marker}'""" if stream_pagination_marker else ""
         )
-        return f"""SHOW STREAMS IN DATABASE {db_name} LIMIT {limit} {from_clause};"""
-
-    @staticmethod
-    def streams_for_schema(schema_name: str, db_name: str) -> str:
-        return f"""SHOW STREAMS IN SCHEMA {db_name}.{schema_name}"""
+        return f"""SHOW STREAMS IN DATABASE {db_name} LIMIT {limit} {from_clause};"""
diff --git a/metadata-ingestion/tests/integration/snowflake/common.py b/metadata-ingestion/tests/integration/snowflake/common.py
@@ -10,6 +10,7 @@
 
 NUM_TABLES = 10
 NUM_VIEWS = 2
+NUM_STREAMS = 1
 NUM_COLS = 10
 NUM_OPS = 10
 NUM_USAGE = 0
@@ -175,6 +176,7 @@ def default_query_results(  # noqa: C901
     query,
     num_tables=NUM_TABLES,
     num_views=NUM_VIEWS,
+    num_streams=NUM_STREAMS,
     num_cols=NUM_COLS,
     num_ops=NUM_OPS,
     num_usages=NUM_USAGE,
@@ -293,6 +295,28 @@ def default_query_results(  # noqa: C901
             )
             for col_idx in range(1, num_cols + 1)
         ]
+    elif query == SnowflakeQuery.streams_for_database("TEST_DB"):
+        # TODO: Add tests for stream pagination.
+        return [
+            {
+                "created_on": datetime(2021, 6, 8, 0, 0, 0, 0),
+                "name": f"STREAM_{stream_idx}",
+                "database_name": "TEST_DB",
+                "schema_name": "TEST_SCHEMA",
+                "owner": "ACCOUNTADMIN",
+                "comment": f"Comment for Stream {stream_idx}",
+                "table_name": f"TEST_DB.TEST_SCHEMA.TABLE_{stream_idx}",
+                "source_type": "Table",
+                "base_tables": f"TEST_DB.TEST_SCHEMA.TABLE_{stream_idx}",
+                "type": "DELTA",
+                "stale": "false",
+                "mode": "DEFAULT",
+                "stale_after": datetime(2021, 6, 22, 0, 0, 0, 0),
+                "invalid_reason": None,
+                "owner_role_type": "ROLE",
+            }
+            for stream_idx in range(1, num_streams + 1)
+        ]
     elif query in (
         SnowflakeQuery.use_database("TEST_DB"),
         SnowflakeQuery.show_primary_keys_for_schema("TEST_SCHEMA", "TEST_DB"),