From 89f97d6e1875dfe785b335c7d9423dabbb54ccc5 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 12:45:28 -0700 Subject: [PATCH 1/6] feat: ORDER BY CASE --- integration/pgdog.toml | 11 +- .../client/query_engine/route_query.rs | 7 +- pgdog/src/frontend/router/parser/order_by.rs | 69 +++++++++++ pgdog/src/frontend/router/parser/query/mod.rs | 4 +- .../frontend/router/parser/query/select.rs | 86 +------------- .../parser/rewrite/statement/order_by.rs | 36 ++++-- .../parser/rewrite/statement/order_by/case.rs | 105 +++++++++++++++++ .../parser/rewrite/statement/projection.rs | 110 +++++++++++------- 8 files changed, 289 insertions(+), 139 deletions(-) create mode 100644 pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 429b7ffee..918ead29d 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -25,8 +25,8 @@ tls_certificate = "integration/tls/cert.pem" tls_private_key = "integration/tls/key.pem" query_parser_engine = "pg_query_raw" system_catalogs = "omnisharded_sticky" -reload_schema_on_ddl = false -# idle_healthcheck_delay = 50000000 +reload_schema_on_ddl = true +idle_healthcheck_delay = 50000000 unique_id_function = "standard" auth_type = "scram" workers = 4 @@ -42,7 +42,7 @@ message_buffer = 8096 [rewrite] enabled = false shard_key = "ignore" -split_inserts = "error" +split_inserts = "rewrite" # primary_key = "rewrite" non_deterministic_functions = "rewrite" @@ -190,6 +190,11 @@ database = "pgdog_sharded" column = "customer_id" data_type = "bigint" +[[sharded_tables]] +database = "pgdog_sharded" +column = "tenant_id" +data_type = "bigint" + # ------------------------------------------------------------------------------ # ----- Hash Sharded :: VARCHAR ------------------------------------------------ diff --git a/pgdog/src/frontend/client/query_engine/route_query.rs b/pgdog/src/frontend/client/query_engine/route_query.rs index 5a98d733a..f77ec480c 100644 --- a/pgdog/src/frontend/client/query_engine/route_query.rs +++ b/pgdog/src/frontend/client/query_engine/route_query.rs @@ -152,7 +152,6 @@ impl QueryEngine { Ok(()) => { let command = self.router.command(); client_request.route = Some(command.route().clone()); - trace!("routing {:#?} to {:#?}", client_request.messages, command,); projection::finalize_after_route( client_request, @@ -163,6 +162,12 @@ impl QueryEngine { if let Some(rewrite_result) = rewrite_result { rewrite_result.apply_after_route(client_request)?; } + + trace!( + "routing {:#?} to {:#?}", + client_request.messages, + client_request.route(), + ); } Err(RouterError::Parser(ParserError::OmniWriteWithDirective)) => { diff --git a/pgdog/src/frontend/router/parser/order_by.rs b/pgdog/src/frontend/router/parser/order_by.rs index 60f52fb4e..2c4a85709 100644 --- a/pgdog/src/frontend/router/parser/order_by.rs +++ b/pgdog/src/frontend/router/parser/order_by.rs @@ -2,8 +2,12 @@ use std::fmt::Debug; +use pg_raw_parse::{ConstValue, Node, list::CastNodeList, nodes}; + use crate::net::messages::Vector; +use super::{Column, StatementParameters, Value}; + #[derive(Clone, Debug, PartialEq)] pub(crate) enum OrderBy { Asc(usize), @@ -15,6 +19,71 @@ pub(crate) enum OrderBy { } impl OrderBy { + /// Extract result sort keys from the final, possibly rewritten ORDER BY. + pub(crate) fn parse( + sorts: &CastNodeList, + params: Option>, + ) -> Vec { + sorts + .iter() + .filter_map(|sort| { + use pg_raw_parse::raw::SortByDir::*; + + let asc = matches!(sort.sortby_dir, SORTBY_DEFAULT | SORTBY_ASC); + match sort.node() { + Node::A_Const(c) if let Some(ConstValue::Integer(i)) = c.val() => { + Some(if asc { + Self::Asc(i as _) + } else { + Self::Desc(i as _) + }) + } + Node::ColumnRef(column) => { + let name = column.fields().iter().next_back()?.as_str()?; + Some(if asc { + Self::AscColumn(name.into()) + } else { + Self::DescColumn(name.into()) + }) + } + node => Self::parse_vector(node, params), + } + }) + .collect() + } + + /// Extract a vector distance for shard selection before projection rewriting. + pub(crate) fn parse_vector( + node: Node<'_>, + params: Option>, + ) -> Option { + let Node::A_Expr(expr) = node else { + return None; + }; + if expr.kind != nodes::A_Expr_Kind::AEXPR_OP + || expr.name().iter().next().and_then(Node::as_str) != Some("<->") + { + return None; + } + + let mut vector = None; + let mut column = None; + for node in [expr.lexpr(), expr.rexpr()] { + if let Ok(value) = Value::try_from(node) { + match value { + Value::Placeholder(p) => { + vector = params?.parameter((p - 1) as _).ok()??.vector(); + } + Value::Vector(value) => vector = Some(value), + _ => {} + } + } else if let Ok(value) = Column::try_from(node) { + column = Some(value.name); + } + } + Some(Self::AscVectorL2Column(column?.into(), vector?)) + } + /// ORDER BY x ASC pub(crate) fn asc(&self) -> bool { matches!( diff --git a/pgdog/src/frontend/router/parser/query/mod.rs b/pgdog/src/frontend/router/parser/query/mod.rs index aa3a9f089..dcb55aa0c 100644 --- a/pgdog/src/frontend/router/parser/query/mod.rs +++ b/pgdog/src/frontend/router/parser/query/mod.rs @@ -11,12 +11,12 @@ use crate::{ round_robin, sharding::{self, Centroids, ContextBuilder, ShardOrLookup}, }, - net::{messages::Vector, parameter::ParameterValue}, + net::parameter::ParameterValue, plugin::plugins, }; #[cfg(test)] -use crate::net::messages::Bind; +use crate::net::messages::{Bind, Vector}; use super::{ explain_trace::{ExplainRecorder, ExplainSummary}, diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index 7614f3ecf..bf5e5b207 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -188,14 +188,17 @@ impl QueryParser { )); } - let order_by = Self::select_sort(stmt, context.router_context.bind); let from_clause_table_name = stmt.from_clause().first().and_then(|node| match node { Node::RangeVar(r) => Some(r.relname().expect("RangeVar always has relname")), _ => None, }); // Shard by vector in ORDER BY clause. - for order in &order_by { + for order in stmt + .sort_clause() + .iter() + .filter_map(|sort| OrderBy::parse_vector(sort.node(), context.router_context.bind)) + { if let Some((vector, column_name)) = order.vector() { for table in context.sharding_schema.tables.tables() { if &table.column == column_name @@ -314,7 +317,7 @@ impl QueryParser { let query = Route::select( context.shards_calculator.shard().clone(), - order_by, + Vec::new(), aggregates, limit, distinct, @@ -328,81 +331,4 @@ impl QueryParser { .with_advisory_locks(advisory_locks), )) } - - /// Handle the `ORDER BY` clause of a `SELECT` statement. - /// - /// # Arguments - /// - /// * `nodes`: List of parser-generated nodes from the ORDER BY clause. - /// * `params`: Statement parameters, if any. - /// - fn select_sort( - stmt: &nodes::SelectStmt, - params: Option>, - ) -> Vec { - stmt.sort_clause() - .into_iter() - .filter_map(|sort_by| { - use pg_raw_parse::{ - ConstValue, - raw::{A_Expr_Kind::*, SortByDir::*}, - }; - - let asc = matches!(sort_by.sortby_dir, SORTBY_DEFAULT | SORTBY_ASC); - match sort_by.node() { - Node::A_Const(c) if let Some(ConstValue::Integer(i)) = c.val() => { - if asc { - Some(OrderBy::Asc(i as _)) - } else { - Some(OrderBy::Desc(i as _)) - } - } - - Node::ColumnRef(c) => { - // TODO: save the entire column and disambiguate - // when reading data with RowDescription as context. - let col_name = c.fields().into_iter().next_back()?.as_str()?; - if asc { - Some(OrderBy::AscColumn(col_name.into())) - } else { - Some(OrderBy::DescColumn(col_name.into())) - } - } - - Node::A_Expr(e @ nodes::A_Expr { kind: AEXPR_OP, .. }) - if let Some("<->") = e.name().iter().next().and_then(Node::as_str) => - { - let mut vector: Option = None; - let mut column: Option<&str> = None; - - for e in [e.lexpr(), e.rexpr()] { - if let Ok(vec) = Value::try_from(e) { - match vec { - Value::Placeholder(p) => { - if let Ok(param) = params?.parameter((p - 1) as _) { - vector = param?.vector(); - } - } - Value::Vector(vec) => vector = Some(vec), - _ => (), - } - } else if let Ok(col) = Column::try_from(e) { - column = Some(col.name); - } - } - - if let Some(vector) = vector - && let Some(column) = column - { - Some(OrderBy::AscVectorL2Column(column.into(), vector)) - } else { - None - } - } - - _ => None, - } - }) - .collect() - } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs index 25e3f687f..de6318705 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs @@ -4,6 +4,9 @@ use crate::frontend::router::parser::Column; use super::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan}; +mod case; +pub(super) use case::rewrite_cases; + impl OrderBySource { fn name(&self) -> &str { match self { @@ -163,8 +166,9 @@ pub(super) fn rewrite_select<'a>( plan: &mut ProjectionRewritePlan, ) { let mut helpers = Vec::new(); + let mut sort_rewrites = Vec::new(); let mut sort_position = 0; - for sort in select.sort_clause() { + for (position, sort) in select.sort_clause().into_iter().enumerate() { let node = sort.node(); let source = match node { Node::ColumnRef(column) => { @@ -193,15 +197,19 @@ pub(super) fn rewrite_select<'a>( match node { Node::ColumnRef(column) => match unique_output_name(select, column) { Some(name) if source.as_ref().is_some_and(|source| source.name() == name) => {} - Some(name) => push_helper( - plan, - current_sort_position, - source.expect("column sorts always have a source"), - name, - false, - ), + Some(name) => { + sort_rewrites.push((position, name.clone())); + push_helper( + plan, + current_sort_position, + source.expect("column sorts always have a source"), + name, + false, + ); + } None => { let alias = format!("__pgdog_order_col{current_sort_position}"); + sort_rewrites.push((position, alias.clone())); helpers.push(mem.make_res_target( Some(&alias), mem.empty(), @@ -218,6 +226,7 @@ pub(super) fn rewrite_select<'a>( }, Node::A_Expr(_) if source.is_some() => { let alias = format!("__pgdog_order_col{current_sort_position}"); + sort_rewrites.push((position, alias.clone())); helpers.push(mem.make_res_target( Some(&alias), mem.empty(), @@ -235,6 +244,17 @@ pub(super) fn rewrite_select<'a>( } } + let mut rewrites = sort_rewrites.into_iter().peekable(); + for (position, mut sort) in select.sort_clause_mut().into_iter().enumerate() { + if rewrites.peek().is_some_and(|(index, _)| *index == position) { + let (_, alias) = rewrites.next().expect("sort rewrite exists"); + sort.set_node( + mem.make_column_ref(mem.make_list(&[mem.make_string(Some(&alias)).uncast()])) + .uncast(), + ); + } + } + if !helpers.is_empty() { select .target_list_mut() diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs new file mode 100644 index 000000000..786bb297a --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs @@ -0,0 +1,105 @@ +use pg_raw_parse::{Node, make, nodes, walk}; +use std::collections::HashSet; + +use crate::backend::schema::Schema; + +use super::{OrderBySource, ProjectionRewritePlan, push_helper}; + +/// Project row-level CASE sort keys. Aggregation and DISTINCT need their own +/// rewrite rules because adding a target can change the query's semantics. +pub(in crate::frontend::router::parser::rewrite::statement) fn rewrite_cases<'a>( + select: &mut nodes::SelectStmtMut<'a, '_>, + mem: make::MemoryToken<'a>, + schema: &Schema, + plan: &mut ProjectionRewritePlan, +) { + if !select + .sort_clause() + .iter() + .any(|sort| matches!(sort.node(), Node::CaseExpr(_))) + || !select.distinct_clause().is_empty() + || !select.group_clause().is_empty() + || !matches!(select.having_clause(), Node::None) + || select.op != nodes::SetOperation::SETOP_NONE + { + return; + } + + // Walk expressions as well as targets: an aggregate nested in CASE or + // present only in ORDER BY must also prevent this rewrite. Without a + // function catalog, conservatively leave function calls alone. + let mut aggregate = false; + let mut names = HashSet::new(); + walk::walk(Node::SelectStmt(select), |node| match node { + Node::FuncCall(func) => { + let name = func.funcname().iter().next_back().and_then(Node::as_str); + aggregate |= schema.aggregate_functions.is_empty() + || name.is_some_and(|name| schema.aggregate_functions.contains(name)) + || func.agg_star + || func.agg_distinct + || !func.agg_order().is_empty() + || !matches!(func.agg_filter(), Node::None) + || func.over().is_some(); + } + Node::JsonAggConstructor(_) => aggregate = true, + Node::ResTarget(target) => { + if let Some(name) = target.name() { + names.insert(name.to_owned()); + } + } + Node::String(string) => { + if let Some(name) = string.sval() { + names.insert(name.to_owned()); + } + } + _ => {} + }); + if aggregate { + return; + } + + // Include known star expansions, and keep names deterministic so a cached + // prepared variant still matches if its AST is evicted and rebuilt. + names.extend( + schema + .relations + .values() + .flat_map(|relations| relations.values()) + .flat_map(|relation| relation.columns.keys().cloned()), + ); + + let mut helpers = Vec::new(); + for (position, mut sort) in select.sort_clause_mut().into_iter().enumerate() { + if !matches!(sort.node(), Node::CaseExpr(_)) { + continue; + } + + let mut suffix = position; + let alias = loop { + let alias = format!("__pgdog_order_case{suffix}"); + if names.insert(alias.clone()) { + break alias; + } + suffix += 1; + }; + helpers.push(mem.make_res_target( + Some(&alias), + mem.empty(), + mem.make_unique(sort.node()).uncast(), + )); + sort.set_node( + mem.make_column_ref(mem.make_list(&[mem.make_string(Some(&alias)).uncast()])) + .uncast(), + ); + push_helper( + plan, + position, + OrderBySource::Column(alias.clone()), + alias, + true, + ); + } + select + .target_list_mut() + .extend(mem, mem.make_list(&helpers)); +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 772246586..71325f0a6 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -9,10 +9,12 @@ use super::aggregate::{AggregatesRewrite, HelperKind}; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, OrderBy}; +use crate::frontend::router::parser::{ + Aggregate, Cache, ExecuteParams, OrderBy, StatementParameters, +}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; -use pg_raw_parse::{Node, StmtList, make}; +use pg_raw_parse::{Node, Owned, StmtList, list::CastNodeList, make, nodes}; use std::collections::BTreeSet; use std::sync::Arc; @@ -30,7 +32,7 @@ pub(crate) struct OrderByHelper { pub(crate) source: OrderBySource, pub(crate) alias: String, /// False when the SELECT list already has this unique name and we only - /// remap the route. Those columns must stay in the client result. + /// rewrite its ORDER BY reference. Those columns stay in the client result. pub(crate) injected: bool, } @@ -40,19 +42,6 @@ pub(crate) enum OrderBySource { Vector(String), } -impl OrderByHelper { - fn matches(&self, order_by: &OrderBy) -> bool { - match (&self.source, order_by) { - (OrderBySource::Column(source), OrderBy::AscColumn(column)) - | (OrderBySource::Column(source), OrderBy::DescColumn(column)) - | (OrderBySource::Vector(source), OrderBy::AscVectorL2Column(column, _)) => { - source == column - } - _ => false, - } - } -} - #[derive(Debug, Clone, Default, PartialEq)] pub(crate) struct ProjectionRewritePlan { pub(crate) aggregate_helpers: Vec, @@ -79,10 +68,12 @@ impl ProjectionRewritePlan { } } -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct PostRouteRewrite { sql: Arc, plan: ProjectionRewritePlan, + /// Cache syntax rather than parameter-dependent sort values. + sort_clause: Owned>, } pub(crate) fn finalize_after_route( @@ -90,18 +81,30 @@ pub(crate) fn finalize_after_route( schema: &Schema, offset_plan: Option<&OffsetPlan>, ) -> Result<(), Error> { - if !request.route().is_cross_shard() { - return Ok(()); - } - let Some(ast) = request.ast.as_ref() else { return Ok(()); }; - let rewrite_offset = offset_plan.is_some_and(|plan| !plan.prepare_execute); - let Some(rewrite) = ast - .post_route_rewrite - .get_or_try_init(|| build(&ast.ast, schema, rewrite_offset))? - else { + let rewrite = if request.route().is_cross_shard() { + let rewrite_offset = offset_plan.is_some_and(|plan| !plan.prepare_execute); + ast.post_route_rewrite + .get_or_try_init(|| build(&ast.ast, schema, rewrite_offset))? + .as_ref() + } else { + None + }; + let params = request.messages.iter().find_map(|message| match message { + ProtocolMessage::Bind(bind) => Some(StatementParameters::Bind(bind)), + _ => None, + }); + let order_by = match rewrite { + Some(rewrite) => OrderBy::parse(&rewrite.sort_clause, params), + None => statement_order_by(ast.ast.stmts().next(), params)?, + }; + if let Some(route) = request.route.as_mut() { + route.set_order_by(order_by); + } + + let Some(rewrite) = rewrite else { return Ok(()); }; let base_name = request.messages.iter().find_map(|message| match message { @@ -148,28 +151,37 @@ pub(crate) fn finalize_after_route( && let Some(route) = request.route.as_mut() { route.projection_rewrite_plan = rewrite.plan.clone(); - let mut order_by = route.order_by().to_vec(); - for helper in &rewrite.plan.order_by_helpers { - // Prefer the structural position. The source fallback handles a - // bind-dependent vector sort omitted from this execution's route. - let position = order_by - .get(helper.sort_position) - .filter(|sort| helper.matches(sort)) - .map(|_| helper.sort_position) - .or_else(|| order_by.iter().position(|sort| helper.matches(sort))); - let Some(sort) = position.and_then(|position| order_by.get_mut(position)) else { - continue; + } + + Ok(()) +} + +/// Statements without a projection variant still need sort keys from their +/// original AST, including the SELECT behind a simple-protocol EXECUTE. +fn statement_order_by( + statement: Option>, + params: Option>, +) -> Result, Error> { + match statement { + Some(Node::SelectStmt(select)) => Ok(OrderBy::parse(select.sort_clause(), params)), + Some(Node::ExplainStmt(explain)) => statement_order_by(Some(explain.query()), params), + Some(Node::ExecuteStmt(execute)) => { + let name = execute.name().expect("EXECUTE has a statement name"); + let Some(prepare) = PreparedStatements::global().read().prepare(name) else { + // Direct sessions can pass PREPARE/EXECUTE through without + // storing the statement in PgDog's prepared statement cache. + return Ok(Vec::new()); }; - *sort = if sort.asc() { - OrderBy::AscColumn(helper.alias.clone()) - } else { - OrderBy::DescColumn(helper.alias.clone()) + let ast = Cache::get().record(prepare.query())?; + let statement = match ast.ast.ast.stmts().next() { + Some(Node::PrepareStmt(prepare)) => Some(prepare.query()), + statement => statement, }; + let params = ExecuteParams::new(execute); + statement_order_by(statement, Some(StatementParameters::Execute(¶ms))) } - route.set_order_by(order_by); + _ => Ok(Vec::new()), } - - Ok(()) } fn build( @@ -193,6 +205,9 @@ fn build( plan = AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, &aggregate); } order_by::rewrite_select(&mut select.as_mut(), mem, &mut plan); + if aggregate.is_empty() { + order_by::rewrite_cases(&mut select.as_mut(), mem, schema, &mut plan); + } if rewrite_offset { offset::rewrite_select(&mut select.as_mut(), mem); } @@ -203,7 +218,12 @@ fn build( } let sql: Arc = pg_raw_parse::deparse(&*rewritten)?.as_str().into(); - Ok(Some(PostRouteRewrite { sql, plan })) + let sort_clause = make::owned(|mem| mem.make_unique(rewritten.sort_clause())); + Ok(Some(PostRouteRewrite { + sql, + plan, + sort_clause, + })) } #[cfg(test)] From 58b2ed878961c7bf7e7ac677db6fe36acc76229c Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 12:52:21 -0700 Subject: [PATCH 2/6] tests --- integration/rust/tests/integration/mod.rs | 1 + .../rust/tests/integration/order_by.rs | 104 ++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 integration/rust/tests/integration/order_by.rs diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index f8450f512..5a512de51 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -27,6 +27,7 @@ pub mod non_deterministic_funcs; pub mod notify; mod numeric_infinity; pub mod offset; +mod order_by; pub mod partial_req; pub mod per_stmt_routing; pub mod prepared; diff --git a/integration/rust/tests/integration/order_by.rs b/integration/rust/tests/integration/order_by.rs new file mode 100644 index 000000000..884d8d2a8 --- /dev/null +++ b/integration/rust/tests/integration/order_by.rs @@ -0,0 +1,104 @@ +use crate::setup::connections_sqlx; +use sqlx::{Column, Executor, Row, postgres::PgRow}; + +enum Protocol { + Simple, + Extended, +} + +#[tokio::test] +async fn order_by_case_simple_protocol() -> Result<(), Box> { + check_order_by_case(Protocol::Simple).await +} + +#[tokio::test] +async fn order_by_case_extended_protocol() -> Result<(), Box> { + check_order_by_case(Protocol::Extended).await +} + +async fn check_order_by_case(protocol: Protocol) -> Result<(), Box> { + let pools = connections_sqlx().await; + let mut transaction = pools[1].begin().await?; + + transaction.execute("TRUNCATE sharded").await?; + // Interleave priorities and IDs across shards so neither shard-local + // ordering nor sorting only by ID can produce the expected result. + for (shard, values) in [ + (0, "(1, 'other'), (4, 'pay'), (6, 'later')"), + (1, "(2, 'pay'), (3, 'other'), (5, 'pay')"), + ] { + let inserted = transaction + .execute( + format!( + "/* pgdog_shard: {shard} */ INSERT INTO sharded (id, value) + SELECT batch * 6 + fixture.id, fixture.value + FROM generate_series(0, 9) AS batch + CROSS JOIN (VALUES {values}) AS fixture(id, value)" + ) + .as_str(), + ) + .await?; + assert_eq!(inserted.rows_affected(), 30); + } + + let ascending: Vec<_> = (0_i64..10) + .flat_map(|batch| [2, 4, 5].map(|id| (batch * 6 + id, "pay"))) + .chain((0_i64..10).flat_map(|batch| { + [(1, "other"), (3, "other"), (6, "later")].map(|(id, name)| (batch * 6 + id, name)) + })) + .collect(); + let descending: Vec<_> = (0_i64..10) + .rev() + .flat_map(|batch| [3, 1].map(|id| (batch * 6 + id, "other"))) + .chain((0_i64..10).rev().map(|batch| (batch * 6 + 6, "later"))) + .chain( + (0_i64..10) + .rev() + .flat_map(|batch| [5, 4, 2].map(|id| (batch * 6 + id, "pay"))), + ) + .collect(); + + let cases = [ + ("CASE value WHEN $1 THEN 0 ELSE 1 END, id", ascending), + ( + "CASE value WHEN $1 THEN 0 ELSE 1 END DESC, + CASE value WHEN 'other' THEN 0 ELSE 1 END, id DESC", + descending, + ), + ]; + + for (order_by, expected) in &cases { + let query = format!("SELECT id, value AS name FROM sharded ORDER BY {order_by}"); + let rows = match protocol { + Protocol::Simple => { + sqlx::raw_sql(&query.replace("$1", "'pay'")) + .fetch_all(&mut *transaction) + .await? + } + Protocol::Extended => { + sqlx::query(&query) + .bind("pay") + .fetch_all(&mut *transaction) + .await? + } + }; + assert_rows(&rows, expected); + } + + transaction.rollback().await?; + Ok(()) +} + +fn assert_rows(rows: &[PgRow], expected: &[(i64, &str)]) { + assert_eq!(rows.len(), expected.len()); + for row in rows { + assert_eq!(row.columns().len(), 2, "sort helpers must stay hidden"); + assert_eq!(row.columns()[0].name(), "id"); + assert_eq!(row.columns()[1].name(), "name"); + } + let actual: Vec<_> = rows + .iter() + .map(|row| (row.get::(0), row.get::<&str, _>(1))) + .collect(); + assert_eq!(actual, expected); +} From 86a57388d48e367e06a36cf2c83ffdeb06f6fb43 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 12:58:41 -0700 Subject: [PATCH 3/6] cleanup --- .../rust/tests/integration/order_by.rs | 20 +++++++++++++++-- .../pool/connection/multi_shard/test.rs | 3 +-- .../query_engine/test/rewrite_offset.rs | 1 - .../query_engine/test/rewrite_projection.rs | 22 +++++-------------- .../frontend/router/parser/query/select.rs | 1 - .../router/parser/rewrite/statement/offset.rs | 2 -- pgdog/src/frontend/router/parser/route.rs | 10 ++------- 7 files changed, 26 insertions(+), 33 deletions(-) diff --git a/integration/rust/tests/integration/order_by.rs b/integration/rust/tests/integration/order_by.rs index 884d8d2a8..5d61520d1 100644 --- a/integration/rust/tests/integration/order_by.rs +++ b/integration/rust/tests/integration/order_by.rs @@ -30,8 +30,8 @@ async fn check_order_by_case(protocol: Protocol) -> Result<(), Box Result<(), Box = ascending + .iter() + .copied() + .filter(|(id, _)| ((id - 1) / 6) % 2 == 0) + .chain( + ascending + .iter() + .copied() + .filter(|(id, _)| ((id - 1) / 6) % 2 == 1), + ) + .collect(); + let cases = [ ("CASE value WHEN $1 THEN 0 ELSE 1 END, id", ascending), ( @@ -65,6 +77,10 @@ async fn check_order_by_case(protocol: Protocol) -> Result<(), Box) -> Option { fn cross_shard_route() -> Route { Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit::default(), None, diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index d08c39dd0..cb6b5f0a8 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -7,7 +7,6 @@ use crate::frontend::{ PreparedStatements, router::parser::{Limit, OrderBy}, }; -use pgdog_vector::Vector; use std::sync::Arc; use super::prelude::*; @@ -16,7 +15,6 @@ use super::test_sharded_client; fn route(shard: Shard) -> Route { Route::select( ShardWithPriority::new_table(shard), - vec![], Default::default(), Limit::default(), None, @@ -229,7 +227,6 @@ async fn cross_shard_order_by_projects_missing_sort_column() { .unwrap(); client_request.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::AscColumn("price".into())], Default::default(), Limit::default(), None, @@ -262,14 +259,13 @@ async fn cross_shard_order_by_projects_missing_sort_column() { } #[test] -fn cached_projection_does_not_depend_on_first_route_order() { +fn cached_projection_keeps_all_sort_keys_without_bind_values() { let sql = "SELECT id FROM products ORDER BY embedding <-> $1, price"; let ast = Arc::new(Ast::parse(sql).unwrap()); let mut first = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); first.ast = Some(ast.clone()); first.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::AscColumn("price".into())], Default::default(), Limit::default(), None, @@ -284,17 +280,16 @@ fn cached_projection_does_not_depend_on_first_route_order() { assert!(first_query.query().contains("__pgdog_order_col1")); assert_eq!( first.route().order_by(), - &[OrderBy::AscColumn("__pgdog_order_col1".into())] + &[ + OrderBy::AscColumn("__pgdog_order_col0".into()), + OrderBy::AscColumn("__pgdog_order_col1".into()) + ] ); let mut second = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); second.ast = Some(ast); second.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![ - OrderBy::AscVectorL2Column("embedding".into(), Vector::from(&[1.0, 2.0, 3.0][..])), - OrderBy::AscColumn("price".into()), - ], Default::default(), Limit::default(), None, @@ -317,7 +312,6 @@ fn aliased_projected_sort_column_remaps_route() { request.ast = Some(Arc::new(Ast::parse(sql).unwrap())); request.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::AscColumn("price".into())], Default::default(), Limit::default(), None, @@ -344,7 +338,6 @@ fn duplicate_sort_column_names_use_injected_helper() { request.ast = Some(Arc::new(Ast::parse(sql).unwrap())); request.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::AscColumn("price".into())], Default::default(), Limit::default(), None, @@ -370,10 +363,6 @@ fn helper_replaces_the_matching_duplicate_order_by_position() { request.ast = Some(Arc::new(Ast::parse(sql).unwrap())); request.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![ - OrderBy::AscColumn("price".into()), - OrderBy::AscColumn("price".into()), - ], Default::default(), Limit::default(), None, @@ -405,7 +394,6 @@ async fn aggregate_order_by_and_offset_compose_after_route() { .unwrap(); client_request.route = Some(Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::AscColumn("created_at".into())], Default::default(), Limit::default(), None, diff --git a/pgdog/src/frontend/router/parser/query/select.rs b/pgdog/src/frontend/router/parser/query/select.rs index bf5e5b207..832685150 100644 --- a/pgdog/src/frontend/router/parser/query/select.rs +++ b/pgdog/src/frontend/router/parser/query/select.rs @@ -317,7 +317,6 @@ impl QueryParser { let query = Route::select( context.shards_calculator.shard().clone(), - Vec::new(), aggregates, limit, distinct, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs index 8d35bea45..9531f34b4 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs @@ -269,7 +269,6 @@ mod tests { fn cross_shard_route() -> Route { Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit::default(), None, @@ -279,7 +278,6 @@ mod tests { fn single_shard_route() -> Route { Route::select( ShardWithPriority::new_table(Shard::Direct(0)), - vec![], Default::default(), Limit::default(), None, diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index a0ae3728e..a0d368e51 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -153,14 +153,12 @@ impl Route { /// Create new route for a `SELECT` query. pub(crate) fn select( shard: ShardWithPriority, - order_by: Vec, aggregate: Aggregate, limit: Limit, distinct: Option, ) -> Self { Self { shard, - order_by, read: true, aggregate, limit, @@ -720,13 +718,13 @@ mod test { #[test] fn test_should_buffer_order_by() { - let route = Route::select( + let mut route = Route::select( ShardWithPriority::new_table(Shard::All), - vec![OrderBy::Asc(0)], Default::default(), Limit::default(), None, ); + route.set_order_by(vec![OrderBy::Asc(1)]); assert!(route.requires_post_processing()); } @@ -734,7 +732,6 @@ mod test { fn test_should_buffer_limit_only() { let route = Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit { limit: Some(10), @@ -749,7 +746,6 @@ mod test { fn test_should_buffer_offset_only() { let route = Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit { limit: None, @@ -764,7 +760,6 @@ mod test { fn test_should_buffer_limit_and_offset() { let route = Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit { limit: Some(10), @@ -779,7 +774,6 @@ mod test { fn test_should_buffer_no_limit_no_offset() { let route = Route::select( ShardWithPriority::new_table(Shard::All), - vec![], Default::default(), Limit::default(), None, From 854dce7f87aff9483ee4071cf57315090c7f6f80 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 12:59:10 -0700 Subject: [PATCH 4/6] revert pgdog.toml change --- integration/pgdog.toml | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 918ead29d..429b7ffee 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -25,8 +25,8 @@ tls_certificate = "integration/tls/cert.pem" tls_private_key = "integration/tls/key.pem" query_parser_engine = "pg_query_raw" system_catalogs = "omnisharded_sticky" -reload_schema_on_ddl = true -idle_healthcheck_delay = 50000000 +reload_schema_on_ddl = false +# idle_healthcheck_delay = 50000000 unique_id_function = "standard" auth_type = "scram" workers = 4 @@ -42,7 +42,7 @@ message_buffer = 8096 [rewrite] enabled = false shard_key = "ignore" -split_inserts = "rewrite" +split_inserts = "error" # primary_key = "rewrite" non_deterministic_functions = "rewrite" @@ -190,11 +190,6 @@ database = "pgdog_sharded" column = "customer_id" data_type = "bigint" -[[sharded_tables]] -database = "pgdog_sharded" -column = "tenant_id" -data_type = "bigint" - # ------------------------------------------------------------------------------ # ----- Hash Sharded :: VARCHAR ------------------------------------------------ From e40ebff8f13cac0e80e5a77933e189163569d9fc Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 13:09:37 -0700 Subject: [PATCH 5/6] reusable helper column name --- .../query_engine/test/rewrite_projection.rs | 84 ++++++++++--------- .../rewrite/statement/aggregate/engine.rs | 35 ++++++-- .../parser/rewrite/statement/aggregate/mod.rs | 11 --- .../parser/rewrite/statement/order_by.rs | 36 +++++--- .../parser/rewrite/statement/order_by/case.rs | 35 +------- .../parser/rewrite/statement/projection.rs | 3 + .../rewrite/statement/projection/names.rs | 27 ++++++ 7 files changed, 129 insertions(+), 102 deletions(-) create mode 100644 pgdog/src/frontend/router/parser/rewrite/statement/projection/names.rs diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index cb6b5f0a8..f8455773b 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -82,7 +82,7 @@ async fn cross_shard_aggregate_adds_and_tracks_helpers() { ProtocolMessage::Query(query) => query, _ => panic!("expected Query"), }; - assert!(query.query().contains("__pgdog_count_col0")); + assert!(query.query().contains("__pgdog_count_col")); assert_eq!( client_request .route() @@ -132,7 +132,7 @@ async fn named_prepared_aggregate_uses_cross_shard_variant() { match &client_request.messages[0] { ProtocolMessage::Parse(parse) => { assert_eq!(parse.name(), variant); - assert!(parse.query().contains("__pgdog_count_col0")); + assert!(parse.query().contains("__pgdog_count_col")); } _ => panic!("expected Parse"), } @@ -155,7 +155,7 @@ async fn named_prepared_aggregate_uses_cross_shard_variant() { .rewritten_parse(&variant) .unwrap() .query() - .contains("__pgdog_count_col0") + .contains("__pgdog_count_col") ); } @@ -243,10 +243,15 @@ async fn cross_shard_order_by_projects_missing_sort_column() { ProtocolMessage::Query(query) => query, _ => panic!("expected Query"), }; - assert!(query.query().contains("price AS __pgdog_order_col0")); + let alias = &client_request + .route() + .projection_rewrite_plan + .order_by_helpers[0] + .alias; + assert!(query.query().contains(&format!("price AS {alias}"))); assert_eq!( client_request.route().order_by(), - &[OrderBy::AscColumn("__pgdog_order_col0".into())] + &[OrderBy::AscColumn(alias.clone())] ); assert_eq!( client_request @@ -276,15 +281,16 @@ fn cached_projection_keeps_all_sort_keys_without_bind_values() { ProtocolMessage::Query(query) => query, _ => panic!("expected Query"), }; - assert!(first_query.query().contains("__pgdog_order_col0")); - assert!(first_query.query().contains("__pgdog_order_col1")); - assert_eq!( - first.route().order_by(), - &[ - OrderBy::AscColumn("__pgdog_order_col0".into()), - OrderBy::AscColumn("__pgdog_order_col1".into()) - ] - ); + let helpers = &first.route().projection_rewrite_plan.order_by_helpers; + assert_eq!(helpers.len(), 2); + for helper in helpers { + assert!(first_query.query().contains(&helper.alias)); + } + let expected = [ + OrderBy::AscColumn(helpers[0].alias.clone()), + OrderBy::AscColumn(helpers[1].alias.clone()), + ]; + assert_eq!(first.route().order_by(), &expected); let mut second = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); second.ast = Some(ast); @@ -296,13 +302,7 @@ fn cached_projection_keeps_all_sort_keys_without_bind_values() { )); projection::finalize_after_route(&mut second, &Schema::default(), None).unwrap(); - assert_eq!( - second.route().order_by(), - &[ - OrderBy::AscColumn("__pgdog_order_col0".into()), - OrderBy::AscColumn("__pgdog_order_col1".into()) - ] - ); + assert_eq!(second.route().order_by(), &expected); } #[test] @@ -349,10 +349,11 @@ fn duplicate_sort_column_names_use_injected_helper() { ProtocolMessage::Query(query) => query, _ => panic!("expected Query"), }; - assert!(query.query().contains("b.price AS __pgdog_order_col0")); + let alias = &request.route().projection_rewrite_plan.order_by_helpers[0].alias; + assert!(query.query().contains(&format!("b.price AS {alias}"))); assert_eq!( request.route().order_by(), - &[OrderBy::AscColumn("__pgdog_order_col0".into())] + &[OrderBy::AscColumn(alias.clone())] ); } @@ -374,7 +375,11 @@ fn helper_replaces_the_matching_duplicate_order_by_position() { request.route().order_by(), &[ OrderBy::AscColumn("price".into()), - OrderBy::AscColumn("__pgdog_order_col1".into()) + OrderBy::AscColumn( + request.route().projection_rewrite_plan.order_by_helpers[0] + .alias + .clone() + ) ] ); } @@ -415,23 +420,22 @@ async fn aggregate_order_by_and_offset_compose_after_route() { ProtocolMessage::Query(query) => query, _ => panic!("expected Query"), }; - assert!(query.query().contains("__pgdog_count_col0")); - assert!(query.query().contains("created_at AS __pgdog_order_col0")); + assert!(query.query().contains("__pgdog_count_col")); + let alias = &client_request + .route() + .projection_rewrite_plan + .order_by_helpers[0] + .alias; + assert!(query.query().contains(&format!("created_at AS {alias}"))); assert!(query.query().contains("LIMIT 10::bigint + 5::bigint")); assert!(!query.query().contains("OFFSET")); let route = client_request.route(); - assert_eq!( - route.order_by(), - &[OrderBy::AscColumn("__pgdog_order_col0".into())] - ); - assert_eq!( - route.projection_rewrite_plan.aggregate_helpers[0].alias, - "__pgdog_count_col0" - ); - assert_eq!( - route.projection_rewrite_plan.order_by_helpers[0].alias, - "__pgdog_order_col0" + assert_eq!(route.order_by(), &[OrderBy::AscColumn(alias.clone())]); + assert!( + query + .query() + .contains(&route.projection_rewrite_plan.aggregate_helpers[0].alias) ); assert_eq!( route.limit(), @@ -476,7 +480,7 @@ async fn split_anonymous_prepare_rewrites_each_execution_once() { ProtocolMessage::Parse(parse) => parse, _ => panic!("expected Parse"), }; - assert_eq!(parse.query().matches("__pgdog_count_col0").count(), 1); + assert_eq!(parse.query().matches("__pgdog_count_col").count(), 1); assert!( !client .client_request @@ -484,7 +488,7 @@ async fn split_anonymous_prepare_rewrites_each_execution_once() { .as_ref() .unwrap() .query() - .contains("__pgdog_count_col0") + .contains("__pgdog_count_col") ); client.client_request.clear(); @@ -517,7 +521,7 @@ async fn split_anonymous_prepare_rewrites_each_execution_once() { .as_ref() .unwrap() .query() - .matches("__pgdog_count_col0") + .matches("__pgdog_count_col") .count(), 1 ); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs index c4205f1ed..ff5b5ffb4 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs @@ -4,7 +4,9 @@ use itertools::*; use pg_raw_parse::{Node, make, nodes}; use super::{AggregateHelper, HelperKind}; -use crate::frontend::router::parser::rewrite::statement::projection::ProjectionRewritePlan; +use crate::frontend::router::parser::rewrite::statement::projection::{ + HelperColumnKind, ProjectionRewritePlan, helper_column_name, +}; /// Query rewrite engine. Currently supports injecting helper aggregates for AVG and /// variance-related functions that require additional helper aggregates when run @@ -39,8 +41,7 @@ impl AggregatesRewrite { .map(move |spec| (target, spec)) }) .map(|(target, HelperSpec { func, kind })| { - let helper_alias = - format!("__pgdog_{}_col{}", kind.alias_suffix(), target.column()); + let helper_alias = helper_column_name(HelperColumnKind::Aggregate(kind)); let node = mem.make_res_target(Some(&helper_alias), mem.empty(), func.uncast()); plan.aggregate_helpers.push(AggregateHelper { @@ -195,7 +196,12 @@ mod tests { assert_eq!(plan.aggregate_helpers.len(), 1); let helper = &plan.aggregate_helpers[0]; assert_eq!(helper.target_column, 0); - assert_eq!(helper.alias, "__pgdog_count_col0"); + assert!(helper.alias.starts_with("__pgdog_count_col")); + assert!( + ast.target_list() + .iter() + .any(|target| target.name() == Some(helper.alias.as_str())) + ); assert!(!helper.distinct); assert!(matches!(helper.kind, HelperKind::Count)); @@ -215,7 +221,12 @@ mod tests { assert_eq!(plan.aggregate_helpers.len(), 1); let helper = &plan.aggregate_helpers[0]; assert_eq!(helper.target_column, 1); - assert_eq!(helper.alias, "__pgdog_count_col1"); + assert!(helper.alias.starts_with("__pgdog_count_col")); + assert!( + ast.target_list() + .iter() + .any(|target| target.name() == Some(helper.alias.as_str())) + ); assert!(!helper.distinct); assert!(matches!(helper.kind, HelperKind::Count)); @@ -238,12 +249,22 @@ mod tests { let helper_price = &plan.aggregate_helpers[0]; assert_eq!(helper_price.target_column, 0); - assert_eq!(helper_price.alias, "__pgdog_count_col0"); + assert!(helper_price.alias.starts_with("__pgdog_count_col")); + assert!( + ast.target_list() + .iter() + .any(|target| target.name() == Some(helper_price.alias.as_str())) + ); assert!(matches!(helper_price.kind, HelperKind::Count)); let helper_discount = &plan.aggregate_helpers[1]; assert_eq!(helper_discount.target_column, 1); - assert_eq!(helper_discount.alias, "__pgdog_count_col1"); + assert!(helper_discount.alias.starts_with("__pgdog_count_col")); + assert!( + ast.target_list() + .iter() + .any(|target| target.name() == Some(helper_discount.alias.as_str())) + ); assert!(matches!(helper_discount.kind, HelperKind::Count)); let aggregate = Aggregate::parse(&ast, &Default::default()); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs index 5709b4589..180b842b4 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs @@ -14,14 +14,3 @@ pub(crate) enum HelperKind { /// `SUM(POWER(column, 2))`. SumSquares, } - -impl HelperKind { - /// Suffix used in the projected helper's internal alias. - pub(crate) fn alias_suffix(self) -> &'static str { - match self { - Self::Count => "count", - Self::Sum => "sum", - Self::SumSquares => "sumsq", - } - } -} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs index de6318705..c54f6cc2f 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs @@ -2,7 +2,9 @@ use pg_raw_parse::{ConstValue, Node, make, nodes}; use crate::frontend::router::parser::Column; -use super::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan}; +use super::projection::{ + HelperColumnKind, OrderByHelper, OrderBySource, ProjectionRewritePlan, helper_column_name, +}; mod case; pub(super) use case::rewrite_cases; @@ -208,7 +210,7 @@ pub(super) fn rewrite_select<'a>( ); } None => { - let alias = format!("__pgdog_order_col{current_sort_position}"); + let alias = helper_column_name(HelperColumnKind::OrderBy); sort_rewrites.push((position, alias.clone())); helpers.push(mem.make_res_target( Some(&alias), @@ -225,7 +227,7 @@ pub(super) fn rewrite_select<'a>( } }, Node::A_Expr(_) if source.is_some() => { - let alias = format!("__pgdog_order_col{current_sort_position}"); + let alias = helper_column_name(HelperColumnKind::OrderBy); sort_rewrites.push((position, alias.clone())); helpers.push(mem.make_res_target( Some(&alias), @@ -291,9 +293,13 @@ mod tests { fn projects_missing_sort_column() { let (sql, plan) = rewrite("SELECT id FROM products ORDER BY price"); - assert!(sql.contains("price AS __pgdog_order_col0")); + assert!(sql.contains(&format!("price AS {}", plan.order_by_helpers[0].alias))); assert_eq!(plan.order_by_helpers.len(), 1); - assert_eq!(plan.order_by_helpers[0].alias, "__pgdog_order_col0"); + assert!( + plan.order_by_helpers[0] + .alias + .starts_with("__pgdog_order_col") + ); assert!(plan.order_by_helpers[0].injected); } @@ -328,7 +334,7 @@ mod tests { let (sql, plan) = rewrite("SELECT a.price, b.price FROM a JOIN b ON a.id = b.a_id ORDER BY b.price"); - assert!(sql.contains("b.price AS __pgdog_order_col0")); + assert!(sql.contains(&format!("b.price AS {}", plan.order_by_helpers[0].alias))); assert_eq!(plan.order_by_helpers.len(), 1); assert!(plan.order_by_helpers[0].injected); } @@ -350,7 +356,7 @@ mod tests { let (sql, plan) = rewrite("SELECT a.*, b.price FROM a JOIN b ON a.id = b.a_id ORDER BY b.price"); - assert!(sql.contains("b.price AS __pgdog_order_col0")); + assert!(sql.contains(&format!("b.price AS {}", plan.order_by_helpers[0].alias))); assert!(plan.order_by_helpers[0].injected); } @@ -358,7 +364,7 @@ mod tests { fn injects_helper_for_unqualified_star_join() { let (sql, plan) = rewrite("SELECT * FROM a JOIN b ON a.id = b.a_id ORDER BY b.price"); - assert!(sql.contains("b.price AS __pgdog_order_col0")); + assert!(sql.contains(&format!("b.price AS {}", plan.order_by_helpers[0].alias))); assert!(plan.order_by_helpers[0].injected); } @@ -382,7 +388,7 @@ mod tests { fn projects_column_not_covered_by_qualified_star() { let (sql, plan) = rewrite("SELECT a.* FROM a JOIN b ON a.id = b.a_id ORDER BY b.score"); - assert!(sql.contains("b.score AS __pgdog_order_col0")); + assert!(sql.contains(&format!("b.score AS {}", plan.order_by_helpers[0].alias))); assert_eq!(plan.order_by_helpers.len(), 1); } @@ -390,8 +396,12 @@ mod tests { fn projects_unqualified_sort_for_qualified_star() { let (sql, plan) = rewrite("SELECT p.* FROM products p ORDER BY price"); - assert!(sql.contains("price AS __pgdog_order_col0")); - assert_eq!(plan.order_by_helpers[0].alias, "__pgdog_order_col0"); + assert!(sql.contains(&format!("price AS {}", plan.order_by_helpers[0].alias))); + assert!( + plan.order_by_helpers[0] + .alias + .starts_with("__pgdog_order_col") + ); } #[test] @@ -399,7 +409,7 @@ mod tests { let (sql, plan) = rewrite("SELECT a.price FROM a JOIN b ON a.id = b.a_id ORDER BY a.price, b.price"); - assert!(sql.contains("b.price AS __pgdog_order_col1")); + assert!(sql.contains(&format!("b.price AS {}", plan.order_by_helpers[0].alias))); assert_eq!(plan.order_by_helpers.len(), 1); assert_eq!(plan.order_by_helpers[0].sort_position, 1); } @@ -409,7 +419,7 @@ mod tests { let (sql, plan) = rewrite("SELECT id FROM products ORDER BY embedding <-> $1 LIMIT 5"); assert!(sql.contains("embedding <-> $1")); - assert!(sql.contains("AS __pgdog_order_col0")); + assert!(sql.contains(&format!("AS {}", plan.order_by_helpers[0].alias))); assert_eq!(plan.order_by_helpers.len(), 1); assert!(plan.order_by_helpers[0].injected); } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs index 786bb297a..75feba477 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by/case.rs @@ -1,9 +1,10 @@ use pg_raw_parse::{Node, make, nodes, walk}; -use std::collections::HashSet; use crate::backend::schema::Schema; -use super::{OrderBySource, ProjectionRewritePlan, push_helper}; +use super::{ + HelperColumnKind, OrderBySource, ProjectionRewritePlan, helper_column_name, push_helper, +}; /// Project row-level CASE sort keys. Aggregation and DISTINCT need their own /// rewrite rules because adding a target can change the query's semantics. @@ -29,7 +30,6 @@ pub(in crate::frontend::router::parser::rewrite::statement) fn rewrite_cases<'a> // present only in ORDER BY must also prevent this rewrite. Without a // function catalog, conservatively leave function calls alone. let mut aggregate = false; - let mut names = HashSet::new(); walk::walk(Node::SelectStmt(select), |node| match node { Node::FuncCall(func) => { let name = func.funcname().iter().next_back().and_then(Node::as_str); @@ -42,46 +42,19 @@ pub(in crate::frontend::router::parser::rewrite::statement) fn rewrite_cases<'a> || func.over().is_some(); } Node::JsonAggConstructor(_) => aggregate = true, - Node::ResTarget(target) => { - if let Some(name) = target.name() { - names.insert(name.to_owned()); - } - } - Node::String(string) => { - if let Some(name) = string.sval() { - names.insert(name.to_owned()); - } - } _ => {} }); if aggregate { return; } - // Include known star expansions, and keep names deterministic so a cached - // prepared variant still matches if its AST is evicted and rebuilt. - names.extend( - schema - .relations - .values() - .flat_map(|relations| relations.values()) - .flat_map(|relation| relation.columns.keys().cloned()), - ); - let mut helpers = Vec::new(); for (position, mut sort) in select.sort_clause_mut().into_iter().enumerate() { if !matches!(sort.node(), Node::CaseExpr(_)) { continue; } - let mut suffix = position; - let alias = loop { - let alias = format!("__pgdog_order_case{suffix}"); - if names.insert(alias.clone()) { - break alias; - } - suffix += 1; - }; + let alias = helper_column_name(HelperColumnKind::Case); helpers.push(mem.make_res_target( Some(&alias), mem.empty(), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 71325f0a6..c0b59ec45 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -18,6 +18,9 @@ use pg_raw_parse::{Node, Owned, StmtList, list::CastNodeList, make, nodes}; use std::collections::BTreeSet; use std::sync::Arc; +mod names; +pub(super) use names::{HelperColumnKind, helper_column_name}; + #[derive(Debug, Clone, PartialEq)] pub(crate) struct AggregateHelper { pub(crate) target_column: usize, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection/names.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection/names.rs new file mode 100644 index 000000000..b0a77c8b6 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection/names.rs @@ -0,0 +1,27 @@ +use rand::{Rng, rng}; +use std::sync::atomic::{AtomicU64, Ordering}; + +use super::HelperKind; + +static COUNTER: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Copy)] +pub(crate) enum HelperColumnKind { + Aggregate(HelperKind), + OrderBy, + Case, +} + +/// Generate a helper alias with eight random digits and a process-wide counter. +pub(crate) fn helper_column_name(kind: HelperColumnKind) -> String { + let prefix = match kind { + HelperColumnKind::Aggregate(HelperKind::Count) => "count_col", + HelperColumnKind::Aggregate(HelperKind::Sum) => "sum_col", + HelperColumnKind::Aggregate(HelperKind::SumSquares) => "sumsq_col", + HelperColumnKind::OrderBy => "order_col", + HelperColumnKind::Case => "order_case", + }; + let suffix = rng().random_range(0..100_000_000u32); // Basically we don't conflict with some existing column name in the schema. + let counter = COUNTER.fetch_add(1, Ordering::Relaxed); // Guarantees against collisions in suffix. + format!("__pgdog_{prefix}_{suffix:08}_{counter}") +} From ad9fedd6994ace0136c89b9f45ffe9ef661039d1 Mon Sep 17 00:00:00 2001 From: Lev Kokotov Date: Tue, 6 Oct 2026 13:10:37 -0700 Subject: [PATCH 6/6] move some stuff --- .../parser/rewrite/statement/{order_by.rs => order_by/mod.rs} | 0 .../parser/rewrite/statement/{projection.rs => projection/mod.rs} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename pgdog/src/frontend/router/parser/rewrite/statement/{order_by.rs => order_by/mod.rs} (100%) rename pgdog/src/frontend/router/parser/rewrite/statement/{projection.rs => projection/mod.rs} (100%) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by/mod.rs similarity index 100% rename from pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs rename to pgdog/src/frontend/router/parser/rewrite/statement/order_by/mod.rs diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection/mod.rs similarity index 100% rename from pgdog/src/frontend/router/parser/rewrite/statement/projection.rs rename to pgdog/src/frontend/router/parser/rewrite/statement/projection/mod.rs