fix(ingest): do not cache temporary tables schema resolvers

metadata-ingestion/src/datahub/ingestion/source/dbt/dbt_common.py

@@ -1273,7 +1273,10 @@ def _parse_cll(
             )
             return SqlParsingResult.make_from_error(e)
 
-        sql_result = sqlglot_lineage(preprocessed_sql, schema_resolver=schema_resolver)
+        sql_result = sqlglot_lineage(
+            preprocessed_sql,  # type:ignore
+            schema_resolver=schema_resolver,
+        )
         if sql_result.debug_info.table_error:
             self.report.sql_parser_table_errors += 1
             logger.info(

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

@@ -39,10 +39,12 @@
     ColumnRef,
     DownstreamColumnRef,
     SqlParsingResult,
+    _sqlglot_lineage_cached,
     infer_output_schema,
     sqlglot_lineage,
 )
 from datahub.sql_parsing.sqlglot_utils import (
+    _parse_statement,
     generate_hash,
     get_query_fingerprint,
     try_format_query,
@@ -222,6 +224,9 @@ class SqlAggregatorReport(Report):
     sql_parsing_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
     sql_fingerprinting_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
     sql_formatting_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
+    sql_parsing_cache_stats: dict = dataclasses.field(default_factory=dict)
+    parse_statement_cache_stats: dict = dataclasses.field(default_factory=dict)
+    format_query_cache_stats: dict = dataclasses.field(default_factory=dict)
 
     # Other lineage loading metrics.
     num_known_query_lineage: int = 0
@@ -272,6 +277,10 @@ def compute_stats(self) -> None:
         self.num_temp_sessions = len(self._aggregator._temp_lineage_map)
         self.num_inferred_temp_schemas = len(self._aggregator._inferred_temp_schemas)
 
+        self.sql_parsing_cache_stats = _sqlglot_lineage_cached.cache_info()._asdict()
+        self.parse_statement_cache_stats = _parse_statement.cache_info()._asdict()
+        self.format_query_cache_stats = try_format_query.cache_info()._asdict()
+
         return super().compute_stats()
 
 

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

@@ -1020,8 +1020,7 @@ def _sqlglot_lineage_inner(
     )
 
 
-@functools.lru_cache(maxsize=SQL_PARSE_RESULT_CACHE_SIZE)
-def sqlglot_lineage(
+def _sqlglot_lineage_nocache(
     sql: str,
     schema_resolver: SchemaResolverInterface,
     default_db: Optional[str] = None,
@@ -1091,6 +1090,28 @@ def sqlglot_lineage(
         return SqlParsingResult.make_from_error(e)
 
 
+_sqlglot_lineage_cached = functools.lru_cache(maxsize=SQL_PARSE_RESULT_CACHE_SIZE)(
+    _sqlglot_lineage_nocache
+)
+
+
+def sqlglot_lineage(
+    sql: str,
+    schema_resolver: SchemaResolverInterface,
+    default_db: Optional[str] = None,
+    default_schema: Optional[str] = None,
+    default_dialect: Optional[str] = None,
+) -> SqlParsingResult:
+    if schema_resolver.includes_temp_tables():
+        return _sqlglot_lineage_nocache(
+            sql, schema_resolver, default_db, default_schema, default_dialect
+        )
+    else:
+        return _sqlglot_lineage_cached(
+            sql, schema_resolver, default_db, default_schema, default_dialect
+        )
+
+
 def create_lineage_sql_parsed_result(
     query: str,
     default_db: Optional[str],