diff --git a/integration/pgdog.toml b/integration/pgdog.toml index 429b7ffee..e748f0bf2 100644 --- a/integration/pgdog.toml +++ b/integration/pgdog.toml @@ -59,6 +59,17 @@ 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. + +[[databases]] +name = "pgdog_leak" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + # ------------------------------------------------------------------------------ # ----- Database :: pgdog_sharded ---------------------------------------------- @@ -505,3 +516,9 @@ 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" diff --git a/integration/python/gate/pgdog.toml b/integration/python/gate/pgdog.toml new file mode 100644 index 000000000..6da5a5a34 --- /dev/null +++ b/integration/python/gate/pgdog.toml @@ -0,0 +1,18 @@ +[general] +prepared_statements = "extended" + +[[databases]] +name = "pgdog" +host = "127.0.0.1" + +[[databases]] +name = "pgdog_leak_auto" +host = "127.0.0.1" +database_name = "pgdog" +pool_size = 1 +min_pool_size = 1 + +[[query_parsers]] +database = "pgdog_leak_auto" +level = "auto" +engine = "pg_query_raw" diff --git a/integration/python/gate/users.toml b/integration/python/gate/users.toml new file mode 100644 index 000000000..2a86d85b4 --- /dev/null +++ b/integration/python/gate/users.toml @@ -0,0 +1,9 @@ +[[users]] +name = "pgdog" +database = "pgdog" +password = "pgdog" + +[[users]] +name = "pgdog" +database = "pgdog_leak_auto" +password = "pgdog" diff --git a/integration/python/run.sh b/integration/python/run.sh index 89760d383..efb59a3b2 100644 --- a/integration/python/run.sh +++ b/integration/python/run.sh @@ -11,3 +11,13 @@ wait_for_pgdog source ${SCRIPT_DIR}/dev.sh stop_pgdog +run_pgdog integration/python/gate +wait_for_pgdog + +pushd ${SCRIPT_DIR} +source venv/bin/activate +PGDOG_LEAK_DATABASES=pgdog_leak_auto pytest -x test_session_params_leak.py +deactivate +popd + +stop_pgdog diff --git a/integration/python/test_session_params_leak.py b/integration/python/test_session_params_leak.py new file mode 100644 index 000000000..6e9bb084e --- /dev/null +++ b/integration/python/test_session_params_leak.py @@ -0,0 +1,184 @@ +"""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 at two parser levels: "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. "auto" needs its own config (gate/), where nothing else turns the parser +on; run.sh runs these tests a second time against it. +""" + +import os +import uuid + +import psycopg +import pytest + +DATABASES = os.environ.get("PGDOG_LEAK_DATABASES", "pgdog_leak").split(",") + + +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 + + +@pytest.fixture +def gated(request, dbname): + if dbname == "pgdog_leak_auto": + request.node.add_marker( + pytest.mark.xfail( + strict=True, + reason="the regex gate doesn't match set_config() yet (#1327)", + ) + ) + + +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, gated): + """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, gated): + """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..360246b5b 100644 --- a/integration/users.toml +++ b/integration/users.toml @@ -3,6 +3,11 @@ name = "pgdog" database = "pgdog" password = "pgdog" +[[users]] +name = "pgdog" +database = "pgdog_leak" +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 4c32e8868..d54a5e504 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}, @@ -443,6 +443,26 @@ impl Binding { } } + pub(crate) 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)), + _ => (), + } + } + + pub(crate) 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(crate) 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 08fa087ac..325a21ebc 100644 --- a/pgdog/src/backend/server.rs +++ b/pgdog/src/backend/server.rs @@ -21,7 +21,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::{ @@ -141,6 +141,9 @@ pub(crate) 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, @@ -432,6 +435,7 @@ impl Server { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, statement_executed: false, re_synced: false, @@ -617,6 +621,8 @@ impl Server { match message.code() { 'Z' => { + self.params_recorded = false; + let now = Instant::now(); let rfq = ReadyForQuery::from_bytes(message.payload())?; @@ -677,7 +683,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); @@ -772,6 +780,36 @@ impl Server { Ok(executed) } + /// Record a client's parameter change, so we know what to undo before + /// this connection goes to somebody else. + pub(crate) 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; + } + + pub(crate) 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(crate) fn transaction_params_hook(&mut self, rollback: bool) { if rollback { @@ -1327,7 +1365,7 @@ pub(crate) mod test { backend::pool::token_cache::TokenCache, config::Memory, frontend::{PreparedStatements, RewritePlan}, - net::{Prepare, *}, + net::{Prepare, parameter::ParameterValue, *}, }; use super::{Error, *}; @@ -1365,6 +1403,7 @@ pub(crate) mod test { streaming: false, schema_changed: false, sync_prepared: false, + params_recorded: false, in_transaction: false, re_synced: false, replication_mode: false, @@ -3150,6 +3189,155 @@ pub(crate) 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/set.rs b/pgdog/src/frontend/client/query_engine/set.rs index 6bde6f2b7..496599396 100644 --- a/pgdog/src/frontend/client/query_engine/set.rs +++ b/pgdog/src/frontend/client/query_engine/set.rs @@ -24,6 +24,20 @@ impl QueryEngine { return Ok(()); } + // Take a server before touching the parameters: syncing a change the + // statement is about to make itself would send it twice. + if set_config && !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; @@ -46,7 +60,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; } @@ -58,6 +76,9 @@ 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, None).await?; } else { let fake_response = set_config @@ -99,7 +120,9 @@ impl QueryEngine { &mut self, context: &mut QueryEngineContext<'_>, ) -> Result<(), Error> { - if context.in_transaction() || self.backend.connected() { + if context.in_transaction() { + context.params.reset_all_transaction(); + } else if self.backend.connected() { context.params.reset_all(); } else { context.params.restore_startup(context.startup_params); @@ -107,6 +130,7 @@ impl QueryEngine { } if self.backend.connected() { + self.backend.record_reset_all(context.in_transaction()); self.execute(context, None).await?; } else { self.fake_command_response(context, "RESET", None).await?; diff --git a/pgdog/src/frontend/router/parser/query/set_config.rs b/pgdog/src/frontend/router/parser/query/set_config.rs index 7aad92b0b..d0b6f1409 100644 --- a/pgdog/src/frontend/router/parser/query/set_config.rs +++ b/pgdog/src/frontend/router/parser/query/set_config.rs @@ -1,16 +1,17 @@ use super::*; +use crate::net::messages::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()), @@ -24,26 +25,65 @@ 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, + params: Option>, +) -> Option { + let name = parse_config_name(fcall.args().first()?, params)?; + let value = parse_config_value(fcall.args().get(1)?, params)?; + let local = parse_is_local(fcall.args().get(2)?, params)?; Some(SetParam { name, value, local }) } +/// Value bound to `$number`; the inner Option is the SQL NULL. +fn bound_text(params: Option>, number: i32) -> Option> { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = params?.parameter(index).ok()??; + + if param.is_null() { + Some(None) + } else { + Some(Some(param.text()?.to_owned())) + } +} + +fn bound_bool(params: Option>, number: i32) -> Option { + let index = usize::try_from(number).ok()?.checked_sub(1)?; + let param = params?.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<'_>, params: Option>) -> 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(params, *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<'_>, + params: Option>, +) -> Option> { match arg { Node::A_Const(c) => match c.val() { Some(value) => Some(Some(ParameterValue::String( @@ -51,14 +91,18 @@ fn parse_config_value(arg: Node<'_>) -> Option> { ))), None => Some(None), }, + Node::ParamRef(nodes::ParamRef { number, .. }) => { + Some(bound_text(params, *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<'_>, params: Option>) -> Option { match arg { Node::A_Const(c) => c.val()?.bool_value(), + Node::ParamRef(nodes::ParamRef { number, .. }) => bound_bool(params, *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 8917f2993..9c344a97e 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::bind::Parameter, parameter::ParameterValue}, }; use super::Error; @@ -284,3 +284,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, + set_config, + .. + } => { + 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!(set_config); + } + _ => 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 fc808f47e..801339782 100644 --- a/pgdog/src/net/parameter.rs +++ b/pgdog/src/net/parameter.rs @@ -247,11 +247,24 @@ impl Parameters { } } - /// Remove parameter from params temporarily. The transaction - /// is comitted, it will be removed permanently. + /// Remove a parameter permanently. pub(crate) fn reset(&mut self, name: impl AsRef) { let name = name.as_ref().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(crate) fn reset_transaction(&mut self, name: impl AsRef) { + let name = name.as_ref().to_lowercase(); + if let Some(value) = self.params.remove(&name) { self.reset_params.insert(name.clone(), value); self.hash = Self::compute_hash(&self.params); @@ -271,17 +284,31 @@ impl Parameters { /// Reset all tracked parameters. pub(crate) 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); + } + } + + pub(crate) 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. @@ -951,11 +978,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(); @@ -1057,7 +1110,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);