This is an automated email from the ASF dual-hosted git repository. yuqi4733 pushed a commit to branch internal-main in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit e917ee1634471c407ad5b1a789ba1da96fae7b1c Author: geyanggang <[email protected]> AuthorDate: Fri Feb 6 16:59:02 2026 +0800 [#104] feat(bigquery-catalog): Add support proxy for BQ catalog. --- .../catalog/bigquery/BigQueryCatalog.java | 44 +++- .../BigQueryCatalogPropertiesMetadata.java | 37 ++++ .../catalog/bigquery/BigQueryClientPool.java | 83 ++++++- .../src/main/resources/jdbc-bigquery.conf | 9 +- .../catalog/bigquery/TestBigQueryCatalogProxy.java | 246 +++++++++++++++++++++ 5 files changed, 410 insertions(+), 9 deletions(-) diff --git a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalog.java b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalog.java index 85dee82a15..da6fc13183 100644 --- a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalog.java +++ b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalog.java @@ -93,7 +93,7 @@ public class BigQueryCatalog extends JdbcCatalog { /** * Build JDBC URL from individual BigQuery configuration components. If jdbc-url contains * authentication parameters, use it directly. Otherwise, enhance the base URL with project-id, - * jdbc-user (service account email), and jdbc-password (key file path). + * jdbc-user (service account email), and jdbc-password (key file path), and proxy configuration. */ private Map<String, String> buildJdbcUrl(Map<String, String> config) { @@ -112,6 +112,12 @@ public class BigQueryCatalog extends JdbcCatalog { String keyFilePath = config.get(JdbcConfig.PASSWORD.getKey()); // Key file path String serviceAccountEmail = config.get(JdbcConfig.USERNAME.getKey()); // Service account email + // Extract proxy configuration + String proxyHost = config.get(BigQueryCatalogPropertiesMetadata.PROXY_HOST); + String proxyPort = config.get(BigQueryCatalogPropertiesMetadata.PROXY_PORT); + String proxyUsername = config.get(BigQueryCatalogPropertiesMetadata.PROXY_USERNAME); + String proxyPassword = config.get(BigQueryCatalogPropertiesMetadata.PROXY_PASSWORD); + // Validate required properties if (StringUtils.isBlank(projectId)) { throw new IllegalArgumentException("project-id is required for BigQuery catalog"); @@ -121,6 +127,21 @@ public class BigQueryCatalog extends JdbcCatalog { "jdbc-password (key file path) is required for BigQuery catalog"); } + // Validate proxy configuration if provided + if (StringUtils.isNotBlank(proxyHost)) { + if (StringUtils.isBlank(proxyPort)) { + throw new IllegalArgumentException("proxy-port is required when proxy-host is specified"); + } + try { + int port = Integer.parseInt(proxyPort); + if (port <= 0 || port > 65535) { + throw new IllegalArgumentException("proxy-port must be between 1 and 65535"); + } + } catch (NumberFormatException e) { + throw new IllegalArgumentException("proxy-port must be a valid integer", e); + } + } + // Build complete JDBC URL with authentication parameters // Use OAuthType=0 for service account authentication StringBuilder urlBuilder = new StringBuilder(); @@ -145,6 +166,23 @@ public class BigQueryCatalog extends JdbcCatalog { urlBuilder.append("OAuthPvtKeyPath=").append(keyFilePath).append(";"); + // Add proxy configuration if provided + if (StringUtils.isNotBlank(proxyHost)) { + urlBuilder.append("ProxyHost=").append(proxyHost).append(";"); + urlBuilder.append("ProxyPort=").append(proxyPort).append(";"); + + // Add proxy authentication if provided + if (StringUtils.isNotBlank(proxyUsername)) { + urlBuilder.append("ProxyUid=").append(proxyUsername).append(";"); + + if (StringUtils.isNotBlank(proxyPassword)) { + urlBuilder.append("ProxyPwd=").append(proxyPassword).append(";"); + } + } + + LOG.info("Added proxy configuration: {}:{}", proxyHost, proxyPort); + } + String completeJdbcUrl = urlBuilder.toString(); // Create new config map with complete jdbc-url @@ -155,7 +193,9 @@ public class BigQueryCatalog extends JdbcCatalog { LOG.info("Built BigQuery JDBC URL for project: {}", projectId); LOG.debug( "Complete JDBC URL: {}", - completeJdbcUrl.replaceAll("OAuthPvtKeyPath=[^;]+", "OAuthPvtKeyPath=[REDACTED]")); + completeJdbcUrl + .replaceAll("OAuthPvtKeyPath=[^;]+", "OAuthPvtKeyPath=[REDACTED]") + .replaceAll("ProxyPwd=[^;]+", "ProxyPwd=[REDACTED]")); return newConfig; } diff --git a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalogPropertiesMetadata.java b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalogPropertiesMetadata.java index 8d356159be..065f229c7e 100644 --- a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalogPropertiesMetadata.java +++ b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryCatalogPropertiesMetadata.java @@ -33,6 +33,11 @@ public class BigQueryCatalogPropertiesMetadata extends JdbcCatalogPropertiesMeta // BigQuery catalog property keys public static final String PROJECT_ID = "project-id"; + // Proxy configuration properties + public static final String PROXY_HOST = "proxy-host"; + public static final String PROXY_PORT = "proxy-port"; + public static final String PROXY_USERNAME = "proxy-username"; + public static final String PROXY_PASSWORD = "proxy-password"; private static final Map<String, PropertyEntry<?>> BIGQUERY_CATALOG_PROPERTY_ENTRIES = ImmutableMap.<String, PropertyEntry<?>>builder() @@ -40,6 +45,38 @@ public class BigQueryCatalogPropertiesMetadata extends JdbcCatalogPropertiesMeta PROJECT_ID, PropertyEntry.stringRequiredPropertyEntry( PROJECT_ID, "Google Cloud Project ID", false /* immutable */, false /* hidden */)) + .put( + PROXY_HOST, + PropertyEntry.stringOptionalPropertyEntry( + PROXY_HOST, + "Proxy server hostname or IP address", + false /* immutable */, + null, + false /* hidden */)) + .put( + PROXY_PORT, + PropertyEntry.integerOptionalPropertyEntry( + PROXY_PORT, + "Proxy server port number", + false /* immutable */, + null, + false /* hidden */)) + .put( + PROXY_USERNAME, + PropertyEntry.stringOptionalPropertyEntry( + PROXY_USERNAME, + "Proxy authentication username", + false /* immutable */, + null, + false /* hidden */)) + .put( + PROXY_PASSWORD, + PropertyEntry.stringOptionalPropertyEntry( + PROXY_PASSWORD, + "Proxy authentication password", + false /* immutable */, + null, + true /* hidden */)) .build(); @Override diff --git a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryClientPool.java b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryClientPool.java index 05683bfb5e..e69456ba4f 100644 --- a/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryClientPool.java +++ b/catalogs/catalog-jdbc-bigquery/src/main/java/org/apache/gravitino/catalog/bigquery/BigQueryClientPool.java @@ -20,6 +20,7 @@ package org.apache.gravitino.catalog.bigquery; import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport; import com.google.api.client.http.HttpTransport; +import com.google.api.client.http.javanet.NetHttpTransport; import com.google.api.client.json.JsonFactory; import com.google.api.client.json.gson.GsonFactory; import com.google.api.services.bigquery.Bigquery; @@ -28,6 +29,10 @@ import com.google.auth.oauth2.GoogleCredentials; import com.google.auth.oauth2.ServiceAccountCredentials; import java.io.FileInputStream; import java.io.IOException; +import java.net.Authenticator; +import java.net.InetSocketAddress; +import java.net.PasswordAuthentication; +import java.net.Proxy; import java.security.GeneralSecurityException; import java.util.Map; import org.apache.commons.lang3.StringUtils; @@ -39,8 +44,9 @@ import org.slf4j.LoggerFactory; * BigQuery client pool for managing BigQuery API clients. * * <p>This class provides a centralized way to create and manage BigQuery API clients with proper - * authentication using service account credentials. It uses the Google API Client Library for Java - * (google-api-services-bigquery) which is already included in the Simba JDBC driver dependencies. + * authentication using service account credentials and optional proxy configuration. It uses the + * Google API Client Library for Java (google-api-services-bigquery) which is already included in + * the Simba JDBC driver dependencies. */ public class BigQueryClientPool { @@ -50,16 +56,28 @@ public class BigQueryClientPool { private final String projectId; private final String keyFilePath; + private final String proxyHost; + private final Integer proxyPort; + private final String proxyUsername; + private final String proxyPassword; private volatile Bigquery bigQueryClient; /** * Creates a new BigQuery client pool. * - * @param config catalog configuration containing project-id and jdbc-password (key file path) + * @param config catalog configuration containing project-id, jdbc-password (key file path), and + * optional proxy configuration */ public BigQueryClientPool(Map<String, String> config) { this.projectId = config.get(BigQueryCatalogPropertiesMetadata.PROJECT_ID); this.keyFilePath = config.get(JdbcConfig.PASSWORD.getKey()); // Key file path + this.proxyHost = config.get(BigQueryCatalogPropertiesMetadata.PROXY_HOST); + + String proxyPortStr = config.get(BigQueryCatalogPropertiesMetadata.PROXY_PORT); + this.proxyPort = StringUtils.isNotBlank(proxyPortStr) ? Integer.parseInt(proxyPortStr) : null; + + this.proxyUsername = config.get(BigQueryCatalogPropertiesMetadata.PROXY_USERNAME); + this.proxyPassword = config.get(BigQueryCatalogPropertiesMetadata.PROXY_PASSWORD); if (StringUtils.isBlank(projectId)) { throw new IllegalArgumentException("project-id is required for BigQuery catalog"); @@ -70,6 +88,12 @@ public class BigQueryClientPool { } LOG.info("Initialized BigQuery client pool for project: {}", projectId); + // Log proxy configuration + LOG.info( + "Proxy: {}:{}, Auth: {}", + StringUtils.isNotBlank(proxyHost) ? proxyHost : "none", + proxyPort != null ? proxyPort : "none", + StringUtils.isNotBlank(proxyUsername) ? "yes" : "no"); } /** @@ -93,7 +117,8 @@ public class BigQueryClientPool { } /** - * Creates a new BigQuery API client with service account authentication. + * Creates a new BigQuery API client with service account authentication and optional proxy + * support. * * @return BigQuery API client * @throws RuntimeException if client creation fails @@ -108,8 +133,8 @@ public class BigQueryClientPool { credentials = ServiceAccountCredentials.fromStream(serviceAccountStream); } - // Create HTTP transport - HttpTransport httpTransport = GoogleNetHttpTransport.newTrustedTransport(); + // Create HTTP transport with optional proxy support + HttpTransport httpTransport = createHttpTransport(); // Build BigQuery client with credentials Bigquery client = @@ -142,6 +167,52 @@ public class BigQueryClientPool { } } + /** + * Creates HTTP transport with optional proxy configuration. + * + * @return HTTP transport + * @throws GeneralSecurityException if transport creation fails + * @throws IOException if transport creation fails + */ + private HttpTransport createHttpTransport() throws GeneralSecurityException, IOException { + if (StringUtils.isBlank(proxyHost) || proxyPort == null) { + LOG.debug("Creating HTTP transport without proxy"); + return GoogleNetHttpTransport.newTrustedTransport(); + } + + LOG.debug("Creating HTTP transport with proxy: {}:{}", proxyHost, proxyPort); + + // Create proxy configuration + Proxy proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyHost, proxyPort)); + + // Set up proxy authentication if credentials are provided + if (StringUtils.isNotBlank(proxyUsername) && StringUtils.isNotBlank(proxyPassword)) { + LOG.debug("Setting up proxy authentication for user: {}", proxyUsername); + + // Create a custom authenticator for this proxy + Authenticator proxyAuthenticator = + new Authenticator() { + @Override + protected PasswordAuthentication getPasswordAuthentication() { + // Only provide credentials for our specific proxy + if (getRequestorType() == RequestorType.PROXY + && proxyHost.equals(getRequestingHost()) + && proxyPort.equals(getRequestingPort())) { + return new PasswordAuthentication(proxyUsername, proxyPassword.toCharArray()); + } + return null; + } + }; + + // Set the authenticator for this thread + // Note: This is a global setting, but we only respond to our specific proxy + Authenticator.setDefault(proxyAuthenticator); + } + + // Create NetHttpTransport with proxy + return new NetHttpTransport.Builder().setProxy(proxy).build(); + } + /** * Closes the BigQuery client and releases resources. * diff --git a/catalogs/catalog-jdbc-bigquery/src/main/resources/jdbc-bigquery.conf b/catalogs/catalog-jdbc-bigquery/src/main/resources/jdbc-bigquery.conf index 0557c12d62..08b6013b0f 100644 --- a/catalogs/catalog-jdbc-bigquery/src/main/resources/jdbc-bigquery.conf +++ b/catalogs/catalog-jdbc-bigquery/src/main/resources/jdbc-bigquery.conf @@ -26,4 +26,11 @@ # Note: jdbc-user and jdbc-password are not typically used with BigQuery # jdbc-user = <your-service-account>@<your-project-id>.iam.gserviceaccount.com # Add /path/to/key.json -# jdbc-password = <your-key-json-path> \ No newline at end of file +# jdbc-password = <your-key-json-path> +# proxy-host = Proxy server hostname or IP address (optional) +# proxy-port = Proxy server port number (optional) +# proxy-username = Proxy authentication username (optional) +# proxy-password = Proxy authentication password (optional) +# When using a proxy that requires credentials,the following JVM arguments must be used: +# -Djdk.http.auth.tunneling.disabledSchemes= +# -Djdk.http.auth.proxying.disabledSchemes= \ No newline at end of file diff --git a/catalogs/catalog-jdbc-bigquery/src/test/java/org/apache/gravitino/catalog/bigquery/TestBigQueryCatalogProxy.java b/catalogs/catalog-jdbc-bigquery/src/test/java/org/apache/gravitino/catalog/bigquery/TestBigQueryCatalogProxy.java new file mode 100644 index 0000000000..4fc3e84a6b --- /dev/null +++ b/catalogs/catalog-jdbc-bigquery/src/test/java/org/apache/gravitino/catalog/bigquery/TestBigQueryCatalogProxy.java @@ -0,0 +1,246 @@ +/* + * 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.gravitino.catalog.bigquery; + +import static org.apache.gravitino.catalog.bigquery.BigQueryCatalogPropertiesMetadata.PROJECT_ID; +import static org.apache.gravitino.catalog.bigquery.BigQueryCatalogPropertiesMetadata.PROXY_HOST; +import static org.apache.gravitino.catalog.bigquery.BigQueryCatalogPropertiesMetadata.PROXY_PASSWORD; +import static org.apache.gravitino.catalog.bigquery.BigQueryCatalogPropertiesMetadata.PROXY_PORT; +import static org.apache.gravitino.catalog.bigquery.BigQueryCatalogPropertiesMetadata.PROXY_USERNAME; +import static org.apache.gravitino.catalog.jdbc.config.JdbcConfig.PASSWORD; +import static org.apache.gravitino.catalog.jdbc.config.JdbcConfig.USERNAME; + +import com.google.common.collect.Maps; +import java.util.Map; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Unit tests for BigQuery catalog proxy configuration. */ +public class TestBigQueryCatalogProxy { + + @Test + void testProxyConfigurationInJdbcUrl() { + // Test the proxy configuration by creating a catalog and checking the processed config + // We'll use reflection to access the buildJdbcUrl method or test through withCatalogConf + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + config.put(USERNAME.getKey(), "[email protected]"); + + // Proxy configuration + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "3128"); + config.put(PROXY_USERNAME, "proxyuser"); + config.put(PROXY_PASSWORD, "proxypass"); + + // Create catalog and process configuration + BigQueryCatalog catalog = new BigQueryCatalog(); + catalog.withCatalogConf(config); + + // Since we can't access properties() without entity, we'll verify the configuration + // was processed correctly by checking that no exceptions were thrown and the catalog + // was initialized successfully. The actual JDBC URL construction is tested in integration + // tests. + + // Verify the catalog was created successfully (no exceptions thrown) + Assertions.assertNotNull(catalog); + + // Test that proxy validation works correctly + Assertions.assertDoesNotThrow( + () -> { + BigQueryCatalog testCatalog = new BigQueryCatalog(); + testCatalog.withCatalogConf(config); + }); + } + + @Test + void testProxyConfigurationWithoutAuthentication() { + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Proxy configuration without authentication + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "8080"); + + // Create catalog and process configuration + BigQueryCatalog catalog = new BigQueryCatalog(); + + // Verify the catalog processes the configuration without errors + Assertions.assertDoesNotThrow( + () -> { + catalog.withCatalogConf(config); + }); + + // Verify the catalog was created successfully + Assertions.assertNotNull(catalog); + } + + @Test + void testNoProxyConfiguration() { + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties only + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Create catalog and process configuration + BigQueryCatalog catalog = new BigQueryCatalog(); + + // Verify the catalog processes the configuration without errors + Assertions.assertDoesNotThrow( + () -> { + catalog.withCatalogConf(config); + }); + + // Verify the catalog was created successfully + Assertions.assertNotNull(catalog); + } + + @Test + void testInvalidProxyConfiguration() { + BigQueryCatalog catalog = new BigQueryCatalog(); + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Invalid proxy configuration - host without port + config.put(PROXY_HOST, "proxy.example.com"); + // Missing PROXY_PORT + + // Should throw exception + Assertions.assertThrows( + IllegalArgumentException.class, + () -> catalog.withCatalogConf(config), + "proxy-port is required when proxy-host is specified"); + } + + @Test + void testInvalidProxyPort() { + BigQueryCatalog catalog = new BigQueryCatalog(); + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Invalid proxy port + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "invalid-port"); + + // Should throw exception + Assertions.assertThrows( + IllegalArgumentException.class, + () -> catalog.withCatalogConf(config), + "proxy-port must be a valid integer"); + } + + @Test + void testProxyPortOutOfRange() { + BigQueryCatalog catalog = new BigQueryCatalog(); + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Proxy port out of range + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "70000"); + + // Should throw exception + Assertions.assertThrows( + IllegalArgumentException.class, + () -> catalog.withCatalogConf(config), + "proxy-port must be between 1 and 65535"); + } + + @Test + void testBigQueryClientPoolWithAuthenticatedProxy() { + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Proxy configuration with authentication + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "3128"); + config.put(PROXY_USERNAME, "proxyuser"); + config.put(PROXY_PASSWORD, "proxypass"); + + // Create client pool with authenticated proxy configuration + BigQueryClientPool clientPool = new BigQueryClientPool(config); + + // Verify proxy configuration is stored + Assertions.assertEquals("test-project", clientPool.getProjectId()); + + // Note: We can't easily test the actual HTTP transport proxy authentication + // without making real network calls, but we can verify the client pool + // initializes correctly with proxy authentication parameters + clientPool.close(); + } + + @Test + void testBigQueryClientPoolWithProxy() { + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Proxy configuration + config.put(PROXY_HOST, "proxy.example.com"); + config.put(PROXY_PORT, "3128"); + + // Create client pool with proxy configuration + BigQueryClientPool clientPool = new BigQueryClientPool(config); + + // Verify proxy configuration is stored + Assertions.assertEquals("test-project", clientPool.getProjectId()); + + // Note: We can't easily test the actual HTTP transport proxy configuration + // without making real network calls, but we can verify the client pool + // initializes correctly with proxy parameters + clientPool.close(); + } + + @Test + void testBigQueryClientPoolWithoutProxy() { + Map<String, String> config = Maps.newHashMap(); + + // Basic required properties only + config.put(PROJECT_ID, "test-project"); + config.put(PASSWORD.getKey(), "/path/to/key.json"); + + // Create client pool without proxy configuration + BigQueryClientPool clientPool = new BigQueryClientPool(config); + + // Verify basic configuration + Assertions.assertEquals("test-project", clientPool.getProjectId()); + + clientPool.close(); + } +}
