You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2026/03/30 22:23:26 UTC
(pinot) branch master updated: Unify `headers.*` property handling across HTTP and gRPC JDBC drivers… (#18029)
This is an automated email from the ASF dual-hosted git repository.
xiangfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 71220243d97 Unify `headers.*` property handling across HTTP and gRPC JDBC drivers… (#18029)
71220243d97 is described below
commit 71220243d97ae6e09d2ca8da0a4f2792ea8dcfd8
Author: Nihal Jain <ni...@apache.org>
AuthorDate: Tue Mar 31 03:53:20 2026 +0530
Unify `headers.*` property handling across HTTP and gRPC JDBC drivers… (#18029)
* Unify `headers.*` property handling across HTTP and gRPC JDBC drivers and add comprehensive unit tests for auth and header propagation paths.
Fixes #18028
The gRPC JDBC driver (`jdbc:pinotgrpc://`) silently ignored `headers.*` prefixed properties (e.g., `headers.Authorization`), while the HTTP driver (`jdbc:pinot://`) handled them correctly. This caused auth failures when using `headers.Authorization` with the gRPC path.
With this change we do following:
- Add `headers.*` property extraction to `PinotGrpcConnection` via `ConnectionUtils.getHeadersFromProperties()`, matching HTTP driver behavior.
- Add comprehensive unit tests for both HTTP (jdbc:pinot://) and gRPC (jdbc:pinotgrpc://) JDBC driver paths covering all auth scenarios with full precedence chain: Authorization (direct property) > headers.Authorization > URL user/password > property user/password
* Update pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/AbstractJdbcDriverAuthTest.java
Co-authored-by: Copilot <17...@users.noreply.github.com>
* Update pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/HttpJdbcDriverAuthTest.java
Co-authored-by: Copilot <17...@users.noreply.github.com>
---------
Co-authored-by: Xiang Fu <xi...@gmail.com>
Co-authored-by: Copilot <17...@users.noreply.github.com>
---
.../pinot/client/grpc/PinotGrpcConnection.java | 7 +
.../pinot/client/AbstractJdbcDriverAuthTest.java | 348 +++++++++++++++++++++
.../pinot/client/HttpJdbcDriverAuthTest.java | 74 +++++
.../pinot/client/grpc/GrpcJdbcDriverAuthTest.java | 90 ++++++
4 files changed, 519 insertions(+)
diff --git a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcConnection.java b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcConnection.java
index a28f3679af0..1450f7d54d7 100644
--- a/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcConnection.java
+++ b/pinot-clients/pinot-jdbc-client/src/main/java/org/apache/pinot/client/grpc/PinotGrpcConnection.java
@@ -32,6 +32,7 @@ import org.apache.pinot.client.base.AbstractBaseConnection;
import org.apache.pinot.client.controller.PinotControllerTransport;
import org.apache.pinot.client.controller.PinotControllerTransportFactory;
import org.apache.pinot.client.controller.response.ControllerTenantBrokerResponse;
+import org.apache.pinot.client.utils.ConnectionUtils;
import org.apache.pinot.client.utils.DriverUtils;
import org.apache.pinot.spi.utils.CommonConstants;
import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
@@ -74,12 +75,18 @@ public class PinotGrpcConnection extends AbstractBaseConnection {
brokers = getBrokerGrpcList(controllerURL, tenant);
}
+
+ // Extract headers.* prefixed properties first (same behavior as HTTP path)
+ // This enables consistent header propagation across both HTTP and gRPC paths
+ _metadataMap.putAll(ConnectionUtils.getHeadersFromProperties(properties));
+
for (String possibleQueryOption : POSSIBLE_QUERY_OPTIONS) {
Object property = properties.getProperty(possibleQueryOption);
if (property != null) {
_queryOptions.put(possibleQueryOption, DriverUtils.parseOptionValue(property));
}
}
+
for (String possibleMetadataMapOption : POSSIBLE_METADATA_MAP_OPTIONS) {
for (Object key : properties.keySet()) {
if (key.toString().equalsIgnoreCase(possibleMetadataMapOption)) {
diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/AbstractJdbcDriverAuthTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/AbstractJdbcDriverAuthTest.java
new file mode 100644
index 00000000000..0ff27dcc073
--- /dev/null
+++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/AbstractJdbcDriverAuthTest.java
@@ -0,0 +1,348 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.client;
+
+import com.google.protobuf.ByteString;
+import io.grpc.Server;
+import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
+import io.grpc.stub.StreamObserver;
+import java.io.IOException;
+import java.net.ServerSocket;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.concurrent.CopyOnWriteArrayList;
+import org.apache.pinot.common.proto.Broker;
+import org.apache.pinot.common.proto.PinotQueryBrokerGrpc;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Abstract base class for testing auth and header propagation through
+ * {@link DriverManager#getConnection(String, Properties)} for both
+ * HTTP ({@code jdbc:pinot://}) and gRPC ({@code jdbc:pinotgrpc://}) driver paths.
+ *
+ * <p>Subclasses provide the JDBC URL scheme and implement a single method to extract
+ * headers from their respective connection types. All assertions are handled in this
+ * base class for consistency.
+ *
+ * <p>A mock gRPC server is started for the pinotgrpc path validation query;
+ * the HTTP path does not require a mock server.
+ */
+public abstract class AbstractJdbcDriverAuthTest {
+
+ protected Server _mockGrpcServer;
+ protected int _mockGrpcPort;
+ protected MetadataCapturingService _mockService;
+
+ /**
+ * Mock gRPC service that captures metadata from each incoming BrokerRequest.
+ */
+ protected static class MetadataCapturingService extends PinotQueryBrokerGrpc.PinotQueryBrokerImplBase {
+ private final List<Map<String, String>> _capturedMetadata = new CopyOnWriteArrayList<>();
+ private final List<String> _capturedSql = new CopyOnWriteArrayList<>();
+
+ @Override
+ public void submit(Broker.BrokerRequest request,
+ StreamObserver<Broker.BrokerResponse> responseObserver) {
+ _capturedMetadata.add(new HashMap<>(request.getMetadataMap()));
+ _capturedSql.add(request.getSql());
+
+ String responseJson = "{\"resultTable\":null,\"exceptions\":[],\"numDocsScanned\":0,"
+ + "\"totalDocs\":0,\"timeUsedMs\":1,\"numServersQueried\":1,\"numServersResponded\":1}";
+ responseObserver.onNext(Broker.BrokerResponse.newBuilder()
+ .setPayload(ByteString.copyFromUtf8(responseJson)).build());
+ responseObserver.onCompleted();
+ }
+
+ public List<Map<String, String>> getCapturedMetadata() {
+ return _capturedMetadata;
+ }
+
+ public List<String> getCapturedSql() {
+ return _capturedSql;
+ }
+
+ public Map<String, String> getLastCapturedMetadata() {
+ return _capturedMetadata.get(_capturedMetadata.size() - 1);
+ }
+
+ public void clear() {
+ _capturedMetadata.clear();
+ _capturedSql.clear();
+ }
+ }
+
+ @BeforeClass
+ public void setUp()
+ throws IOException {
+ _mockGrpcPort = findFreePort();
+ _mockService = new MetadataCapturingService();
+ _mockGrpcServer = NettyServerBuilder.forPort(_mockGrpcPort)
+ .addService(_mockService)
+ .build()
+ .start();
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws InterruptedException {
+ if (_mockGrpcServer != null) {
+ _mockGrpcServer.shutdown().awaitTermination();
+ }
+ }
+
+ /** Returns the JDBC URL for this driver variant. */
+ protected abstract String getJdbcUrl();
+
+ /** Returns base connection properties including brokers. */
+ protected abstract Properties getBaseProperties();
+
+ /**
+ * Extracts headers/metadata from the connection for verification.
+ * <p>For HTTP: extracts from {@code JsonAsyncHttpPinotClientTransport._headers}
+ * <p>For gRPC: extracts from {@code PinotGrpcConnection.getMetadataMap()}
+ *
+ * @param conn the JDBC connection
+ * @return map of header names to values
+ */
+ protected abstract Map<String, String> extractHeadersFromConnection(Connection conn)
+ throws Exception;
+
+ @Test
+ public void testUserPasswordAuth()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("user", "admin");
+ props.setProperty("password", "secret");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ String authValue = headers.get("Authorization");
+ assertNotNull(authValue, "Authorization should be present");
+ assertTrue(authValue.startsWith("Basic "),
+ "Authorization value should start with 'Basic ' but was: " + authValue);
+ }
+ }
+
+ @Test
+ public void testExplicitAuthorizationProperty()
+ throws Exception {
+ Properties props = getBaseProperties();
+ // Direct "Authorization" property — works for both HTTP and gRPC paths
+ props.setProperty("Authorization", "Bearer my-jwt-token-123");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ String authValue = headers.get("Authorization");
+ assertNotNull(authValue, "Authorization should be present");
+ assertTrue(authValue.startsWith("Bearer "),
+ "Authorization value should start with 'Bearer ' but was: " + authValue);
+ }
+ }
+
+ @Test
+ public void testExplicitAuthOverridesUserPassword()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("user", "admin");
+ props.setProperty("password", "secret");
+ props.setProperty("Authorization", "Bearer explicit-override");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ String authValue = headers.get("Authorization");
+ assertNotNull(authValue, "Authorization should be present");
+ assertTrue(authValue.startsWith("Bearer "),
+ "Explicit Authorization should override user/password. Expected 'Bearer ' but was: " + authValue);
+ }
+ }
+
+ /**
+ * Tests that {@code headers.Authorization} (prefixed form) is handled by both driver paths.
+ */
+ @Test
+ public void testHeadersPrefixAuthorizationPropagated()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("headers.Authorization", "Bearer prefixed-jwt-token");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ assertEquals(headers.get("Authorization"), "Bearer prefixed-jwt-token",
+ "headers.Authorization should be extracted and set as Authorization");
+ }
+ }
+
+ /**
+ * Tests that multiple {@code headers.*} prefixed properties are handled consistently.
+ * Uses dummy header names to avoid confusion with real Pinot headers.
+ */
+ @Test
+ public void testMultipleHeadersPrefixProperties()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("headers.Authorization", "Basic dummy-auth-token");
+ props.setProperty("headers.X-Dummy-Header-1", "dummy-value-1");
+ props.setProperty("headers.X-Dummy-Header-2", "dummy-value-2");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+
+ assertEquals(headers.get("Authorization"), "Basic dummy-auth-token",
+ "headers.Authorization should be extracted");
+ assertTrue(headers.containsKey("X-Dummy-Header-1"),
+ "Custom header 'X-Dummy-Header-1' should be present");
+ assertEquals(headers.get("X-Dummy-Header-1"), "dummy-value-1",
+ "Custom header 'X-Dummy-Header-1' should have expected value");
+ assertTrue(headers.containsKey("X-Dummy-Header-2"),
+ "Custom header 'X-Dummy-Header-2' should be present");
+ assertEquals(headers.get("X-Dummy-Header-2"), "dummy-value-2",
+ "Custom header 'X-Dummy-Header-2' should have expected value");
+ }
+ }
+
+ /**
+ * Tests that {@code headers.Authorization} takes precedence over {@code user}/{@code password}.
+ */
+ @Test
+ public void testHeadersPrefixAuthOverridesUserPassword()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("user", "admin");
+ props.setProperty("password", "secret");
+ props.setProperty("headers.Authorization", "Bearer prefix-wins");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ assertEquals(headers.get("Authorization"), "Bearer prefix-wins",
+ "headers.Authorization should override user/password");
+ }
+ }
+
+ /**
+ * Tests that username and password specified in the JDBC URL query string
+ * (e.g., {@code jdbc:pinot://localhost:9000?user=Foo&password=Bar}) produce
+ * a valid Basic auth header.
+ */
+ @Test
+ public void testUserPasswordInUrl()
+ throws Exception {
+ Properties props = getBaseProperties();
+ String urlWithCreds = getJdbcUrl() + "?user=UrlUser&password=UrlPass";
+
+ try (Connection conn = DriverManager.getConnection(urlWithCreds, props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ String authValue = headers.get("Authorization");
+ assertNotNull(authValue, "Authorization should be present from URL credentials");
+ assertTrue(authValue.startsWith("Basic "),
+ "URL user/password should produce Basic auth but was: " + authValue);
+ }
+ }
+
+ /**
+ * Tests that URL credentials override property credentials.
+ * Per the Javadoc: (username and password in URL) > (user and password specified in properties)
+ */
+ @Test
+ public void testUrlCredentialsOverridePropertyCredentials()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("user", "PropUser");
+ props.setProperty("password", "PropPass");
+ String urlWithCreds = getJdbcUrl() + "?user=UrlUser&password=UrlPass";
+
+ try (Connection conn = DriverManager.getConnection(urlWithCreds, props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ String authValue = headers.get("Authorization");
+ assertNotNull(authValue, "Authorization should be present");
+ assertTrue(authValue.startsWith("Basic "),
+ "Should produce Basic auth but was: " + authValue);
+ // Decode and verify the URL credentials won, not property credentials
+ String expectedToken =
+ org.apache.pinot.common.auth.BasicAuthTokenUtils.toBasicAuthToken("UrlUser", "UrlPass");
+ assertEquals(authValue, expectedToken,
+ "URL credentials should override property credentials");
+ }
+ }
+
+ /**
+ * Tests that a direct {@code Authorization} property takes precedence over
+ * {@code headers.Authorization}. In {@link org.apache.pinot.client.utils.DriverUtils#handleAuth},
+ * the direct {@code Authorization} property is applied last, unconditionally overriding
+ * any value already in the headers map (including one extracted from {@code headers.Authorization}).
+ */
+ @Test
+ public void testExplicitAuthOverridesHeadersPrefix()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("headers.Authorization", "Bearer from-headers-prefix");
+ props.setProperty("Authorization", "Bearer direct-wins");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ assertEquals(headers.get("Authorization"), "Bearer direct-wins",
+ "Direct Authorization property should override headers.Authorization");
+ }
+ }
+
+ /**
+ * Tests that {@code headers.Authorization} takes precedence over URL credentials.
+ * Per the Javadoc: (headers.Authorization property) > (username and password in URL)
+ */
+ @Test
+ public void testHeadersAuthOverridesUrlCredentials()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("headers.Authorization", "Bearer headers-wins-over-url");
+ String urlWithCreds = getJdbcUrl() + "?user=UrlUser&password=UrlPass";
+
+ try (Connection conn = DriverManager.getConnection(urlWithCreds, props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> headers = extractHeadersFromConnection(conn);
+ assertEquals(headers.get("Authorization"), "Bearer headers-wins-over-url",
+ "headers.Authorization should override URL credentials");
+ }
+ }
+
+ protected static int findFreePort()
+ throws IOException {
+ try (ServerSocket socket = new ServerSocket(0)) {
+ return socket.getLocalPort();
+ }
+ }
+}
diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/HttpJdbcDriverAuthTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/HttpJdbcDriverAuthTest.java
new file mode 100644
index 00000000000..f2fe2f9427d
--- /dev/null
+++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/HttpJdbcDriverAuthTest.java
@@ -0,0 +1,74 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.client;
+
+import java.lang.reflect.Field;
+import java.sql.Connection;
+import java.util.Map;
+import java.util.Properties;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertTrue;
+
+
+/**
+ * Tests for {@code jdbc:pinot://} (HTTP) driver path through {@link java.sql.DriverManager}.
+ *
+ * <p>Verifies that auth credentials are correctly propagated through the JDBC driver into
+ * the HTTP transport headers of {@link JsonAsyncHttpPinotClientTransport}.
+ *
+ * <p>Unlike the gRPC path, the HTTP connection does not send a validation query on creation,
+ * so no mock server is needed. Headers are verified via reflection on the transport's
+ * private {@code _headers} field.
+ *
+ * <p>The HTTP path (via {@link org.apache.pinot.client.utils.ConnectionUtils#getHeadersFromProperties})
+ * extracts and strips the {@code headers.*} prefix, setting custom headers on the transport.
+ */
+@Test
+public class HttpJdbcDriverAuthTest extends AbstractJdbcDriverAuthTest {
+
+ @Override
+ protected String getJdbcUrl() {
+ // Controller URL — not actually contacted since we set "brokers" property
+ return "jdbc:pinot://localhost:9000";
+ }
+
+ @Override
+ protected Properties getBaseProperties() {
+ Properties props = new Properties();
+ // Dummy broker to bypass controller discovery (no HTTP call is made on connect)
+ props.setProperty("brokers", "localhost:8099");
+ return props;
+ }
+
+ @Override
+ @SuppressWarnings("unchecked")
+ protected Map<String, String> extractHeadersFromConnection(Connection conn)
+ throws Exception {
+ PinotConnection pinotConn = (PinotConnection) conn;
+ org.apache.pinot.client.Connection session = pinotConn.getSession();
+ PinotClientTransport<?> transport = session.getTransport();
+ assertTrue(transport instanceof JsonAsyncHttpPinotClientTransport,
+ "Transport should be JsonAsyncHttpPinotClientTransport but was: " + transport.getClass().getName());
+
+ Field headersField = JsonAsyncHttpPinotClientTransport.class.getDeclaredField("_headers");
+ headersField.setAccessible(true);
+ return (Map<String, String>) headersField.get(transport);
+ }
+}
diff --git a/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/GrpcJdbcDriverAuthTest.java b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/GrpcJdbcDriverAuthTest.java
new file mode 100644
index 00000000000..05bbd9795bb
--- /dev/null
+++ b/pinot-clients/pinot-jdbc-client/src/test/java/org/apache/pinot/client/grpc/GrpcJdbcDriverAuthTest.java
@@ -0,0 +1,90 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.client.grpc;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.util.Map;
+import java.util.Properties;
+import org.apache.pinot.client.AbstractJdbcDriverAuthTest;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+
+/**
+ * Tests for {@code jdbc:pinotgrpc://} driver path through {@link java.sql.DriverManager}.
+ *
+ * <p>Verifies that auth credentials and gRPC metadata options (blockRowSize, encoding,
+ * compression) are correctly propagated through the JDBC driver into
+ * {@link PinotGrpcConnection#getMetadataMap()}.
+ *
+ * <p>Uses the mock gRPC server from {@link AbstractJdbcDriverAuthTest} to handle
+ * the validation query that the gRPC connection sends on construction.
+ *
+ * <p>The gRPC path extracts {@code headers.*} prefixed properties (same behavior as HTTP path),
+ * providing consistent header propagation across both JDBC driver variants.
+ */
+@Test
+public class GrpcJdbcDriverAuthTest extends AbstractJdbcDriverAuthTest {
+
+ @Override
+ protected String getJdbcUrl() {
+ // Controller URL — not actually contacted since we set "brokers" property
+ return "jdbc:pinotgrpc://localhost:" + _mockGrpcPort;
+ }
+
+ @Override
+ protected Properties getBaseProperties() {
+ Properties props = new Properties();
+ // Point brokers directly to the mock gRPC server so the validation query succeeds
+ props.setProperty("brokers", "localhost:" + _mockGrpcPort);
+ return props;
+ }
+
+ @Override
+ protected Map<String, String> extractHeadersFromConnection(Connection conn)
+ throws Exception {
+ PinotGrpcConnection grpcConn = (PinotGrpcConnection) conn;
+ return grpcConn.getMetadataMap();
+ }
+
+ @Test
+ public void testGrpcOptionsBlockRowSizeEncodingCompression()
+ throws Exception {
+ Properties props = getBaseProperties();
+ props.setProperty("Authorization", "Basic token");
+ props.setProperty("blockRowSize", "10000");
+ props.setProperty("encoding", "JSON");
+ props.setProperty("compression", "ZSTD");
+
+ try (Connection conn = DriverManager.getConnection(getJdbcUrl(), props)) {
+ assertNotNull(conn, "Connection should be created successfully");
+ Map<String, String> metadata = extractHeadersFromConnection(conn);
+
+ assertEquals(metadata.get("blockRowSize"), "10000",
+ "blockRowSize should be propagated into metadata map");
+ assertEquals(metadata.get("encoding"), "JSON",
+ "encoding should be propagated into metadata map");
+ assertEquals(metadata.get("compression"), "ZSTD",
+ "compression should be propagated into metadata map");
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org