From 218be9eb32788b6ed46f56515bacfbba12b97ae5 Mon Sep 17 00:00:00 2001 From: George MacKerron Date: Tue, 15 Aug 2023 14:52:00 +0100 Subject: [PATCH] Added deferrable transaction option to http batch queries (#4993) ## Problem HTTP batch queries currently allow us to set the isolation level and read only, but not deferrable. ## Summary of changes Add support for deferrable. Echo deferrable status in response headers only if true. Likewise, now echo read-only status in response headers only if true. --- proxy/src/http/sql_over_http.rs | 30 +++++++++++++++++++++--------- test_runner/regress/test_proxy.py | 12 ++++++++++-- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/proxy/src/http/sql_over_http.rs b/proxy/src/http/sql_over_http.rs index aa06a91bba..33375e63e9 100644 --- a/proxy/src/http/sql_over_http.rs +++ b/proxy/src/http/sql_over_http.rs @@ -47,6 +47,7 @@ static ARRAY_MODE: HeaderName = HeaderName::from_static("neon-array-mode"); static ALLOW_POOL: HeaderName = HeaderName::from_static("neon-pool-opt-in"); static TXN_ISOLATION_LEVEL: HeaderName = HeaderName::from_static("neon-batch-isolation-level"); static TXN_READ_ONLY: HeaderName = HeaderName::from_static("neon-batch-read-only"); +static TXN_DEFERRABLE: HeaderName = HeaderName::from_static("neon-batch-deferrable"); static HEADER_VALUE_TRUE: HeaderValue = HeaderValue::from_static("true"); @@ -195,7 +196,7 @@ pub async fn handle( // Allow connection pooling only if explicitly requested let allow_pool = headers.get(&ALLOW_POOL) == Some(&HEADER_VALUE_TRUE); - // isolation level and read only + // isolation level, read only and deferrable let txn_isolation_level_raw = headers.get(&TXN_ISOLATION_LEVEL).cloned(); let txn_isolation_level = match txn_isolation_level_raw { @@ -209,8 +210,8 @@ pub async fn handle( None => None, }; - let txn_read_only_raw = headers.get(&TXN_READ_ONLY).cloned(); - let txn_read_only = txn_read_only_raw.as_ref() == Some(&HEADER_VALUE_TRUE); + let txn_read_only = headers.get(&TXN_READ_ONLY) == Some(&HEADER_VALUE_TRUE); + let txn_deferrable = headers.get(&TXN_DEFERRABLE) == Some(&HEADER_VALUE_TRUE); let request_content_length = match request.body().size_hint().upper() { Some(v) => v, @@ -247,6 +248,9 @@ pub async fn handle( if txn_read_only { builder = builder.read_only(true); } + if txn_deferrable { + builder = builder.deferrable(true); + } let transaction = builder.start().await?; for query in batch_query.queries { let result = query_to_json(&transaction, query, raw_output, array_mode).await; @@ -260,12 +264,20 @@ pub async fn handle( } transaction.commit().await?; let mut headers = HashMap::default(); - headers.insert( - TXN_READ_ONLY.clone(), - HeaderValue::try_from(txn_read_only.to_string())?, - ); - if let Some(txn_isolation_level_raw) = txn_isolation_level_raw { - headers.insert(TXN_ISOLATION_LEVEL.clone(), txn_isolation_level_raw); + if txn_read_only { + headers.insert( + TXN_READ_ONLY.clone(), + HeaderValue::try_from(txn_read_only.to_string())?, + ); + } + if txn_deferrable { + headers.insert( + TXN_DEFERRABLE.clone(), + HeaderValue::try_from(txn_deferrable.to_string())?, + ); + } + if let Some(txn_isolation_level) = txn_isolation_level_raw { + headers.insert(TXN_ISOLATION_LEVEL.clone(), txn_isolation_level); } Ok((json!({ "results": results }), headers)) } diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py index fabec6b5bc..dd767e14b7 100644 --- a/test_runner/regress/test_proxy.py +++ b/test_runner/regress/test_proxy.py @@ -265,7 +265,11 @@ def test_sql_over_http_output_options(static_proxy: NeonProxy): def test_sql_over_http_batch(static_proxy: NeonProxy): static_proxy.safe_psql("create role http with login password 'http' superuser") - def qq(queries: List[Tuple[str, Optional[List[Any]]]], read_only: bool = False) -> Any: + def qq( + queries: List[Tuple[str, Optional[List[Any]]]], + read_only: bool = False, + deferrable: bool = False, + ) -> Any: connstr = f"postgresql://http:http@{static_proxy.domain}:{static_proxy.proxy_port}/postgres" response = requests.post( f"https://{static_proxy.domain}:{static_proxy.external_http_port}/sql", @@ -277,6 +281,7 @@ def test_sql_over_http_batch(static_proxy: NeonProxy): "Neon-Connection-String": connstr, "Neon-Batch-Isolation-Level": "Serializable", "Neon-Batch-Read-Only": "true" if read_only else "false", + "Neon-Batch-Deferrable": "true" if deferrable else "false", }, verify=str(static_proxy.test_output_dir / "proxy.crt"), ) @@ -299,7 +304,8 @@ def test_sql_over_http_batch(static_proxy: NeonProxy): ) assert headers["Neon-Batch-Isolation-Level"] == "Serializable" - assert headers["Neon-Batch-Read-Only"] == "false" + assert "Neon-Batch-Read-Only" not in headers + assert "Neon-Batch-Deferrable" not in headers assert result[0]["rows"] == [{"answer": 42}] assert result[1]["rows"] == [{"answer": "42"}] @@ -327,8 +333,10 @@ def test_sql_over_http_batch(static_proxy: NeonProxy): ("select 42 as answer", None), ], True, + True, ) assert headers["Neon-Batch-Isolation-Level"] == "Serializable" assert headers["Neon-Batch-Read-Only"] == "true" + assert headers["Neon-Batch-Deferrable"] == "true" assert result[0]["rows"] == [{"answer": 42}]