Add support for region use with endpoint discovery with datasource co…

…nnection. (#47)

Signed-off-by: forestmvey <[email protected]>

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();

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);
     }
 

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<Properties> 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(