diff --git a/integrationtest/src/main/java/software/amazon/timestream/integrationtest/ConnectionTest.java b/integrationtest/src/main/java/software/amazon/timestream/integrationtest/ConnectionTest.java index c7bfe30..edec7e6 100644 --- a/integrationtest/src/main/java/software/amazon/timestream/integrationtest/ConnectionTest.java +++ b/integrationtest/src/main/java/software/amazon/timestream/integrationtest/ConnectionTest.java @@ -82,6 +82,46 @@ void testConnectionWithRegion() throws SQLException { } } + @Test + void testConnectionWithRegionOverride() throws SQLException { + final Properties properties = new Properties(); + final String urlWithRegion = Constants.URL + "://Region=us-east-1"; + + // URL should override property value for region. + properties.put("region", "us-west-2"); + + try (Connection connection = DriverManager.getConnection(urlWithRegion, properties)) { + validateConnection( + connection, + EXPECTED_DEFAULT_CONNECTION_URL + ); + } + + final TimestreamDataSource dataSource = new TimestreamDataSource(); + dataSource.setRegion("us-east-1"); + try (Connection connection = dataSource.getConnection()) { + validateConnection( + connection, + EXPECTED_DEFAULT_CONNECTION_URL + ); + } + } + + @Test + void testConnectionWithEmptyRegion() { + final Properties properties = new Properties(); + final String urlWithEmptyRegion = Constants.URL + "://Region="; + + Assertions.assertThrows( + SQLException.class, + () -> DriverManager.getConnection(urlWithEmptyRegion, properties)); + + final TimestreamDataSource dataSource = new TimestreamDataSource(); + dataSource.setEndpoint(urlWithEmptyRegion); + dataSource.setRegion(""); + Assertions.assertThrows(SQLException.class, dataSource::getConnection); + } + @Test void testConnectionWithSDKOptions() throws SQLException { final Properties properties = new Properties(); diff --git a/jdbc/src/main/java/software/amazon/timestream/jdbc/TimestreamDataSource.java b/jdbc/src/main/java/software/amazon/timestream/jdbc/TimestreamDataSource.java index bb28eb1..af2d32e 100644 --- a/jdbc/src/main/java/software/amazon/timestream/jdbc/TimestreamDataSource.java +++ b/jdbc/src/main/java/software/amazon/timestream/jdbc/TimestreamDataSource.java @@ -16,6 +16,7 @@ package software.amazon.timestream.jdbc; import com.amazonaws.ClientConfiguration; +import com.amazonaws.util.StringUtils; import com.google.common.annotations.VisibleForTesting; import javax.sql.ConnectionEvent; @@ -712,13 +713,16 @@ private Properties getProperties(final String accessKey, final String secretKey) throw new SQLException(error); } - if (this.region == null) { + if (StringUtils.isNullOrEmpty(this.region)) { final String error = Error.lookup(Error.MISSING_SERVICE_REGION); LOGGER.severe(error); throw new SQLException(error); } properties.put(TimestreamConnectionProperty.ENDPOINT.getConnectionProperty(), this.endpoint); + } + + if (!StringUtils.isNullOrEmpty(this.region)) { properties.put(TimestreamConnectionProperty.REGION.getConnectionProperty(), this.region); } diff --git a/jdbc/src/test/java/software/amazon/timestream/jdbc/TimestreamDataSourceTest.java b/jdbc/src/test/java/software/amazon/timestream/jdbc/TimestreamDataSourceTest.java index 0f43e10..6cb236a 100644 --- a/jdbc/src/test/java/software/amazon/timestream/jdbc/TimestreamDataSourceTest.java +++ b/jdbc/src/test/java/software/amazon/timestream/jdbc/TimestreamDataSourceTest.java @@ -104,6 +104,24 @@ void testDataSourceWithEndpoint() throws SQLException { Assertions.assertEquals(expected, constructedConnectionProperties); } + @Test + void testDataSourceWithEndpointDiscoveryAndRegion() throws SQLException { + final ArgumentCaptor propertiesArgumentCaptor = ArgumentCaptor.forClass(Properties.class); + final MockTimestreamDataSource mockTimestreamDataSource = new MockTimestreamDataSource( + mockTimestreamConnection); + + final Properties expected = new Properties(); + expected.put(TimestreamConnectionProperty.REGION.getConnectionProperty(), "east"); + + mockTimestreamDataSource.setRegion("east"); + mockTimestreamDataSource.getConnection(); + Mockito.verify(mockTimestreamConnection).setClientInfo(propertiesArgumentCaptor.capture()); + + final Properties constructedConnectionProperties = propertiesArgumentCaptor.getValue(); + Assertions.assertNotNull(constructedConnectionProperties); + Assertions.assertEquals(expected, constructedConnectionProperties); + } + @Test void testDataSourceWithEmptyEndpoint() { final MockTimestreamDataSource mockTimestreamDataSource = new MockTimestreamDataSource(