diff --git a/integration/pgdog.toml b/integration/pgdog.toml index d5fd64d75..e2cacb8b0 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -53,6 +53,26 @@ host = "127.0.0.1" role = "replica" read_only = true +# ------------------------------------------------------------------------------ +# ----- Database :: pgdog_leak ------------------------------------------------- +# One server connection, never reaped: the next client has to get the same one. +# Two copies of the same single-primary database, differing only in parser level, +# so the leak tests cover both ways a statement reaches the parser. + +[[databases]] +name = "pgdog_leak" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + +[[databases]] +name = "pgdog_leak_auto" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + # ------------------------------------------------------------------------------ # ----- Database :: pgdog_sharded ---------------------------------------------- @@ -495,3 +515,17 @@ password = "pgdog" database = "pgdog" level = "auto" engine = "pg_query_raw" + +# These tests are about statements the parser would otherwise opt out of. +[[query_parsers]] +database = "pgdog_leak" +level = "on" +engine = "pg_query_raw" + +# The same tests at the default level. A single primary with no replicas is the +# topology where "auto" doesn't force the parser on, so the statement has to get +# past the regex gate on its own -- the path "on" skips. +[[query_parsers]] +database = "pgdog_leak_auto" +level = "auto" +engine = "pg_query_raw" diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py new file mode 100644 index 000000000..e044146a1 --- /dev/null +++ b/integration/python/test_session_params_leak.py @@ -0,0 +1,171 @@ +"""Session parameters must not survive the client that set them. + +Runs against a database with a single server connection, so the next client +always gets the connection the previous one used. current_setting() is used +instead of SHOW because SHOW can be answered by PgDog itself. + +Every test runs against two copies of that database that differ only in parser +level: "on" always parses, "auto" leaves a single-primary cluster to the regex +gate. A statement that only the gate can let through -- set_config(), a function +call inside a SELECT rather than a statement-start keyword -- reaches the parser +on one and has to earn it on the other. +""" + +import uuid + +import psycopg +import pytest + +DATABASES = ["pgdog_leak", "pgdog_leak_auto"] + + +def connect(dbname): + conn = psycopg.connect( + user="pgdog", + password="pgdog", + dbname=dbname, + host="127.0.0.1", + port=6432, + ) + # Otherwise psycopg wraps each statement in a transaction and rolls it + # back on close, undoing the state we're testing for. + conn.autocommit = True + return conn + + +def read(dbname, setting): + conn = connect(dbname) + value = conn.execute(f"SELECT current_setting('{setting}')").fetchone()[0] + conn.close() + + return value + + +@pytest.fixture(params=DATABASES) +def dbname(request): + return request.param + + +def test_reset_rolled_back(dbname): + """A ROLLBACK brings back the value the RESET cleared. + + This is the sequence pg_dump -t emits. SET search_path TO '' leaves + an empty quoted identifier, hence the two spellings of "empty". + """ + conn = connect(dbname) + conn.execute("SET search_path TO ''") + with conn.transaction(): + conn.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY") + conn.execute("RESET search_path") + conn.execute("SELECT 1") + raise psycopg.Rollback() + conn.close() + + assert read(dbname, "search_path") not in ("", '""') + + +def test_set_committed_after_connecting(dbname): + """A SET that lands once the connection is already ours.""" + conn = connect(dbname) + with conn.transaction(): + conn.execute("SELECT 1") + conn.execute("SET statement_timeout TO '5s'") + conn.close() + + assert read(dbname, "statement_timeout") == "0" + + +def test_set_config_bound_params(dbname): + """set_config() arguments arrive in the Bind message, not as constants.""" + conn = connect(dbname) + conn.execute("SELECT pg_catalog.set_config(%s, %s, false)", ("search_path", "")) + conn.close() + + assert read(dbname, "search_path") not in ("", '""') + + +def test_reset_committed(dbname): + """A committed RESET is permanent and needs no undoing.""" + conn = connect(dbname) + conn.execute("SET search_path TO public") + with conn.transaction(): + conn.execute("SELECT 1") + conn.execute("RESET search_path") + conn.close() + + assert read(dbname, "search_path") == '"$user", public' + + +TENANT_A = "11111111-1111-1111-1111-111111111111" +TENANT_B = "22222222-2222-2222-2222-222222222222" + + +@pytest.fixture +def tenants(dbname): + """A table whose rows are visible only to the tenant named in a GUC. + + Reads go through a plain role: the pooler's own user is a superuser, and + superusers ignore row-level security however the table is configured. + + NULLIF() is deliberate: a targeted RESET leaves the placeholder GUC as an + empty string rather than unset, and '' would fail the ::uuid cast. + """ + table = "rls_probe_" + uuid.uuid4().hex[:8] + conn = connect(dbname) + conn.execute( + "DO $$ BEGIN CREATE ROLE rls_tenant NOLOGIN; " + "EXCEPTION WHEN duplicate_object THEN NULL; END $$" + ) + conn.execute(f"CREATE TABLE public.{table} (org_id uuid, note text)") + conn.execute( + f"INSERT INTO public.{table} VALUES ('{TENANT_A}', 'a'), ('{TENANT_B}', 'b')" + ) + conn.execute(f"GRANT SELECT ON public.{table} TO rls_tenant") + conn.execute(f"ALTER TABLE public.{table} ENABLE ROW LEVEL SECURITY") + conn.execute( + f"CREATE POLICY tenant_isolation ON public.{table} USING " + "(org_id = NULLIF(current_setting('app.current_org_id', true), '')::uuid)" + ) + conn.close() + + yield table + + conn = connect(dbname) + conn.execute("RESET ROLE") + conn.execute(f"DROP TABLE public.{table}") + conn.close() + + +def test_tenant_guc_does_not_outlive_its_client(dbname, tenants): + """The silent half of the leak: no error, just another tenant's rows. + + Row-level security keyed on a custom GUC is how multi-tenant applications + isolate tenants, and set_config() with a bound parameter is how that GUC + gets set. A value that survives checkin makes the next client read as the + previous tenant. + """ + first = connect(dbname) + first.execute("SET ROLE rls_tenant") + first.execute( + "SELECT pg_catalog.set_config('app.current_org_id', %s, false)", (TENANT_A,) + ) + mine = first.execute(f"SELECT note FROM public.{tenants}").fetchall() + served_by = first.execute("SELECT pg_backend_pid()").fetchone()[0] + first.close() + + assert mine == [("a",)], "the tenant that set the GUC sees its own row" + + second = connect(dbname) + second.execute("SET ROLE rls_tenant") + theirs = second.execute(f"SELECT note FROM public.{tenants}").fetchall() + leaked = second.execute( + "SELECT current_setting('app.current_org_id', true)" + ).fetchone()[0] + same_server = second.execute("SELECT pg_backend_pid()").fetchone()[0] + second.close() + + # Without this the test passes whenever the pool happens to hand out a + # different connection, which proves nothing about what the first one left. + assert same_server == served_by, "clients did not share a server connection" + assert leaked in (None, ""), f"the tenant GUC outlived its client: {leaked!r}" + assert theirs == [], "next client read the previous tenant's rows" diff --git a/integration/rust/tests/integration/set_config.rs b/integration/rust/tests/integration/set_config.rs index 8d294d1a1..15e9b03d6 100644 --- a/integration/rust/tests/integration/set_config.rs +++ b/integration/rust/tests/integration/set_config.rs @@ -28,12 +28,11 @@ async fn test_set_config_behaves_like_set() { .unwrap(); assert_eq!(lock_timeout, "500s"); - let set_config: Option = - query_scalar("SELECT set_config('lock_timeout', NULL, false);") - .fetch_one(&mut *conn) - .await - .unwrap(); - assert_eq!(set_config, None); + let set_config: String = query_scalar("SELECT set_config('lock_timeout', NULL, false);") + .fetch_one(&mut *conn) + .await + .unwrap(); + assert_eq!(set_config, "0"); let lock_timeout: String = query_scalar("SHOW lock_timeout") .fetch_one(&mut *conn) diff --git a/integration/users.toml b/integration/users.toml index bba115a85..25a24d84f 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -3,6 +3,16 @@ name = "pgdog" database = "pgdog" password = "pgdog" +[[users]] +name = "pgdog" +database = "pgdog_leak" +password = "pgdog" + +[[users]] +name = "pgdog" +database = "pgdog_leak_auto" +password = "pgdog" + [[users]] name = "pgdog_migrator" database = "pgdog" diff --git a/pgdog/src/backend/pool/connection/binding.rs b/pgdog/src/backend/pool/connection/binding.rs index faa5fe28a..57a05cff9 100644 --- a/pgdog/src/backend/pool/connection/binding.rs +++ b/pgdog/src/backend/pool/connection/binding.rs @@ -2,7 +2,7 @@ use crate::{ frontend::{ - ClientRequest, + ClientRequest, SetParam, client::query_engine::{ TwoPcPhase, two_pc::{TwoPcTransaction, statement::phase_control}, @@ -444,6 +444,28 @@ impl Binding { } } + /// Record a client parameter change on every server we hold. + pub fn record_params(&mut self, params: &[SetParam], in_transaction: bool) { + match self { + Binding::Direct(server, ..) => server.record_params(params, in_transaction), + Binding::MultiShard(servers, _) => servers + .iter_mut() + .for_each(|server| server.record_params(params, in_transaction)), + _ => (), + } + } + + /// Record a client `RESET ALL` on every server we hold. + pub fn record_reset_all(&mut self, in_transaction: bool) { + match self { + Binding::Direct(server, ..) => server.record_reset_all(in_transaction), + Binding::MultiShard(servers, _) => servers + .iter_mut() + .for_each(|server| server.record_reset_all(in_transaction)), + _ => (), + } + } + /// Handle transaction end. pub fn transaction_params_hook(&mut self, rollback: bool) { match self { diff --git a/pgdog/src/backend/server.rs b/pgdog/src/backend/server.rs index 96d9d0e7f..30b38b552 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -20,7 +20,7 @@ use crate::{ auth::{md5, scram::Client}, backend::pool::stats::MemoryStats, config::AuthType, - frontend::ClientRequest, + frontend::{ClientRequest, SetParam}, net::{ Close, Liveness, MessageBuffer, Parameter, ProtocolMessage, Sync, messages::{ @@ -140,6 +140,9 @@ pub struct Server { streaming: bool, schema_changed: bool, sync_prepared: bool, + // The change in flight is already recorded, so its CommandComplete + // tells us nothing new. + params_recorded: bool, in_transaction: bool, re_synced: bool, replication_mode: bool, @@ -431,6 +434,7 @@ impl Server { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, statement_executed: false, re_synced: false, @@ -597,6 +601,8 @@ impl Server { match message.code() { 'Z' => { + self.params_recorded = false; + let now = Instant::now(); let rfq = ReadyForQuery::from_bytes(message.payload())?; @@ -654,7 +660,9 @@ impl Server { self.prepared_statements.clear(); self.client_params.clear(); } - "RESET" => self.client_params.clear(), // Someone reset params, we're gonna need to re-sync. + // A RESET nobody told us the contents of: we can't tell + // what it touched, so drop everything. + "RESET" if !self.params_recorded => self.client_params.clear(), _ => (), } self.stats.rows_affected(&cmd); @@ -746,6 +754,37 @@ impl Server { Ok(executed) } + /// Record a client's parameter change, so we know what to undo before + /// this connection goes to somebody else. + pub fn record_params(&mut self, params: &[SetParam], in_transaction: bool) { + for param in params { + match (¶m.value, in_transaction) { + (Some(value), true) => { + self.client_params + .insert_transaction(¶m.name, value.clone(), param.local); + } + (Some(value), false) => { + self.client_params.insert(¶m.name, value.clone()); + } + (None, true) => self.client_params.reset_transaction(¶m.name), + (None, false) => self.client_params.reset(¶m.name), + } + } + + self.params_recorded = true; + } + + /// Record a client's `RESET ALL`. + pub fn record_reset_all(&mut self, in_transaction: bool) { + if in_transaction { + self.client_params.reset_all_transaction(); + } else { + self.client_params.reset_all(); + } + + self.params_recorded = true; + } + // Handle COMMIT/ROLLBACK for in-transaction params tracking. pub fn transaction_params_hook(&mut self, rollback: bool) { if rollback { @@ -1320,8 +1359,10 @@ pub mod test { }; use crate::{ - backend::pool::token_cache::TokenCache, config::Memory, frontend::PreparedStatements, - net::*, + backend::pool::token_cache::TokenCache, + config::Memory, + frontend::PreparedStatements, + net::{parameter::ParameterValue, *}, }; use super::{Error, *}; @@ -1365,6 +1406,7 @@ pub mod test { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, re_synced: false, replication_mode: false, @@ -3096,6 +3138,155 @@ pub mod test { ) } + #[tokio::test] + async fn test_recorded_reset_survives_rollback() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", ""); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + server.execute("BEGIN").await.unwrap(); + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + true, + ); + server.execute("RESET search_path").await.unwrap(); + server.execute("ROLLBACK").await.unwrap(); + server.transaction_params_hook(true); + + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "search_path""#]); + } + + #[tokio::test] + async fn test_recorded_reset_committed_is_permanent() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", ""); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + server.execute("BEGIN").await.unwrap(); + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + true, + ); + server.execute("RESET search_path").await.unwrap(); + server.execute("COMMIT").await.unwrap(); + server.transaction_params_hook(false); + + assert!( + server + .client_params + .reset_queries(&Parameters::default()) + .is_empty() + ); + } + + #[tokio::test] + async fn test_recorded_set_is_undone_for_the_next_client() { + let mut server = test_server().await; + server + .link_client(FrontendPid::new(), &Parameters::default(), None) + .await + .unwrap(); + + server.record_params( + &[SetParam { + name: "statement_timeout".into(), + value: Some(ParameterValue::String("5s".into())), + local: false, + }], + false, + ); + server + .execute("SET statement_timeout TO '5s'") + .await + .unwrap(); + + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "statement_timeout""#]); + } + + #[tokio::test] + async fn test_reset_keeps_other_recorded_params() { + let mut server = test_server().await; + server + .link_client(FrontendPid::new(), &Parameters::default(), None) + .await + .unwrap(); + + server.record_params( + &[SetParam { + name: "statement_timeout".into(), + value: Some(ParameterValue::String("5s".into())), + local: false, + }], + false, + ); + server + .execute("SET statement_timeout TO '5s'") + .await + .unwrap(); + + server.record_params( + &[SetParam { + name: "search_path".into(), + value: None, + local: false, + }], + false, + ); + server.execute("RESET search_path").await.unwrap(); + + let queries = server + .client_params + .reset_queries(&Parameters::default()) + .into_iter() + .map(|query| query.query().to_string()) + .collect::>(); + assert_eq!(queries, vec![r#"RESET "statement_timeout""#]); + } + + #[tokio::test] + async fn test_untracked_reset_still_clears_client_params() { + let mut server = test_server().await; + let mut params = Parameters::default(); + params.insert("search_path", "public"); + server + .link_client(FrontendPid::new(), ¶ms, None) + .await + .unwrap(); + + // A RESET we never recorded, as when the query parser is off. + server.execute("RESET search_path").await.unwrap(); + + assert!(server.client_params.is_empty()); + } + #[tokio::test] async fn test_reset_clears_client_params() { let mut server = test_server().await; diff --git a/pgdog/src/frontend/client/query_engine/mod.rs b/pgdog/src/frontend/client/query_engine/mod.rs index d90994c3a..47366627c 100644 --- a/pgdog/src/frontend/client/query_engine/mod.rs +++ b/pgdog/src/frontend/client/query_engine/mod.rs @@ -236,12 +236,10 @@ impl QueryEngine { } Command::Unlisten(channel) => self.unlisten(context, &channel.clone()).await?, Command::Set { - params, - behave_like_select, - .. + params, is_select, .. } => { let params = params.clone(); - self.set(context, ¶ms, *behave_like_select).await?; + self.set(context, ¶ms, *is_select).await?; } Command::ResetAll => { self.reset_all(context).await?; diff --git a/pgdog/src/frontend/client/query_engine/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index dff7cfde8..a9f76eae2 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -15,13 +15,27 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, params: &[SetParam], - behave_like_select: bool, + is_select: bool, ) -> Result<(), Error> { // Make sure client isn't changing route mid-transaction. if self.route_change_check(context, params).await? { return Ok(()); } + // Take a server before touching the parameters: syncing a change the + // statement is about to make itself would send it twice. + if is_select && !self.backend.connected() { + let connected = if context.in_transaction() { + self.connect_transaction(context).await? + } else { + self.connect(context, None).await? + }; + + if !connected { + return Ok(()); + } + } + let mut fake_command = "SET"; for param in params { let is_pin = param.name == PGDOG_PIN; @@ -44,7 +58,11 @@ impl QueryEngine { } } else { fake_command = "RESET"; - context.params.reset(¶m.name); + if context.in_transaction() { + context.params.reset_transaction(¶m.name); + } else { + context.params.reset(¶m.name); + } if is_pin { self.manual_lock = false; } @@ -56,11 +74,12 @@ impl QueryEngine { } if self.backend.connected() { + // The server is ours, so its session changes with the client's: + // record it or we won't know what to undo for the next client. + self.backend.record_params(params, context.in_transaction()); self.execute(context).await?; } else { - let values_to_return = - behave_like_select.then(|| params.iter().map(|p| p.value.as_ref())); - self.fake_command_response(context, fake_command, values_to_return) + self.fake_command_response(context, fake_command, None::>) .await?; } @@ -96,9 +115,14 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, ) -> Result<(), Error> { - context.params.reset_all(); + if context.in_transaction() { + context.params.reset_all_transaction(); + } else { + context.params.reset_all(); + } if self.backend.connected() { + self.backend.record_reset_all(context.in_transaction()); self.execute(context).await?; } else { self.fake_command_response(context, "RESET", None::>) diff --git a/pgdog/src/frontend/router/parser/command.rs b/pgdog/src/frontend/router/parser/command.rs index 9a50cb9cc..2526e6811 100644 --- a/pgdog/src/frontend/router/parser/command.rs +++ b/pgdog/src/frontend/router/parser/command.rs @@ -32,7 +32,8 @@ pub enum Command { Set { params: Vec, route: Route, - behave_like_select: bool, + /// `SELECT set_config(...)`, not `SET`: Postgres has to answer it. + is_select: bool, }, ResetAll, InternalField { diff --git a/pgdog/src/frontend/router/parser/query/set.rs b/pgdog/src/frontend/router/parser/query/set.rs index 08b652409..b28d86602 100644 --- a/pgdog/src/frontend/router/parser/query/set.rs +++ b/pgdog/src/frontend/router/parser/query/set.rs @@ -27,7 +27,7 @@ impl QueryParser { Ok(Command::Set { params: vec![param], route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, + is_select: false, }) } } @@ -99,7 +99,7 @@ impl QueryParser { Ok(Some(Command::Set { params, route: Route::write(context.shards_calculator.shard()), - behave_like_select: false, + is_select: false, })) } } diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 7f3540f4e..583965405 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -1,20 +1,21 @@ use super::*; +use crate::net::messages::{Bind, Format}; impl QueryParser { /// Handle SELECT set_config('key', 'value', is_local) /// - /// If the function arguments are a form we cannot handle, we warn and - /// pass through + /// Arguments we can't resolve leave it an ordinary query: it runs, but + /// we don't learn what it changed. pub(super) fn set_config( &mut self, fcall: &nodes::FuncCall, context: &QueryParserContext, ) -> Command { - if let Some(param) = parse_args(fcall) { + if let Some(param) = parse_args(fcall, context.router_context.bind) { Command::Set { params: vec![param], route: Route::write(context.shards_calculator.shard()), - behave_like_select: true, + is_select: true, } } else { Command::Query( @@ -25,25 +26,59 @@ impl QueryParser { } /// Returns None if the arguments could not be parsed -fn parse_args(fcall: &nodes::FuncCall) -> Option { - let name = parse_config_name(fcall.args().first()?)?; - let value = parse_config_value(fcall.args().get(1)?)?; - let local = parse_is_local(fcall.args().get(2)?)?; +fn parse_args(fcall: &nodes::FuncCall, bind: Option<&Bind>) -> Option { + let name = parse_config_name(fcall.args().first()?, bind)?; + let value = parse_config_value(fcall.args().get(1)?, bind)?; + let local = parse_is_local(fcall.args().get(2)?, bind)?; Some(SetParam { name, value, local }) } +/// Value bound to `$number`; the inner Option is the SQL NULL. +fn bound_text(bind: Option<&Bind>, number: i32) -> Option> { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = bind?.parameter(index).ok()??; + + if param.is_null() { + Some(None) + } else { + Some(Some(param.text()?.to_owned())) + } +} + +fn bound_bool(bind: Option<&Bind>, number: i32) -> Option { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = bind?.parameter(index).ok()??; + + if param.is_null() { + return None; + } + + match param.format() { + Format::Binary => match param.data() { + [0] => Some(false), + [1] => Some(true), + _ => None, + }, + Format::Text => match param.text()?.trim().to_lowercase().as_str() { + "t" | "true" | "y" | "yes" | "on" | "1" => Some(true), + "f" | "false" | "n" | "no" | "off" | "0" => Some(false), + _ => None, + }, + } +} + /// Returns None if the name could not be parsed -fn parse_config_name(arg: Node<'_>) -> Option { +fn parse_config_name(arg: Node<'_>, bind: Option<&Bind>) -> Option { match arg { Node::A_Const(c) => c.val()?.string_value().map(ToOwned::to_owned), - // Only constant strings can be handled for now + Node::ParamRef(nodes::ParamRef { number, .. }) => bound_text(bind, *number)?, _ => None, } } /// Returns None if the value could not be parsed, Some(None) if the value /// is NULL, and Some if the value was successfully parsed -fn parse_config_value(arg: Node<'_>) -> Option> { +fn parse_config_value(arg: Node<'_>, bind: Option<&Bind>) -> Option> { match arg { Node::A_Const(c) => match c.val() { Some(value) => Some(Some(ParameterValue::String( @@ -51,14 +86,18 @@ fn parse_config_value(arg: Node<'_>) -> Option> { ))), None => Some(None), }, + Node::ParamRef(nodes::ParamRef { number, .. }) => { + Some(bound_text(bind, *number)?.map(ParameterValue::String)) + } _ => None, } } /// Returns None if the node was not a constant boolean -fn parse_is_local(arg: Node<'_>) -> Option { +fn parse_is_local(arg: Node<'_>, bind: Option<&Bind>) -> Option { match arg { Node::A_Const(c) => c.val()?.bool_value(), + Node::ParamRef(nodes::ParamRef { number, .. }) => bound_bool(bind, *number), _ => None, } } diff --git a/pgdog/src/frontend/router/parser/query/test/test_set.rs b/pgdog/src/frontend/router/parser/query/test/test_set.rs index d8094e4cc..5586defba 100644 --- a/pgdog/src/frontend/router/parser/query/test/test_set.rs +++ b/pgdog/src/frontend/router/parser/query/test/test_set.rs @@ -7,7 +7,7 @@ use crate::{ route::{OverrideReason, ShardSource}, }, }, - net::parameter::ParameterValue, + net::{Format, messages::Parameter, parameter::ParameterValue}, }; use super::setup::*; @@ -66,15 +66,13 @@ fn test_set_config_null_value() { match command { Command::Set { - params, - behave_like_select, - .. + params, is_select, .. } => { assert_eq!(params.len(), 1); assert_eq!(params[0].name, "lock_timeout"); assert_eq!(params[0].value, None); assert!(!params[0].local); - assert!(behave_like_select); + assert!(is_select); } _ => panic!("expected Command::Set, got {command:#?}"), } @@ -252,3 +250,157 @@ fn test_single_shard_set() { _ => panic!("not a set"), } } + +#[test] +fn test_set_config_bound_params() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named( + "__test_set_config", + "SELECT pg_catalog.set_config($1, $2, $3)", + ) + .into(), + Bind::new_params( + "__test_set_config", + &[ + Parameter::new(b"search_path"), + Parameter::new(b""), + Parameter::new(b"f"), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { + ref params, + is_select, + .. + } => { + assert_eq!(params.len(), 1); + assert_eq!(params[0].name, "search_path"); + assert_eq!(params[0].value, Some(ParameterValue::String("".into()))); + assert!(!params[0].local); + assert!(is_select); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_bound_null_value() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_set_config_null", "SELECT set_config($1, $2, false)").into(), + Bind::new_params( + "__test_set_config_null", + &[Parameter::new(b"lock_timeout"), Parameter::new_null()], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { ref params, .. } => { + assert_eq!(params[0].name, "lock_timeout"); + assert_eq!(params[0].value, None); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_unresolvable_args_stay_a_query() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT pg_catalog.set_config($1, $2, false)").into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} + +#[test] +fn test_set_config_expression_stays_a_query() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Query::new("SELECT set_config('search_path', current_setting('search_path'), false)") + .into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} + +#[test] +fn test_set_config_bound_binary_is_local() { + let mut test = QueryParserTest::new(); + + let command = test.execute(vec![ + Parse::named("__test_set_config_bin", "SELECT set_config($1, $2, $3)").into(), + Bind::new_params_codes( + "__test_set_config_bin", + &[ + Parameter::new(b"statement_timeout"), + Parameter::new(b"1000"), + Parameter::new(&[1]), + ], + &[Format::Text, Format::Text, Format::Binary], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + match command { + Command::Set { ref params, .. } => { + assert_eq!(params[0].name, "statement_timeout"); + assert_eq!(params[0].value, Some(ParameterValue::String("1000".into()))); + assert!(params[0].local, "binary true must be read as SET LOCAL"); + } + _ => panic!("expected Command::Set, got {command:#?}"), + } +} + +#[test] +fn test_set_config_bound_non_utf8_stays_a_query() { + let mut test = QueryParserTest::new(); + + // set_config() takes text; what we can't read as text we can't track. + let command = test.execute(vec![ + Parse::named( + "__test_set_config_bytes", + "SELECT set_config($1, $2, false)", + ) + .into(), + Bind::new_params( + "__test_set_config_bytes", + &[ + Parameter::new(b"search_path"), + Parameter::new(&[0xff, 0xfe]), + ], + ) + .into(), + Execute::new().into(), + Sync.into(), + ]); + + assert!( + matches!(command, Command::Query(_)), + "expected Command::Query, got {command:#?}", + ); + assert!(command.route().is_write()); +} diff --git a/pgdog/src/net/parameter.rs b/pgdog/src/net/parameter.rs index 6dbe03a0f..eb5654c0e 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -267,11 +267,24 @@ impl Parameters { } } - /// Remove parameter from params temporarily. The transaction - /// is comitted, it will be removed permanently. + /// Remove a parameter permanently. pub fn reset(&mut self, name: impl ToString) { let name = name.to_string().to_lowercase(); + if self.params.remove(&name).is_some() { + self.hash = Self::compute_hash(&self.params); + } + + self.transaction_params.remove(&name); + self.transaction_local_params.remove(&name); + // Nothing left to restore on a rollback. + self.reset_params.remove(&name); + } + + /// Remove a parameter until the transaction ends: a ROLLBACK brings it back. + pub fn reset_transaction(&mut self, name: impl ToString) { + let name = name.to_string().to_lowercase(); + if let Some(value) = self.params.remove(&name) { self.reset_params.insert(name.clone(), value); self.hash = Self::compute_hash(&self.params); @@ -283,17 +296,32 @@ impl Parameters { /// Reset all tracked parameters. pub fn reset_all(&mut self) { - let mut keys: Vec = self.params.keys().cloned().collect(); - keys.extend(self.transaction_params.keys().cloned()); - keys.extend(self.transaction_local_params.keys().cloned()); + for key in self.resettable_keys() { + self.reset(&key); + } + } + + /// Reset all tracked parameters until the transaction ends. + pub fn reset_all_transaction(&mut self) { + for key in self.resettable_keys() { + self.reset_transaction(&key); + } + } + + fn resettable_keys(&self) -> Vec { + // Lifted out first: resetting borrows the maps we'd be iterating. + let mut keys: Vec = self + .params + .keys() + .chain(self.transaction_params.keys()) + .chain(self.transaction_local_params.keys()) + .filter(|key| !UNTRACKED_PARAMS.contains(key)) + .cloned() + .collect(); keys.sort(); keys.dedup(); - for key in keys { - if !UNTRACKED_PARAMS.contains(&key) { - self.reset(&key); - } - } + keys } /// Commit params we saved during the transaction. @@ -998,11 +1026,37 @@ mod test { } #[test] - fn test_reset_rollback_restores_param() { + fn test_reset_outside_transaction_is_permanent() { let mut params = Parameters::default(); params.insert("search_path", "public"); params.reset("search_path"); + + // A later transaction has nothing to do with that RESET. + params.rollback(); + + assert_eq!(params.get("search_path"), None); + } + + #[test] + fn test_reset_all_outside_transaction_is_permanent() { + let mut params = Parameters::default(); + params.insert("search_path", "public"); + params.insert("timezone", "UTC"); + + params.reset_all(); + params.rollback(); + + assert_eq!(params.get("search_path"), None); + assert_eq!(params.get("timezone"), None); + } + + #[test] + fn test_reset_rollback_restores_param() { + let mut params = Parameters::default(); + params.insert("search_path", "public"); + + params.reset_transaction("search_path"); assert_eq!(params.get("search_path"), None); params.rollback(); @@ -1104,7 +1158,7 @@ mod test { params.insert("search_path", "public"); params.insert("timezone", "UTC"); - params.reset_all(); + params.reset_all_transaction(); assert_eq!(params.get("search_path"), None); assert_eq!(params.get("timezone"), None);