From c103af74a26c95c5ae3b18653876055f29bbd736 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Tue, 3 Dec 2024 15:35:32 +0000 Subject: [PATCH] add test for neon_proxy_params_compat option --- proxy/src/compute.rs | 10 ++++++---- proxy/src/console_redirect_proxy.rs | 1 + proxy/src/proxy/connect_compute.rs | 4 +++- proxy/src/proxy/mod.rs | 14 ++++++++++++-- test_runner/fixtures/neon_fixtures.py | 2 +- test_runner/regress/test_proxy.py | 19 +++++++++++++++++++ 6 files changed, 42 insertions(+), 8 deletions(-) diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs index 6a9c5b4ce848b..08ddb6f1fe81e 100644 --- a/proxy/src/compute.rs +++ b/proxy/src/compute.rs @@ -131,9 +131,11 @@ impl ConnCfg { } /// Apply startup message params to the connection config. - pub(crate) fn set_startup_params(&mut self, params: &StartupMessageParams) { - let arbitrary_params = false; - + pub(crate) fn set_startup_params( + &mut self, + params: &StartupMessageParams, + arbitrary_params: bool, + ) { for (k, v) in params.iter() { match k { // Only set `user` if it's not present in the config. @@ -426,7 +428,7 @@ mod tests { let params = "project = foo"; assert_eq!(filtered_options(params).as_deref(), Some("project = foo")); - let params = "project = foo neon_endpoint_type:read_write neon_lsn:0/2"; + let params = "project = foo neon_endpoint_type:read_write neon_lsn:0/2 neon_proxy_params_compat:true"; assert_eq!(filtered_options(params).as_deref(), Some("project = foo")); } } diff --git a/proxy/src/console_redirect_proxy.rs b/proxy/src/console_redirect_proxy.rs index 8f78df19649ba..7db1179eeae83 100644 --- a/proxy/src/console_redirect_proxy.rs +++ b/proxy/src/console_redirect_proxy.rs @@ -206,6 +206,7 @@ pub(crate) async fn handle_client( let mut node = connect_to_compute( ctx, &TcpMechanism { + params_compat: true, params: ¶ms, locks: &config.connect_compute_locks, }, diff --git a/proxy/src/proxy/connect_compute.rs b/proxy/src/proxy/connect_compute.rs index 585dce7baeb8c..a3027abd7cae9 100644 --- a/proxy/src/proxy/connect_compute.rs +++ b/proxy/src/proxy/connect_compute.rs @@ -66,6 +66,8 @@ pub(crate) trait ComputeConnectBackend { } pub(crate) struct TcpMechanism<'a> { + pub(crate) params_compat: bool, + /// KV-dictionary with PostgreSQL connection params. pub(crate) params: &'a StartupMessageParams, @@ -92,7 +94,7 @@ impl ConnectMechanism for TcpMechanism<'_> { } fn update_connect_config(&self, config: &mut compute::ConnCfg) { - config.set_startup_params(self.params); + config.set_startup_params(self.params, self.params_compat); } } diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index af97fb3d71590..e8c7988fc948f 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -338,9 +338,15 @@ pub(crate) async fn handle_client( } }; + let params_compat = match &user_info { + auth::Backend::ControlPlane(_, info) => info.info.options.get("proxy_params_compat").is_some(), + auth::Backend::Local(_) => false, + }; + let mut node = connect_to_compute( ctx, &TcpMechanism { + params_compat, params: ¶ms, locks: &config.connect_compute_locks, }, @@ -415,13 +421,17 @@ impl NeonOptions { .map(Self::parse_from_iter) .unwrap_or_default() } + pub(crate) fn parse_options_raw(options: &str) -> Self { Self::parse_from_iter(StartupMessageParams::parse_options_raw(options)) } + pub(crate) fn get(&self, key: &str) -> Option { + self.0.iter().find_map(|(k, v)| (k == key).then_some(v)).cloned() + } + pub(crate) fn is_ephemeral(&self) -> bool { - // Currently, neon endpoint options are all reserved for ephemeral endpoints. - !self.0.is_empty() + self.0.iter().filter(|(k, _)| k != "proxy_params_compat").count() == 0 } fn parse_from_iter<'a>(options: impl Iterator) -> Self { diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index f55f06bebc000..721a617ea2e49 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -269,7 +269,7 @@ async def connect_async(self, **kwargs: Any) -> asyncpg.Connection: for match in re.finditer(r"-c(\w*)=(\w*)", options): key = match.group(1) val = match.group(2) - if "server_options" in conn_options: + if "server_settings" in conn_options: conn_options["server_settings"].update({key: val}) else: conn_options["server_settings"] = {key: val} diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index 5a01d90d85489..4baf5531cdecc 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from contextlib import closing import json import subprocess import time @@ -131,6 +132,24 @@ def test_proxy_options(static_proxy: NeonProxy, option_name: str): assert out[0][0] == " str" +@pytest.mark.asyncio +async def test_proxy_arbitrary_params(static_proxy: NeonProxy): + with closing( + await static_proxy.connect_async(server_settings={"IntervalStyle": "iso_8601"}) + ) as conn: + out = await conn.fetchval("select to_json('0 seconds'::interval)") + assert out == '"00:00:00"' + + options = "neon_proxy_params_compat:true" + with closing( + await static_proxy.connect_async( + server_settings={"IntervalStyle": "iso_8601", "options": options} + ) + ) as conn: + out = await conn.fetchval("select to_json('0 seconds'::interval)") + assert out == '"PT0S"' + + def test_auth_errors(static_proxy: NeonProxy): """ Check that we throw very specific errors in some unsuccessful auth scenarios.