fix(promql): project metric identity at native query boundaries

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/frontend/src/instance.rs

@@ -79,6 +79,7 @@ use promql_parser::label::Matcher;
 use query::QueryEngineRef;
 use query::metrics::OnDone;
 use query::parser::{PromQuery, QueryStatement};
+use query::promql::planner::PromPlanner;
 use query::query_engine::DescribeResult;
 use query::query_engine::options::{QueryOptions, validate_catalog_and_schema};
 use servers::batcher::logical_table::LogicalTablePendingRowsBatcher;
@@ -1074,12 +1075,16 @@ impl Instance {
             return vec![error::SuspendedSnafu {}.fail()];
         }
 
-        // check will be done in prometheus handler's do_query
-        let result = PrometheusHandler::do_query(self, query, query_ctx)
-            .await
-            .with_context(|_| ExecutePromqlSnafu {
-                query: format!("{query:?}"),
-            });
+        // The PromQL query API is tabular: the internal metric-name column kept
+        // for the Prometheus HTTP API is stripped from the plan before execution.
+        let result = async {
+            let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
+            self.do_query_parsed_inner(query, query_ctx, true).await
+        }
+        .await
+        .with_context(|_| ExecutePromqlSnafu {
+            query: format!("{query:?}"),
+        });
         vec![result]
     }
 
@@ -1445,25 +1450,19 @@ impl Instance {
             PermissionTableTargets::Unresolved
         })
     }
-}
 
-#[async_trait]
-impl PrometheusHandler for Instance {
+    /// Runs an already parsed PromQL query: privilege checks, target permissions,
+    /// planning, interception, execution, cancellation and post-execution.
+    ///
+    /// `tabular` marks the query as one whose result is a table rather than the
+    /// Prometheus HTTP API response, so the internal metric-name column kept for
+    /// the latter is stripped from the plan before execution.
     #[tracing::instrument(skip_all)]
-    async fn do_query(
-        &self,
-        query: &PromQuery,
-        query_ctx: QueryContextRef,
-    ) -> server_error::Result<Output> {
-        let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
-        self.do_query_parsed(query, query_ctx).await
-    }
-
-    #[tracing::instrument(skip_all)]
-    async fn do_query_parsed(
+    async fn do_query_parsed_inner(
         &self,
         query: ParsedPromQuery,
         query_ctx: QueryContextRef,
+        tabular: bool,
     ) -> server_error::Result<Output> {
         let interceptor = self
             .plugins
@@ -1498,6 +1497,14 @@ impl PrometheusHandler for Instance {
 
         interceptor.pre_execute(&query, &eval_stmt.expr, Some(&plan), query_ctx.clone())?;
 
+        let plan = if tabular {
+            PromPlanner::strip_metric_name_column(plan)
+                .map_err(BoxedError::new)
+                .context(ExecuteQuerySnafu)?
+        } else {
+            plan
+        };
+
         // Take the EvalStmt from the original QueryStatement and use it to create the CatalogQueryStatement.
         let query_statement = if let QueryStatement::Promql(eval_stmt, alias) = stmt {
             CatalogQueryStatement::Promql(eval_stmt, alias)
@@ -1550,6 +1557,28 @@ impl PrometheusHandler for Instance {
         let output = Output { data, meta };
         Ok(interceptor.post_execute(output, query_ctx)?)
     }
+}
+
+#[async_trait]
+impl PrometheusHandler for Instance {
+    #[tracing::instrument(skip_all)]
+    async fn do_query(
+        &self,
+        query: &PromQuery,
+        query_ctx: QueryContextRef,
+    ) -> server_error::Result<Output> {
+        let query = ParsedPromQuery::parse(query.clone(), &query_ctx)?;
+        self.do_query_parsed(query, query_ctx).await
+    }
+
+    #[tracing::instrument(skip_all)]
+    async fn do_query_parsed(
+        &self,
+        query: ParsedPromQuery,
+        query_ctx: QueryContextRef,
+    ) -> server_error::Result<Output> {
+        self.do_query_parsed_inner(query, query_ctx, false).await
+    }
 
     async fn check_query_permission(
         &self,

src/query/src/promql/planner.rs

@@ -6835,6 +6835,46 @@ impl PromPlanner {
         }
     }
 
+    /// The names the labels of a match group are reported under.
+    ///
+    /// A group is built from physical columns, and the marked metric-name identity is one of
+    /// them. The identity is not a label of a series: Prometheus names the metric of a series
+    /// `__name__`, so a group that matches on the identity is reported under that name rather
+    /// than under the column it is carried in. The identity of a column is its own metadata,
+    /// resolved with the column itself: a column of another qualifier is never inferred for it,
+    /// an unresolvable or ambiguous column keeps its label, and `group_exprs` is untouched, so
+    /// matching still compares what it did.
+    fn match_group_labels(
+        schema: &DFSchemaRef,
+        group_exprs: &[DfExpr],
+        group_labels: Vec<String>,
+    ) -> Vec<String> {
+        group_exprs
+            .iter()
+            .zip(group_labels)
+            .map(|(expr, label)| {
+                let DfExpr::Column(column) = expr else {
+                    return label;
+                };
+                let is_metric_name =
+                    schema
+                        .qualified_field_from_column(column)
+                        .is_ok_and(|(_, field)| {
+                            field
+                                .metadata()
+                                .get(PROMQL_FIELD_ROLE_KEY)
+                                .map(String::as_str)
+                                == Some(PROMQL_METRIC_NAME_ROLE)
+                        });
+                if is_metric_name {
+                    METRIC_NAME.to_string()
+                } else {
+                    label
+                }
+            })
+            .collect()
+    }
+
     /// Wrap `plan` in a check that fails the query when a match group holds more than one row at
     /// a timestamp. `group_exprs` are the label columns of the group, resolved against `plan`.
     fn assert_unique_match_group(
@@ -6844,6 +6884,7 @@ impl PromPlanner {
         time_index_expr: DfExpr,
         violation: MatchGroupViolation,
     ) -> Result<LogicalPlan> {
+        let group_labels = Self::match_group_labels(plan.schema(), &group_exprs, group_labels);
         let mut partition_by = group_exprs.clone();
         partition_by.push(time_index_expr);
         // A label may carry the generated name, and the count column has to stay unambiguous