fix(ingest/sql-parsing): attribute query usage to temp-table composite queries

docs/how/updating-datahub.md

@@ -131,6 +131,8 @@ Requirements:
 
 - #18065 **(Ingestion / Redshift)** The Redshift connector now emits per-query popularity statistics (`queryUsageStatistics`) on Query entities, matching Snowflake, BigQuery, and Databricks. This is controlled by a new `include_query_usage_statistics` flag, which defaults to `true` but, like `include_column_usage_stats`, only takes effect when `include_usage_statistics` is enabled (off by default) and requires `lineage_generate_queries` (default `true`). **If you already run Redshift usage ingestion** (`include_usage_statistics: true`), query popularity stats will start being emitted after upgrade with no config change; set `include_query_usage_statistics: false` to opt out. Enabling query usage causes read queries to be parsed by the SQL lineage aggregator, which is heavier than the default `stl_scan` table-usage path.
 
+- #18130 **(Ingestion / SQL parsing)** Fixed per-query popularity statistics (`queryUsageStatistics`) being dropped or misattributed for queries that load data through temporary/staging tables — the common ELT/dbt pattern — affecting every connector that uses the shared SQL parsing aggregator (Snowflake, BigQuery, Redshift, Databricks Unity Catalog). Such a query is represented by a single merged ("composite") Query entity that lineage points to. Previously its usage stats either went missing entirely (when the temporary tables were excluded by table filters) or were emitted on separate per-statement Query entities that lineage never referenced. Usage now lands on the composite Query, the redundant per-statement Query entities are no longer emitted, and operations (`operationType`) on these tables reference the same composite Query as lineage. **No action required** — query popularity for temp-table-fed tables becomes correct after upgrade; if you previously saw usage on stray Query entities, it now consolidates onto the Query that lineage actually links to.
+
 - **(GMS / search indexing)** Inline Elasticsearch indexing and entity-graph cache invalidation on ingest now require explicit `systemMetadata.properties.appSource = ui` (or the sync-index header), not merely a GraphQL request context. `GroupService` and `RoleService` stamp UI source via `AspectUtils.buildSynchronousMetadataChangeProposal`; other GraphQL resolvers that bypass `MutationUtils` revert to async MAE indexing. See [GMS Entity Graph Cache — Invalidation (sync writes)](../deploy/gms-entity-graph-cache.md#invalidation-sync-writes).
 
 - #17376: **(Ingestion / Hex)** Major in-place upgrade of the `hex` connector: upstream lineage (table-level and column-level), Project → Component links, run history (`lastRefreshed`), and optional AI context documents are now extracted directly from Hex REST APIs — no external CLI, warehouse-side ingestion dependency, or query-tag scraping required. See the [Hex connector docs](https://docs.datahub.com/docs/generated/ingestion/sources/hex) for the new `include_lineage`, `use_queried_tables_lineage`, `connection_platform_map`, and `include_context_documents` options.

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

@@ -129,6 +129,22 @@ class ViewDefinition:
     default_schema: Optional[str] = None
 
 
+@dataclasses.dataclass
+class QueryComposition:
+    """The raw component statements merged into a composite query during
+    temp-table resolution."""
+
+    # The base statement: the one that writes the real (non-temp) downstream.
+    # Its usage counts represent how often the pipeline ran.
+    base: QueryId
+    # The merged temp-loading statements that feed the base statement.
+    others: List[QueryId]
+
+    @property
+    def all_queries(self) -> List[QueryId]:
+        return [self.base, *self.others]
+
+
 @dataclasses.dataclass
 class QueryMetadata:
     query_id: QueryId
@@ -149,9 +165,22 @@ class QueryMetadata:
 
     used_temp_tables: bool = True
 
+    # Set only on composite queries produced by temp-table resolution: the raw
+    # component statements that were merged into this query. Tracked so their
+    # usage can be attributed to the composite and so they can be suppressed from
+    # being emitted as separate (orphan) Query entities.
+    composed_of: Optional[QueryComposition] = None
+
     extra_info: Optional[dict] = None
     origin: Optional[Urn] = None
 
+    @property
+    def usage_query_id(self) -> QueryId:
+        """Usage counts are keyed by raw statement fingerprint. For a composite
+        query (temp-table resolution) the emitted query_id is the composite hash,
+        so usage must be looked up under the base component fingerprint instead."""
+        return self.composed_of.base if self.composed_of else self.query_id
+
     def make_created_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=make_ts_millis(self.latest_timestamp) or 0,
@@ -1597,6 +1626,11 @@ def _gen_lineage_for_downstream(
             queries_generated.add(query_id)
 
             query = queries_map[query_id]
+            # A composite query subsumes its raw component statements. Mark those
+            # components generated too, so _gen_remaining_queries doesn't re-emit
+            # them as orphan Query entities carrying the same usage.
+            if query.composed_of:
+                queries_generated.update(query.composed_of.all_queries)
             yield from self._gen_query(query, downstream_urn)
 
     @classmethod
@@ -1679,7 +1713,7 @@ def _gen_query(
             # of users / lastExecutedAt timestamps per bucket.
             user = query.actor
 
-            query_counter = self._query_usage_counts.get(query_id)
+            query_counter = self._query_usage_counts.get(query.usage_query_id)
             if not query_counter:
                 return
 
@@ -1876,13 +1910,24 @@ def _recurse_into_query(
             deduplicate_list([q.formatted_query_string for q in ordered_queries])
         )
 
+        # Preserve the raw component statements so query usage can be attributed
+        # back to them. The base statement (the real downstream writer) carries the
+        # usage counts that reflect how often the pipeline executed.
+        composed_of = QueryComposition(
+            base=base_query.query_id,
+            others=[
+                q.query_id for q in ordered_queries if q.query_id != base_query.query_id
+            ],
+        )
+
         resolved_query = dataclasses.replace(
             base_query,
             query_id=composite_query_id,
             formatted_query_string=merged_query_text,
             upstreams=list(resolved_lineage_info.upstreams),
             column_lineage=list(resolved_lineage_info.column_lineage),
             confidence_score=resolved_lineage_info.confidence_score,
+            composed_of=composed_of,
         )
 
         return resolved_query
@@ -1916,18 +1961,26 @@ def _gen_operation_mcps(
 
         for downstream_urn, query_ids in self._lineage_map.items():
             for query_id in query_ids:
-                yield from self._gen_operation_for_downstream(downstream_urn, query_id)
-
-                # Avoid generating the same query twice.
-                if query_id in queries_generated:
+                # Resolve temp tables so operations reference the same (possibly
+                # composite) Query entity that lineage does - otherwise the raw
+                # component id is referenced but never emitted (it's subsumed by
+                # the composite), leaving a dangling query reference.
+                query = self._resolve_query_with_temp_tables(self._query_map[query_id])
+                yield from self._gen_operation_for_downstream(downstream_urn, query)
+
+                # Avoid generating the same query twice. A composite subsumes its
+                # raw component statements, so mark those generated too.
+                gen_id = query.query_id
+                if gen_id in queries_generated:
                     continue
-                queries_generated.add(query_id)
-                yield from self._gen_query(self._query_map[query_id], downstream_urn)
+                queries_generated.add(gen_id)
+                if query.composed_of:
+                    queries_generated.update(query.composed_of.all_queries)
+                yield from self._gen_query(query, downstream_urn)
 
     def _gen_operation_for_downstream(
-        self, downstream_urn: UrnStr, query_id: QueryId
+        self, downstream_urn: UrnStr, query: QueryMetadata
     ) -> Iterable[MetadataChangeProposalWrapper]:
-        query = self._query_map[query_id]
         if query.latest_timestamp is None:
             return
 
@@ -1948,8 +2001,8 @@ def _gen_operation_for_downstream(
             actor=query.actor.urn() if query.actor else None,
             sourceType=models.OperationSourceTypeClass.DATA_PLATFORM,
             queries=(
-                [self._query_urn(query_id)]
-                if self.can_generate_query(query_id)
+                [self._query_urn(query.query_id)]
+                if self.can_generate_query(query.query_id)
                 else None
             ),
         )