From f88ab6b8528d9384cf23839babd820578d67baf9 Mon Sep 17 00:00:00 2001
From: Forest Vey <forestmvey@gmail.com>
Date: Wed, 23 Aug 2023 21:52:50 -0700
Subject: [PATCH] Add support for region use with endpoint discovery with
 datasource connection. (#47)

Signed-off-by: forestmvey <forestv@bitquilltech.com>
---
 .../integrationtest/ConnectionTest.java       | 40 +++++++++++++++++++
 .../timestream/jdbc/TimestreamDataSource.java |  6 ++-
 .../jdbc/TimestreamDataSourceTest.java        | 18 +++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

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