From f7bc8a8653ea0b5c3236c602413b944eb50a0655 Mon Sep 17 00:00:00 2001
From: aaronchung-bitquill
 <118320132+aaronchung-bitquill@users.noreply.github.com>
Date: Fri, 9 Feb 2024 00:57:49 +0000
Subject: [PATCH] feat: HostSpec to not default lastUpdateTime. Handle null
 lastUpdateTime in RdsHostListProvider (#877)

---
 .../software/amazon/jdbc/HostSpecBuilder.java |  7 ---
 .../hostlistprovider/RdsHostListProvider.java |  2 +-
 .../RdsHostListProviderTest.java              | 61 +++++++++++++++++++
 3 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java
index 156989b8a..f9908bff5 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java
@@ -89,7 +89,6 @@ public HostSpecBuilder lastUpdateTime(Timestamp lastUpdateTime) {
 
   public HostSpec build() {
     checkHostIsSet();
-    setDefaultLastUpdateTime();
     return new HostSpec(this.host, this.port, this.hostId, this.role, this.availability,
         this.weight, this.lastUpdateTime, this.hostAvailabilityStrategy);
   }
@@ -99,10 +98,4 @@ private void checkHostIsSet() {
       throw new IllegalArgumentException("host parameter must be set.");
     }
   }
-
-  private void setDefaultLastUpdateTime() {
-    if (this.lastUpdateTime == null) {
-      this.lastUpdateTime = Timestamp.from(Instant.now());
-    }
-  }
 }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java
index cf58d064f..6f0e6900b 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java
@@ -404,7 +404,7 @@ private List<HostSpec> processQueryResults(final ResultSet resultSet) throws SQL
     } else {
       // Take the latest updated writer node as the current writer. All others will be ignored.
       List<HostSpec> sortedWriters = writers.stream()
-          .sorted(Comparator.comparing(HostSpec::getLastUpdateTime).reversed())
+          .sorted(Comparator.comparing(HostSpec::getLastUpdateTime, Comparator.nullsLast(Comparator.reverseOrder())))
           .collect(Collectors.toList());
       hosts.add(sortedWriters.get(0));
     }
diff --git a/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java b/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java
index aaef62ac7..ff17abcb7 100644
--- a/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java
+++ b/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java
@@ -578,4 +578,65 @@ void testGetTopology_InvalidLastUpdatedTimestamp() throws SQLException {
         expectedLastUpdatedTimeStampRounded,
         result.hosts.get(0).getLastUpdateTime().toString().substring(0, 16));
   }
+
+  @Test
+  void testGetTpology_returnsLatestWriter() throws SQLException {
+    rdsHostListProvider = Mockito.spy(
+        getRdsHostListProvider(mockHostListProviderService, "jdbc:someprotocol://url"));
+    rdsHostListProvider.isInitialized = true;
+
+    HostSpec expectedWriterHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
+        .host("expectedWriterHost")
+        .role(HostRole.WRITER)
+        .lastUpdateTime(Timestamp.valueOf("3000-01-01 00:00:00"))
+        .build();
+
+    HostSpec unexpectedWriterHost0 = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
+        .host("unexpectedWriterHost0")
+        .role(HostRole.WRITER)
+        .lastUpdateTime(Timestamp.valueOf("1000-01-01 00:00:00"))
+        .build();
+
+    HostSpec unexpectedWriterHost1 = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
+        .host("unexpectedWriterHost1")
+        .role(HostRole.WRITER)
+        .lastUpdateTime(Timestamp.valueOf("2000-01-01 00:00:00"))
+        .build();
+
+    HostSpec unexpectedWriterHostWithNullLastUpdateTime0 = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
+        .host("unexpectedWriterHostWithNullLastUpdateTime0")
+        .role(HostRole.WRITER)
+        .lastUpdateTime(null)
+        .build();
+
+    HostSpec unexpectedWriterHostWithNullLastUpdateTime1 = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
+        .host("unexpectedWriterHostWithNullLastUpdateTime1")
+        .role(HostRole.WRITER)
+        .lastUpdateTime(null)
+        .build();
+
+    when(mockResultSet.next()).thenReturn(true, true, true, true, true, false);
+
+    when(mockResultSet.getString(1)).thenReturn(
+        unexpectedWriterHostWithNullLastUpdateTime0.getHost(),
+        unexpectedWriterHost0.getHost(),
+        expectedWriterHost.getHost(),
+        unexpectedWriterHost1.getHost(),
+        unexpectedWriterHostWithNullLastUpdateTime1.getHost());
+    when(mockResultSet.getBoolean(2)).thenReturn(true, true, true, true, true);
+    when(mockResultSet.getFloat(3)).thenReturn((float) 0, (float) 0, (float) 0, (float) 0, (float) 0);
+    when(mockResultSet.getFloat(4)).thenReturn((float) 0, (float) 0, (float) 0, (float) 0, (float) 0);
+    when(mockResultSet.getTimestamp(5)).thenReturn(
+        unexpectedWriterHostWithNullLastUpdateTime0.getLastUpdateTime(),
+        unexpectedWriterHost0.getLastUpdateTime(),
+        expectedWriterHost.getLastUpdateTime(),
+        unexpectedWriterHost1.getLastUpdateTime(),
+        unexpectedWriterHostWithNullLastUpdateTime1.getLastUpdateTime()
+    );
+
+    final FetchTopologyResult result = rdsHostListProvider.getTopology(mockConnection, true);
+    verify(rdsHostListProvider, atMostOnce()).queryForTopology(mockConnection);
+
+    assertEquals(expectedWriterHost.getHost(), result.hosts.get(0).getHost());
+  }
 }