diff --git a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java index e93938f25..96fe9b3e4 100644 --- a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java +++ b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryConnection.java @@ -130,6 +130,8 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { String sslTrustStorePassword; long maxBytesBilled; Map labels; + Integer httpConnectTimeout; + Integer httpReadTimeout; BigQueryConnection(String url) throws IOException { this.connectionUrl = url; @@ -271,11 +273,25 @@ public class BigQueryConnection extends BigQueryNoOpsConnection { BigQueryJdbcUrlUtility.SSL_TRUST_STORE_PWD_PROPERTY_NAME, null, this.connectionClassName); + this.httpConnectTimeout = + BigQueryJdbcUrlUtility.parseIntProperty( + url, + BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME, + null, + this.connectionClassName); + this.httpReadTimeout = + BigQueryJdbcUrlUtility.parseIntProperty( + url, + BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME, + null, + this.connectionClassName); this.httpTransportOptions = BigQueryJdbcProxyUtility.getHttpTransportOptions( proxyProperties, this.sslTrustStorePath, this.sslTrustStorePassword, + this.httpConnectTimeout, + this.httpReadTimeout, this.connectionClassName); this.transportChannelProvider = BigQueryJdbcProxyUtility.getTransportChannelProvider( @@ -723,6 +739,14 @@ String getSSLTrustStorePassword() { return sslTrustStorePassword; } + Integer getHttpConnectTimeout() { + return httpConnectTimeout; + } + + Integer getHttpReadTimeout() { + return httpReadTimeout; + } + @Override public boolean isValid(int timeout) throws SQLException { if (timeout < 0) { diff --git a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtility.java b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtility.java index ebc5450b6..22400d305 100644 --- a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtility.java +++ b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtility.java @@ -97,7 +97,8 @@ static Map parseProxyProperties(String URL, String callerClassNa (proxyHost == null && proxyPort != null) || (proxyHost != null && proxyPort == null); if (isMissingProxyHostOrPortWhenProxySet) { throw new IllegalArgumentException( - "Both ProxyHost and ProxyPort parameters need to be specified. No defaulting behavior occurs."); + "Both ProxyHost and ProxyPort parameters need to be specified. No defaulting behavior" + + " occurs."); } boolean isMissingProxyUidOrPwdWhenAuthSet = (proxyUid == null && proxyPwd != null) || (proxyUid != null && proxyPwd == null); @@ -117,18 +118,39 @@ static HttpTransportOptions getHttpTransportOptions( Map proxyProperties, String sslTrustStorePath, String sslTrustStorePassword, + Integer connectTimeout, + Integer readTimeout, String callerClassName) { LOG.finest("++enter++\t" + callerClassName); - if (!proxyProperties.containsKey(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME) - && sslTrustStorePath == null) { + boolean hasProxyOrSsl = + proxyProperties.containsKey(BigQueryJdbcUrlUtility.PROXY_HOST_PROPERTY_NAME) + || sslTrustStorePath != null; + boolean hasTimeoutConfig = connectTimeout != null || readTimeout != null; + + if (!hasProxyOrSsl && !hasTimeoutConfig) { return null; } - return HttpTransportOptions.newBuilder() - .setHttpTransportFactory( - getHttpTransportFactory( - proxyProperties, sslTrustStorePath, sslTrustStorePassword, callerClassName)) - .build(); + + HttpTransportOptions.Builder httpTransportOptionsBuilder = HttpTransportOptions.newBuilder(); + if (hasProxyOrSsl) { + httpTransportOptionsBuilder.setHttpTransportFactory( + getHttpTransportFactory( + proxyProperties, sslTrustStorePath, sslTrustStorePassword, callerClassName)); + } + + if (hasTimeoutConfig) { + httpTransportOptionsBuilder.setConnectTimeout( + connectTimeout != null + ? connectTimeout + : BigQueryJdbcUrlUtility.DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE); + httpTransportOptionsBuilder.setReadTimeout( + readTimeout != null + ? readTimeout + : BigQueryJdbcUrlUtility.DEFAULT_HTTP_READ_TIMEOUT_VALUE); + } + + return httpTransportOptionsBuilder.build(); } private static HttpTransportFactory getHttpTransportFactory( diff --git a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java index 3b26f7be5..487b7d038 100644 --- a/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java +++ b/google-cloud-bigquery-jdbc/src/main/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtility.java @@ -98,6 +98,10 @@ final class BigQueryJdbcUrlUtility { static final String PROXY_PORT_PROPERTY_NAME = "ProxyPort"; static final String PROXY_USER_ID_PROPERTY_NAME = "ProxyUid"; static final String PROXY_PASSWORD_PROPERTY_NAME = "ProxyPwd"; + static final String HTTP_CONNECT_TIMEOUT_PROPERTY_NAME = "HttpConnectTimeout"; + static final String HTTP_READ_TIMEOUT_PROPERTY_NAME = "HttpReadTimeout"; + static final int DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE = 60000; + static final int DEFAULT_HTTP_READ_TIMEOUT_VALUE = 60000; static final boolean DEFAULT_ENABLE_HTAPI_VALUE = false; static final boolean DEFAULT_ENABLE_SESSION_VALUE = false; static final int DEFAULT_LOG_LEVEL = 0; @@ -249,7 +253,8 @@ final class BigQueryJdbcUrlUtility { BigQueryConnectionProperty.newBuilder() .setName(OAUTH_SA_IMPERSONATION_CHAIN_PROPERTY_NAME) .setDescription( - "Comma separated list of service account emails in the impersonation chain.") + "Comma separated list of service account emails in the impersonation" + + " chain.") .build(), BigQueryConnectionProperty.newBuilder() .setName(OAUTH_SA_IMPERSONATION_SCOPES_PROPERTY_NAME) @@ -569,6 +574,18 @@ final class BigQueryJdbcUrlUtility { .setDescription( "The password for accessing the Java TrustStore that is specified using" + " the property SSLTrustStore.") + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(HTTP_CONNECT_TIMEOUT_PROPERTY_NAME) + .setDescription( + "The timeout (in milliseconds) for establishing a connection to the" + + " server.") + .setDefaultValue(String.valueOf(DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE)) + .build(), + BigQueryConnectionProperty.newBuilder() + .setName(HTTP_READ_TIMEOUT_PROPERTY_NAME) + .setDescription("The timeout (in milliseconds) when reading from the server.") + .setDefaultValue(String.valueOf(DEFAULT_HTTP_READ_TIMEOUT_VALUE)) .build()))); private BigQueryJdbcUrlUtility() {} diff --git a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtilityTest.java b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtilityTest.java index 203502cda..1d83174ea 100644 --- a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtilityTest.java +++ b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcProxyUtilityTest.java @@ -138,7 +138,8 @@ public void testGetHttpTransportOptionsWithAuthenticatedProxy() { Map proxyProperties = BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null); HttpTransportOptions result = - BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null); + BigQueryJdbcProxyUtility.getHttpTransportOptions( + proxyProperties, null, null, null, null, "TestClass"); assertNotNull(result); } @@ -154,7 +155,8 @@ public void testGetHttpTransportOptionsWithNonAuthenticatedProxy() { Map proxyProperties = BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null); HttpTransportOptions result = - BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null); + BigQueryJdbcProxyUtility.getHttpTransportOptions( + proxyProperties, null, null, null, null, "TestClass"); assertNotNull(result); } @@ -168,7 +170,8 @@ public void testGetHttpTransportOptionsWithNoProxySettingsReturnsNull() { Map proxyProperties = BigQueryJdbcProxyUtility.parseProxyProperties(connection_uri, null); HttpTransportOptions result = - BigQueryJdbcProxyUtility.getHttpTransportOptions(proxyProperties, null, null, null); + BigQueryJdbcProxyUtility.getHttpTransportOptions( + proxyProperties, null, null, null, null, "TestClass"); assertNull(result); } @@ -185,7 +188,7 @@ public void testGetHttpTransportOptions_withSslTrustStore_noPassword() throws Ex String trustStorePath = getTestResourcePath("test_truststore_nopass.jks"); HttpTransportOptions options = BigQueryJdbcProxyUtility.getHttpTransportOptions( - Collections.emptyMap(), trustStorePath, null, "TestClass"); + Collections.emptyMap(), trustStorePath, null, null, null, "TestClass"); assertNotNull(options); assertNotNull(options.getHttpTransportFactory()); } @@ -195,7 +198,12 @@ public void testGetHttpTransportOptions_withSslTrustStore_withCorrectPassword() String trustStorePath = getTestResourcePath("test_truststore_withpass.jks"); HttpTransportOptions options = BigQueryJdbcProxyUtility.getHttpTransportOptions( - Collections.emptyMap(), trustStorePath, "testpassword", "TestClass"); + Collections.emptyMap(), + trustStorePath, + "testpassword", + null, + null, + "TestClass"); assertNotNull(options); assertNotNull(options.getHttpTransportFactory()); } @@ -212,6 +220,8 @@ public void testGetHttpTransportOptions_withSslTrustStore_withIncorrectPassword( Collections.emptyMap(), trustStorePath, "wrongpassword", + null, + null, "TestClass")); assertThat(exception.getCause()).isInstanceOf(IOException.class); } @@ -224,7 +234,12 @@ public void testGetHttpTransportOptions_withInvalidSslTrustStorePath() { BigQueryJdbcRuntimeException.class, () -> BigQueryJdbcProxyUtility.getHttpTransportOptions( - Collections.emptyMap(), invalidPath, null, "TestClass")); + Collections.emptyMap(), + invalidPath, + null, + null, + null, + "TestClass")); assertThat(exception.getCause()).isInstanceOf(FileNotFoundException.class); } @@ -238,7 +253,7 @@ public void testGetHttpTransportOptions_withSslAndProxy() throws Exception { HttpTransportOptions options = BigQueryJdbcProxyUtility.getHttpTransportOptions( - proxyProperties, trustStorePath, null, "TestClass"); + proxyProperties, trustStorePath, null, null, null, "TestClass"); assertNotNull(options); assertNotNull(options.getHttpTransportFactory()); } @@ -287,7 +302,39 @@ public void testGetTransportChannelProvider_noProxyNoSsl_returnsNull() { public void testGetHttpTransportOptions_noProxyNoSsl_returnsNull() { HttpTransportOptions options = BigQueryJdbcProxyUtility.getHttpTransportOptions( - Collections.emptyMap(), null, null, "TestClass"); + Collections.emptyMap(), null, null, null, null, "TestClass"); assertNull(options); } + + @Test + public void testGetHttpTransportOptions_withTimeouts_returnsOptions() { + HttpTransportOptions options = + BigQueryJdbcProxyUtility.getHttpTransportOptions( + Collections.emptyMap(), null, null, 10000, 20000, "TestClass"); + assertNotNull(options); + assertThat(options.getConnectTimeout()).isEqualTo(10000); + assertThat(options.getReadTimeout()).isEqualTo(20000); + } + + @Test + public void testGetHttpTransportOptions_withConnectTimeoutOnly_returnsOptions() { + HttpTransportOptions options = + BigQueryJdbcProxyUtility.getHttpTransportOptions( + Collections.emptyMap(), null, null, 10000, null, "TestClass"); + assertNotNull(options); + assertThat(options.getConnectTimeout()).isEqualTo(10000); + assertThat(options.getReadTimeout()) + .isEqualTo(BigQueryJdbcUrlUtility.DEFAULT_HTTP_READ_TIMEOUT_VALUE); + } + + @Test + public void testGetHttpTransportOptions_withReadTimeoutOnly_returnsOptions() { + HttpTransportOptions options = + BigQueryJdbcProxyUtility.getHttpTransportOptions( + Collections.emptyMap(), null, null, null, 20000, "TestClass"); + assertNotNull(options); + assertThat(options.getReadTimeout()).isEqualTo(20000); + assertThat(options.getConnectTimeout()) + .isEqualTo(BigQueryJdbcUrlUtility.DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE); + } } diff --git a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java index 86f087bf3..a99829e9e 100644 --- a/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java +++ b/google-cloud-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/BigQueryJdbcUrlUtilityTest.java @@ -800,4 +800,63 @@ public void testParseLabelsEmpty() { Map labels = BigQueryJdbcUrlUtility.parseLabels(connection_uri, null); assertNull(labels); } + + @Test + public void testParseHttpConnectTimeout() { + String connection_uri = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject;" + + "HttpConnectTimeout=10000"; + + Integer timeout = + BigQueryJdbcUrlUtility.parseIntProperty( + connection_uri, + BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME, + BigQueryJdbcUrlUtility.DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE, + null); + assertEquals(Integer.valueOf(10000), timeout); + } + + @Test + public void testParseHttpConnectTimeoutDefault() { + String connection_uri = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject"; + + Integer timeout = + BigQueryJdbcUrlUtility.parseIntProperty( + connection_uri, + BigQueryJdbcUrlUtility.HTTP_CONNECT_TIMEOUT_PROPERTY_NAME, + BigQueryJdbcUrlUtility.DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE, + null); + assertEquals( + Integer.valueOf(BigQueryJdbcUrlUtility.DEFAULT_HTTP_CONNECT_TIMEOUT_VALUE), timeout); + } + + @Test + public void testParseHttpReadTimeout() { + String connection_uri = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject;" + + "HttpReadTimeout=20000"; + + Integer timeout = + BigQueryJdbcUrlUtility.parseIntProperty( + connection_uri, + BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME, + BigQueryJdbcUrlUtility.DEFAULT_HTTP_READ_TIMEOUT_VALUE, + null); + assertEquals(Integer.valueOf(20000), timeout); + } + + @Test + public void testParseHttpReadTimeoutDefault() { + String connection_uri = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;OAuthType=3;ProjectId=testProject"; + + Integer timeout = + BigQueryJdbcUrlUtility.parseIntProperty( + connection_uri, + BigQueryJdbcUrlUtility.HTTP_READ_TIMEOUT_PROPERTY_NAME, + BigQueryJdbcUrlUtility.DEFAULT_HTTP_READ_TIMEOUT_VALUE, + null); + assertEquals(Integer.valueOf(BigQueryJdbcUrlUtility.DEFAULT_HTTP_READ_TIMEOUT_VALUE), timeout); + } }