http2 test

neondatabase · Jan 12, 2024 · 0f428a6 · 0f428a6
1 parent e5ecf35
commit 0f428a6
Show file tree

Hide file tree

Showing 5 changed files with 170 additions and 1 deletion.
diff --git a/poetry.lock b/poetry.lock
diff --git a/proxy/src/serverless.rs b/proxy/src/serverless.rs
@@ -107,6 +107,7 @@ pub async fn task_main(
             let client_addr = io.client_addr();
             let remote_addr = io.inner.remote_addr();
             let sni_name = tls.server_name().map(|s| s.to_string());
+            let protocol = tls.alpn_protocol().map(|s| String::from_utf8_lossy(s).into_owned());
             let conn_pool = conn_pool.clone();
             let ws_connections = ws_connections.clone();
             let endpoint_rate_limiter = endpoint_rate_limiter.clone();
@@ -120,6 +121,7 @@ pub async fn task_main(
                 Ok(MetricService::new(hyper::service::service_fn(
                     move |req: Request<Body>| {
                         let sni_name = sni_name.clone();
+                        let protocol = protocol.clone();
                         let conn_pool = conn_pool.clone();
                         let ws_connections = ws_connections.clone();
                         let endpoint_rate_limiter = endpoint_rate_limiter.clone();
@@ -144,6 +146,7 @@ pub async fn task_main(
                                 "serverless",
                                 session = %session_id,
                                 %peer_addr,
+                                http_protocol = ?protocol,
                             ))
                             .await
                         }

diff --git a/pyproject.toml b/pyproject.toml
@@ -39,6 +39,7 @@ types-pytest-lazy-fixture = "^0.6.3.3"
 pytest-split = "^0.8.1"
 zstandard = "^0.21.0"
 websockets = "^12.0"
+httpx = {extras = ["http2"], version = "^0.26.0"}
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.3.0"

diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
@@ -24,6 +24,7 @@
 
 import asyncpg
 import backoff
+import httpx
 import jwt
 import psycopg2
 import pytest
@@ -2471,6 +2472,33 @@ def http_query(self, query, args, **kwargs):
             assert response.status_code == kwargs["expected_code"], f"response: {response.json()}"
         return response.json()
 
+    async def http2_query(self, query, args, **kwargs):
+        # TODO maybe use default values if not provided
+        user = kwargs["user"]
+        password = kwargs["password"]
+        expected_code = kwargs.get("expected_code")
+
+        connstr = f"postgresql://{user}:{password}@{self.domain}:{self.proxy_port}/postgres"
+        async with httpx.AsyncClient(
+            http2=True, verify=str(self.test_output_dir / "proxy.crt")
+        ) as client:
+            response = await client.post(
+                f"https://{self.domain}:{self.external_http_port}/sql",
+                data=json.dumps({"query": query, "params": args}),
+                headers={
+                    "Content-Type": "application/sql",
+                    "Neon-Connection-String": connstr,
+                    "Neon-Pool-Opt-In": "true",
+                },
+            )
+            assert response.http_version == "HTTP/2"
+
+            if expected_code is not None:
+                assert (
+                    response.status_code == kwargs["expected_code"]
+                ), f"response: {response.json()}"
+            return response.json()
+
     def get_metrics(self) -> str:
         request_result = requests.get(f"http://{self.host}:{self.http_port}/metrics")
         request_result.raise_for_status()

diff --git a/test_runner/regress/test_proxy.py b/test_runner/regress/test_proxy.py
@@ -500,3 +500,13 @@ def query(status: int, query: str) -> Any:
         "select array['foo'::foo, 'bar'::foo, 'baz'::foo] as data",
     )
     assert response["rows"][0]["data"] == ["foo", "bar", "baz"]
+
+
+@pytest.mark.asyncio
+async def test_sql_over_http2(static_proxy: NeonProxy):
+    static_proxy.safe_psql("create role http with login password 'http' superuser")
+
+    resp = await static_proxy.http2_query(
+        "select 42 as answer", [], user="http", password="http", expected_code=200
+    )
+    assert resp["rows"] == [{"answer": 42}]