From 8a26152b5cc174e6d42415ea05b0cc7584b45e33 Mon Sep 17 00:00:00 2001 From: EnRaiha <15997552+EnRaiha@users.noreply.github.com> Date: Sat, 19 Sep 2026 12:23:38 +0800 Subject: [PATCH] fix(pgwire): keep the planner's output keys on the extended protocol MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Describe phase built its OutputSchema with lookup_key equal to the field name, and the execute path replaced the planner's columns with it. Two output columns sharing a name then read the same cell: rows are written under the unique keys cell_keys derives (id, id_1), so the second column never saw its own value on Parse/Bind/Execute while simple query stayed correct. Merge instead of replace: lookup keys come from the planner, the one derivation that runs cell_keys; the announced schema supplies the client-facing display names and catalog types. An announced arity that disagrees with the plan keeps the announced columns but derives their keys with cell_keys so reader and writer still agree. Tests: unit coverage for the merge (duplicate names keep id/id_1; arity mismatch falls back to derived keys) and two wire cases — SELECT 1 AS id, 2 AS id and a w/b join of bare ids — each asserting its own cell on the prepared path. --- .../server/pgwire/handler/routing/execute.rs | 126 ++++++++++++++++-- .../tests/wire/cases/pgwire_extended_query.rs | 68 ++++++++++ 2 files changed, 185 insertions(+), 9 deletions(-) diff --git a/nodedb/src/control/server/pgwire/handler/routing/execute.rs b/nodedb/src/control/server/pgwire/handler/routing/execute.rs index f829e435a..c57a6fb0c 100644 --- a/nodedb/src/control/server/pgwire/handler/routing/execute.rs +++ b/nodedb/src/control/server/pgwire/handler/routing/execute.rs @@ -146,16 +146,17 @@ impl NodeDbPgHandler { // An externally-supplied prepared-statement schema (from the Describe // phase) names the columns; otherwise the planner's fresh output - // schema for this statement does. The Control-Plane computed list is - // known only to this statement's plan, so it rides along under the - // Describe-phase columns: the shaper evaluates it before those - // columns project, and a computed alias never renders as NULL. + // schema for this statement does. The two merge, never replace: the + // lookup keys come from the planner — the single derivation that runs + // `cell_keys`, so two output columns sharing a display name keep + // distinct cells (`id`, `id_1`) — while the Describe phase supplies + // the client-facing display names and the catalog types it announced. + // The Control-Plane computed list is known only to this statement's + // plan, so it rides along under the announced columns: the shaper + // evaluates it before those columns project, and a computed alias + // never renders as NULL. let effective_schema_owned = match shaping.projection { - Some(described) => crate::control::server::response_shape::schema::OutputSchema { - columns: described.columns.clone(), - is_star: described.is_star, - cp_computed: output_schema.cp_computed, - }, + Some(described) => effective_output_schema(&output_schema, described), None => output_schema, }; let effective_schema = Some(&effective_schema_owned); @@ -268,3 +269,110 @@ impl NodeDbPgHandler { .await } } + +/// Merge the planner's output schema with the Describe phase's. +/// +/// Lookup keys come from the planner — the single derivation that runs +/// `cell_keys`, so two output columns sharing a display name keep distinct +/// cells (`id`, `id_1`). The Describe phase supplies only what it knows: the +/// client-facing display names and the catalog types it announced. +fn effective_output_schema( + planner: &crate::control::server::response_shape::schema::OutputSchema, + described: &crate::control::server::response_shape::schema::OutputSchema, +) -> crate::control::server::response_shape::schema::OutputSchema { + use crate::control::server::response_shape::project::cell_keys; + use crate::control::server::response_shape::schema::{OutputColumn, OutputSchema}; + + let mut columns = planner.columns.clone(); + if columns.len() == described.columns.len() { + for (column, announced) in columns.iter_mut().zip(described.columns.iter()) { + column.display_name = announced.display_name.clone(); + column.ty = announced.ty.clone(); + } + } else { + // The announced field list and the plan disagree on arity. Keep the + // columns the client was told about, and derive their keys the same + // way the row writer does, so reader and writer still agree. + let names: Vec = described + .columns + .iter() + .map(|column| column.display_name.clone()) + .collect(); + columns = described + .columns + .iter() + .zip(cell_keys(&names)) + .map(|(column, key)| OutputColumn { + display_name: column.display_name.clone(), + lookup_key: key, + ty: column.ty.clone(), + }) + .collect(); + } + OutputSchema { + columns, + is_star: described.is_star, + cp_computed: planner.cp_computed.clone(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control::server::response_shape::schema::{OutputColumn, OutputSchema}; + use crate::control::server::response_shape::types::DdlColType; + + fn column(name: &str, key: &str) -> OutputColumn { + OutputColumn { + display_name: name.to_owned(), + lookup_key: key.to_owned(), + ty: DdlColType::Text, + } + } + + fn schema(columns: Vec) -> OutputSchema { + OutputSchema { + columns, + is_star: false, + cp_computed: Vec::new(), + } + } + + /// The announced schema never replaces the planner's keys: two columns + /// announced as `id` still read `id` and `id_1`. + #[test] + fn announced_display_names_keep_the_planner_keys() { + let planner = schema(vec![column("id", "id"), column("id", "id_1")]); + let described = schema(vec![column("id", "id"), column("id", "id")]); + + let merged = effective_output_schema(&planner, &described); + let keys: Vec<&str> = merged + .columns + .iter() + .map(|column| column.lookup_key.as_str()) + .collect(); + assert_eq!(keys, ["id", "id_1"]); + let names: Vec<&str> = merged + .columns + .iter() + .map(|column| column.display_name.as_str()) + .collect(); + assert_eq!(names, ["id", "id"]); + } + + /// An announced field count that disagrees with the plan falls back to + /// keys derived from the announced names — reader and writer still agree. + #[test] + fn arity_mismatch_derives_announced_keys() { + let planner = schema(vec![column("one", "one")]); + let described = schema(vec![column("id", "id"), column("id", "id")]); + + let merged = effective_output_schema(&planner, &described); + let keys: Vec<&str> = merged + .columns + .iter() + .map(|column| column.lookup_key.as_str()) + .collect(); + assert_eq!(keys, ["id", "id_1"]); + } +} diff --git a/nodedb/tests/wire/cases/pgwire_extended_query.rs b/nodedb/tests/wire/cases/pgwire_extended_query.rs index 1c309641e..cc5bc208f 100644 --- a/nodedb/tests/wire/cases/pgwire_extended_query.rs +++ b/nodedb/tests/wire/cases/pgwire_extended_query.rs @@ -142,6 +142,74 @@ async fn extended_query_pure_constant_projection() { assert_eq!(y, "hi"); } +/// Two output columns announced under the same name must render their own +/// cells: the Describe phase supplies display names, but the lookup keys come +/// from the planner (`cell_keys`), so `id` and `id_1` stay distinct. +#[tokio::test] +async fn extended_query_duplicate_column_names_keep_distinct_cells() { + let server = TestServer::start().await; + + let rows = server + .client + .query("SELECT 1 AS id, 2 AS id", &[]) + .await + .expect("prepared query should succeed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].len(), 2, "both announced columns must be present"); + + let first: i64 = rows[0].get(0); + let second: i64 = rows[0].get(1); + assert_eq!(first, 1); + assert_eq!( + second, 2, + "the second column must read its own cell, not the first" + ); +} + +/// The issue's second repro: a join whose two sides project a bare `id`. +/// Each column must read its own cell through the prepared path. +#[tokio::test] +async fn extended_query_duplicate_join_columns_keep_distinct_cells() { + let server = TestServer::start().await; + server + .exec( + "CREATE COLLECTION w (id STRING PRIMARY KEY, b_id STRING) \ + WITH (engine='document_strict')", + ) + .await + .unwrap(); + server + .exec("CREATE COLLECTION b (id STRING PRIMARY KEY) WITH (engine='document_strict')") + .await + .unwrap(); + server + .exec("INSERT INTO b (id) VALUES ('b1')") + .await + .unwrap(); + server + .exec("INSERT INTO w (id, b_id) VALUES ('w1', 'b1')") + .await + .unwrap(); + + let rows = server + .client + .query("SELECT w.id, b.id FROM w JOIN b ON w.b_id = b.id", &[]) + .await + .expect("prepared join should succeed"); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].len(), 2, "both joined columns must be present"); + + let first: &str = rows[0].get(0); + let second: &str = rows[0].get(1); + assert_eq!(first, "w1"); + assert_eq!( + second, "b1", + "the joined column must read its own cell, not the first" + ); +} + /// Star projection with a parameterised filter must expand to every /// collection column in the row output. #[tokio::test]