fix(ingest): fix redshift query urns + reduce memory usage (datahub-p…

…roject#10691)

metadata-ingestion/setup.py

@@ -34,7 +34,7 @@
     "importlib_metadata>=4.0.0; python_version < '3.10'",
     "docker",
     "expandvars>=0.6.5",
-    "avro-gen3==0.7.12",
+    "avro-gen3==0.7.13",
     # "avro-gen3 @ git+https://github.com/acryldata/avro_gen@master#egg=avro-gen3",
     "avro>=1.11.3,<1.12",
     "python-dateutil>=2.8.0",

metadata-ingestion/src/datahub/ingestion/source/redshift/lineage.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Dict, List, Optional, Set, Tuple, Union
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 from urllib.parse import urlparse
 
 import humanfriendly
@@ -661,7 +661,7 @@ def populate_lineage(
         if self.config.resolve_temp_table_in_lineage:
             self._init_temp_table_schema(
                 database=database,
-                temp_tables=self.get_temp_tables(connection=connection),
+                temp_tables=list(self.get_temp_tables(connection=connection)),
             )
 
         populate_calls: List[Tuple[str, LineageCollectorType]] = []
@@ -893,23 +893,19 @@ def _process_table_renames(
 
     def get_temp_tables(
         self, connection: redshift_connector.Connection
-    ) -> List[TempTableRow]:
+    ) -> Iterable[TempTableRow]:
         ddl_query: str = self.queries.temp_table_ddl_query(
             start_time=self.config.start_time,
             end_time=self.config.end_time,
         )
 
         logger.debug(f"Temporary table ddl query = {ddl_query}")
 
-        temp_table_rows: List[TempTableRow] = []
-
         for row in RedshiftDataDictionary.get_temporary_rows(
             conn=connection,
             query=ddl_query,
         ):
-            temp_table_rows.append(row)
-
-        return temp_table_rows
+            yield row
 
     def find_temp_tables(
         self, temp_table_rows: List[TempTableRow], temp_table_names: List[str]

metadata-ingestion/src/datahub/ingestion/source/redshift/lineage_v2.py

@@ -116,7 +116,11 @@ def build(
                     default_schema=self.config.default_schema,
                     session_id=temp_row.session_id,
                     query_timestamp=temp_row.start_time,
-                    is_known_temp_table=True,
+                    # The "temp table" query actually returns all CREATE TABLE statements, even if they
+                    # aren't explicitly a temp table. As such, setting is_known_temp_table=True
+                    # would not be correct. We already have mechanisms to autodetect temp tables,
+                    # so we won't lose anything by not setting it.
+                    is_known_temp_table=False,
                 )
 
         populate_calls: List[Tuple[LineageCollectorType, str, Callable]] = []

metadata-ingestion/src/datahub/ingestion/source/redshift/query.py

@@ -502,7 +502,7 @@ def list_insert_create_queries_sql(
                     usename as username,
                     ddl,
                     sq.query as query_id,
-                    min(si.starttime) as starttime,
+                    min(si.starttime) as timestamp,
                     ANY_VALUE(pid) as session_id
                 from
                     stl_insert as si

metadata-ingestion/src/datahub/sql_parsing/sql_parsing_aggregator.py

@@ -870,6 +870,7 @@ def _query_type_precedence(cls, query_type: str) -> int:
             models.DatasetLineageTypeClass.TRANSFORMED,
         ]
 
+        # Lower value = higher precedence.
         idx = query_precedence.index(query_type)
         if idx == -1:
             return len(query_precedence)
@@ -885,13 +886,17 @@ def _gen_lineage_for_downstream(
         ]
 
         # Sort the queries by highest precedence first, then by latest timestamp.
+        # In case of ties, prefer queries with a known query type.
         # Tricky: by converting the timestamp to a number, we also can ignore the
         # differences between naive and aware datetimes.
         queries = sorted(
+            # Sorted is a stable sort, so in the case of total ties, we want
+            # to prefer the most recently added query.
             reversed(queries),
             key=lambda query: (
                 self._query_type_precedence(query.lineage_type),
                 -(make_ts_millis(query.latest_timestamp) or 0),
+                query.query_type == QueryType.UNKNOWN,
             ),
         )
 

metadata-ingestion/tests/unit/test_redshift_lineage.py

@@ -241,7 +241,7 @@ def test_collapse_temp_lineage():
 
     lineage_extractor._init_temp_table_schema(
         database=lineage_extractor.config.database,
-        temp_tables=lineage_extractor.get_temp_tables(connection=connection),
+        temp_tables=list(lineage_extractor.get_temp_tables(connection=connection)),
     )
 
     lineage_extractor._populate_lineage_map(