From 57941cd702657c19ac2cd05df56b2a11ee10c268 Mon Sep 17 00:00:00 2001 From: arman-ddl <122062464+arman-ddl@users.noreply.github.com> Date: Thu, 20 Jun 2024 21:10:32 -0700 Subject: [PATCH] Allow extra headers in Barrage session creation (#5637) * Accept extra headers in Barrage session creation * Preserve extra header order * Update py/server/deephaven/barrage.py * Add extra headers to unit test * Move headers to own test * Force CI to rerun completely --------- Co-authored-by: Jianfeng Mao <4297243+jmao-denver@users.noreply.github.com> Co-authored-by: jianfengmao --- py/server/deephaven/barrage.py | 9 +++++++-- py/server/tests/test_barrage.py | 4 ++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/py/server/deephaven/barrage.py b/py/server/deephaven/barrage.py index 00e4a54d0bb..8680c718756 100644 --- a/py/server/deephaven/barrage.py +++ b/py/server/deephaven/barrage.py @@ -116,6 +116,7 @@ def barrage_session(host: str, auth_token: str = "", use_tls: bool = False, tls_root_certs: bytes = None, + extra_headers: Dict[str, str] = None ) -> BarrageSession: """Returns a Deephaven gRPC session to a remote server if a cached session is available; otherwise, creates a new session. @@ -135,6 +136,7 @@ def barrage_session(host: str, tls_root_certs (bytes): PEM encoded root certificates to use for TLS connection, or None to use system defaults. If not None implies use a TLS connection and the use_tls argument should have been passed as True. Defaults to None + extra_headers (Dict[str, str]): extra headers to set when configuring the gRPC channel. Defaults to None. Returns: a Deephaven Barrage session @@ -152,7 +154,7 @@ def barrage_session(host: str, else: target_uri = f"dh+plain://{target_uri}" - j_client_config = _build_client_config(target_uri, tls_root_certs) + j_client_config = _build_client_config(target_uri, tls_root_certs, extra_headers) auth = f"{auth_type} {auth_token}" try: @@ -209,9 +211,12 @@ def _get_barrage_session_direct(client_config: jpy.JType, auth: str) -> BarrageS return BarrageSession(j_barrage_session, j_channel) -def _build_client_config(target_uri: str, tls_root_certs: bytes) -> jpy.JType: +def _build_client_config(target_uri: str, tls_root_certs: bytes, extra_headers: Dict[str, str] = None) -> jpy.JType: j_client_config_builder = _JClientConfig.builder() j_client_config_builder.target(_JDeephavenTarget.of(_JURI(target_uri))) + if extra_headers: + for header, value in extra_headers.items(): + j_client_config_builder.putExtraHeaders(header, value) if tls_root_certs: j_ssl_config = _JSSLConfig.builder().trust( _JTrustCustom.ofX509(tls_root_certs, 0, len(tls_root_certs))).build() diff --git a/py/server/tests/test_barrage.py b/py/server/tests/test_barrage.py index b3411f60ad7..28c0a49b6aa 100644 --- a/py/server/tests/test_barrage.py +++ b/py/server/tests/test_barrage.py @@ -68,6 +68,10 @@ def test_barrage_session(self): with self.assertRaises(DHError): barrage_session(host="localhost", port=10000, auth_type="Basic", auth_token="user:password") + def test_barrage_session_with_extra_headers(self): + session = barrage_session(host="localhost", port=10000, auth_type="Anonymous", extra_headers={"envoy-prefix": "test"}) + self.assertIsNotNone(session) + def test_subscribe(self): session = barrage_session(host="localhost", port=10000, auth_type="Anonymous") t = session.subscribe(ticket=self.shared_ticket.bytes)