Caching - minor fixes after rebase including doing null check for telemetry.

Author: QuChen88

wrapper/src/main/java/software/amazon/jdbc/plugin/cache/CacheConnection.java

@@ -14,10 +14,12 @@
 import java.security.NoSuchAlgorithmException;
 import java.time.Duration;
 import java.util.Properties;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.logging.Logger;
 import io.lettuce.core.support.ConnectionPoolSupport;
 import org.apache.commons.pool2.impl.GenericObjectPool;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+import software.amazon.jdbc.PropertyDefinition;
 import software.amazon.jdbc.util.StringUtils;
 
 // Abstraction layer on top of a connection to a remote cache server
@@ -36,8 +38,8 @@ public class CacheConnection {
   private static final int DEFAULT_POOL_MIN_IDLE = 0;
   private static final long DEFAULT_MAX_BORROW_WAIT_MS = 50;
 
-  private static final Object READ_LOCK = new Object();
-  private static final Object WRITE_LOCK = new Object();
+  private static final ReentrantLock READ_LOCK = new ReentrantLock();
+  private static final ReentrantLock WRITE_LOCK = new ReentrantLock();
 
   protected static final AwsWrapperProperty CACHE_RW_ENDPOINT_ADDR =
       new AwsWrapperProperty(
@@ -59,6 +61,10 @@ public class CacheConnection {
 
   private final boolean useSSL;
 
+  static {
+    PropertyDefinition.registerPluginProperties(CacheConnection.class);
+  }
+
   public CacheConnection(final Properties properties) {
     this.cacheRwServerAddr = CACHE_RW_ENDPOINT_ADDR.getString(properties);
     this.cacheRoServerAddr = CACHE_RO_ENDPOINT_ADDR.getString(properties);
@@ -84,11 +90,14 @@ private void initializeCacheConnectionIfNeeded(boolean isRead) {
     GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> cacheConnectionPool =
         isRead ? readConnectionPool : writeConnectionPool;
     if (cacheConnectionPool == null) {
-      Object lock = isRead ? READ_LOCK : WRITE_LOCK;
-      synchronized (lock) {
+      ReentrantLock connectionPoolLock = isRead ? READ_LOCK : WRITE_LOCK;
+      connectionPoolLock.lock();
+      try {
         if ((isRead && readConnectionPool == null) || (!isRead && writeConnectionPool == null)) {
           createConnectionPool(isRead);
         }
+      } finally {
+        connectionPoolLock.unlock();
       }
     }
   }
@@ -105,7 +114,7 @@ private void createConnectionPool(boolean isRead) {
       String[] hostnameAndPort = serverAddr.split(":");
       RedisURI redisUriCluster = RedisURI.Builder.redis(hostnameAndPort[0])
           .withPort(Integer.parseInt(hostnameAndPort[1]))
-          .withSsl(useSSL).withVerifyPeer(false).build();
+          .withSsl(useSSL).withVerifyPeer(false).withLibraryName("aws-jdbc-lettuce").build();
 
       RedisClient client = RedisClient.create(resources, redisUriCluster);
       GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> pool =
@@ -219,6 +228,7 @@ public void writeToCache(String key, byte[] value, int expiry) {
       asyncCommands.set(keyHash, value, SetArgs.Builder.ex(expiry))
           .whenComplete((result, exception) -> handleCompletedCacheWrite(finalConn, exception));
     } catch (Exception e) {
+      // Failed to trigger the async write to the cache, return the cache connection to the pool as broken
       LOGGER.warning("Failed to write to cache: " + e.getMessage());
       if (conn != null && writeConnectionPool != null) {
         try {

wrapper/src/main/java/software/amazon/jdbc/plugin/cache/CachedResultSet.java

@@ -40,7 +40,6 @@
 import java.util.Calendar;
 import java.util.GregorianCalendar;
 
-@SuppressWarnings({"RedundantThrows", "checkstyle:OverloadMethodsDeclarationOrder"})
 public class CachedResultSet implements ResultSet {
 
   public static class CachedRow {
@@ -65,8 +64,10 @@ public Object get(final String columnName) {
 
   protected ArrayList<CachedRow> rows;
   protected int currentRow;
+  protected boolean wasNullFlag;
   protected ResultSetMetaData metadata;
   protected static ObjectMapper mapper = new ObjectMapper();
+  protected static boolean mapperInitialized = false;
   protected static final TimeZone defaultTimeZone = TimeZone.getDefault();
   private static final Calendar calendarWithUserTz = new GregorianCalendar();
 
@@ -83,34 +84,45 @@ public CachedResultSet(final ResultSet resultSet) throws SQLException {
       rows.add(row);
     }
     currentRow = -1;
+    initializeObjectMapper();
   }
 
   public CachedResultSet(final List<Map<String, Object>> resultList) {
     rows = new ArrayList<>();
-    CachedResultSetMetaData.Field[] fields = new CachedResultSetMetaData.Field[resultList.get(0).size()];
-    boolean fieldsInitialized = false;
-    for (Map<String, Object> rowMap : resultList) {
-      final CachedRow row = new CachedRow();
-      int i = 0;
-      for (Map.Entry<String, Object> entry : rowMap.entrySet()) {
-        String columnName = entry.getKey();
-        if (!fieldsInitialized) {
-          fields[i] = new CachedResultSetMetaData.Field(columnName, columnName);
+    int numFields = resultList.isEmpty() ? 0 : resultList.get(0).size();
+    CachedResultSetMetaData.Field[] fields = new CachedResultSetMetaData.Field[numFields];
+    if (!resultList.isEmpty()) {
+      boolean fieldsInitialized = false;
+      for (Map<String, Object> rowMap : resultList) {
+        final CachedRow row = new CachedRow();
+        int i = 0;
+        for (Map.Entry<String, Object> entry : rowMap.entrySet()) {
+          String columnName = entry.getKey();
+          if (!fieldsInitialized) {
+            fields[i] = new CachedResultSetMetaData.Field(columnName, columnName);
+          }
+          row.put(++i, columnName, entry.getValue());
         }
-        row.put(++i, columnName, entry.getValue());
+        rows.add(row);
+        fieldsInitialized = true;
       }
-      rows.add(row);
-      fieldsInitialized = true;
     }
     currentRow = -1;
     metadata = new CachedResultSetMetaData(fields);
+    initializeObjectMapper();
   }
 
-  public String serializeIntoJsonString() throws SQLException {
+  // Helper method to initialize the object mapper for serialization of objects
+  private void initializeObjectMapper() {
+    if (mapperInitialized) return;
+    // For serialization of Date/LocalDateTime etc, set up the time module,
+    // and use standard string format (i.e. ISO)
     mapper.registerModule(new JavaTimeModule());
-    // Serialize Date/LocalDateTime etc. into standard string format (i.e. ISO)
     mapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
+    mapperInitialized = true;
+  }
 
+  public String serializeIntoJsonString() throws SQLException {
     List<Map<String, Object>> resultList = new ArrayList<>();
     ResultSetMetaData metaData = this.getMetaData();
     int columns = metaData.getColumnCount();
@@ -156,7 +168,10 @@ public void close() throws SQLException {
 
   @Override
   public boolean wasNull() throws SQLException {
-    throw new UnsupportedOperationException();
+    if (isClosed()) {
+      throw new SQLException("This result set is closed");
+    }
+    return this.wasNullFlag;
   }
 
   // TODO: implement all the getXXX APIs.
@@ -476,12 +491,12 @@ public InputStream getBinaryStream(final String columnLabel) throws SQLException
 
   @Override
   public SQLWarning getWarnings() throws SQLException {
-    throw new UnsupportedOperationException();
+    return null;
   }
 
   @Override
   public void clearWarnings() throws SQLException {
-    throw new UnsupportedOperationException();
+    // no-op
   }
 
   @Override
@@ -494,28 +509,34 @@ public ResultSetMetaData getMetaData() throws SQLException {
     return metadata;
   }
 
-  @Override
-  public Object getObject(final int columnIndex) throws SQLException {
+  private void checkCurrentRow() throws SQLException {
     if (this.currentRow < 0 || this.currentRow >= this.rows.size()) {
-      return null; // out of boundaries
+      throw new SQLException("The current row index " + this.currentRow + " is out of range.");
     }
+  }
+
+  @Override
+  public Object getObject(final int columnIndex) throws SQLException {
+    checkCurrentRow();
     final CachedRow row = this.rows.get(this.currentRow);
     if (!row.columnByIndex.containsKey(columnIndex)) {
-      return null; // column index out of boundaries
+      throw new SQLException("The column index: " + columnIndex + " is out of range, number of columns: " + row.columnByIndex.size());
     }
-    return row.columnByIndex.get(columnIndex);
+    Object obj = row.columnByIndex.get(columnIndex);
+    this.wasNullFlag = (obj == null);
+    return obj;
   }
 
   @Override
   public Object getObject(final String columnLabel) throws SQLException {
-    if (this.currentRow < 0 || this.currentRow >= this.rows.size()) {
-      return null; // out of boundaries
-    }
+    checkCurrentRow();
     final CachedRow row = this.rows.get(this.currentRow);
     if (!row.columnByName.containsKey(columnLabel)) {
-      return null; // column name not found
+      throw new SQLException("The column label: " + columnLabel + " is not found");
     }
-    return row.columnByName.get(columnLabel);
+    Object obj = row.columnByName.get(columnLabel);
+    this.wasNullFlag = (obj == null);
+    return obj;
   }
 
   @Override
@@ -559,12 +580,12 @@ public boolean isAfterLast() throws SQLException {
 
   @Override
   public boolean isFirst() throws SQLException {
-    return this.currentRow == 0 && this.rows.size() > 0;
+    return this.currentRow == 0 && !this.rows.isEmpty();
   }
 
   @Override
   public boolean isLast() throws SQLException {
-    return this.currentRow == (this.rows.size() - 1) && this.rows.size() > 0;
+    return this.currentRow == (this.rows.size() - 1) && !this.rows.isEmpty();
   }
 
   @Override
@@ -596,24 +617,49 @@ public int getRow() throws SQLException {
 
   @Override
   public boolean absolute(final int row) throws SQLException {
-    if (row > 0) {
-      this.currentRow = row - 1;
+    if (row == 0) {
+      this.beforeFirst();
+      return false;
     } else {
-      this.currentRow = this.rows.size() + row;
+      int rowsSize = this.rows.size();
+      if (row < 0) {
+        if (row < -rowsSize) {
+          this.beforeFirst();
+          return false;
+        }
+        this.currentRow = rowsSize + row;
+      } else { // row > 0
+        if (row > rowsSize) {
+          this.afterLast();
+          return false;
+        }
+        this.currentRow = row - 1;
+      }
     }
-    return this.currentRow >= 0 && this.currentRow < this.rows.size();
+    return true;
   }
 
   @Override
   public boolean relative(final int rows) throws SQLException {
     this.currentRow += rows;
-    return this.currentRow >= 0 && this.currentRow < this.rows.size();
+    if (this.currentRow < 0) {
+      this.beforeFirst();
+      return false;
+    } else if (this.currentRow >= this.rows.size()) {
+      this.afterLast();
+      return false;
+    }
+    return true;
   }
 
   @Override
   public boolean previous() throws SQLException {
+    if (this.currentRow < 1) {
+      this.beforeFirst();
+      return false;
+    }
     this.currentRow--;
-    return this.currentRow >= 0 && this.currentRow < this.rows.size();
+    return true;
   }
 
   @Override
@@ -1054,7 +1100,7 @@ public int getHoldability() throws SQLException {
 
   @Override
   public boolean isClosed() throws SQLException {
-    return false;
+    return this.rows == null;
   }
 
   @Override

wrapper/src/main/java/software/amazon/jdbc/plugin/cache/DataRemoteCachePlugin.java

@@ -28,20 +28,25 @@
 import java.util.Set;
 import java.util.logging.Logger;
 import software.amazon.jdbc.JdbcCallable;
+import software.amazon.jdbc.JdbcMethod;
 import software.amazon.jdbc.PluginService;
 import software.amazon.jdbc.PropertyDefinition;
 import software.amazon.jdbc.plugin.AbstractConnectionPlugin;
 import software.amazon.jdbc.util.Messages;
 import software.amazon.jdbc.util.StringUtils;
+import software.amazon.jdbc.util.WrapperUtils;
 import software.amazon.jdbc.util.telemetry.TelemetryCounter;
 import software.amazon.jdbc.util.telemetry.TelemetryFactory;
 
 public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
   private static final Logger LOGGER = Logger.getLogger(DataRemoteCachePlugin.class.getName());
   private static final Set<String> subscribedMethods = Collections.unmodifiableSet(new HashSet<>(
-      Arrays.asList("Statement.executeQuery", "Statement.execute",
-          "PreparedStatement.execute", "PreparedStatement.executeQuery",
-          "CallableStatement.execute", "CallableStatement.executeQuery")));
+      Arrays.asList(JdbcMethod.STATEMENT_EXECUTEQUERY.methodName,
+          JdbcMethod.STATEMENT_EXECUTE.methodName,
+          JdbcMethod.PREPAREDSTATEMENT_EXECUTE.methodName,
+          JdbcMethod.PREPAREDSTATEMENT_EXECUTEQUERY.methodName,
+          JdbcMethod.CALLABLESTATEMENT_EXECUTE.methodName,
+          JdbcMethod.CALLABLESTATEMENT_EXECUTEQUERY.methodName)));
 
   static {
     PropertyDefinition.registerPluginProperties(DataRemoteCachePlugin.class);
@@ -173,8 +178,6 @@ public <T, E extends Exception> T execute(
       final JdbcCallable<T, E> jdbcMethodFunc,
       final Object[] jdbcMethodArgs)
       throws E {
-    totalCallsCounter.inc();
-
     ResultSet result;
     boolean needToCache = false;
     final String sql = getQuery(jdbcMethodArgs);
@@ -196,42 +199,51 @@ public <T, E extends Exception> T execute(
     // Query result can be served from the cache if it has a configured TTL value, and it is
     // not executed in a transaction as a transaction typically need to return consistent results.
     if (!isInTransaction && (configuredQueryTtl != null)) {
+      incrCounter(totalCallsCounter);
       result = fetchResultSetFromCache(mainQuery);
       if (result == null) {
         // Cache miss. Need to fetch result from the database
         needToCache = true;
-        missCounter.inc();
+        incrCounter(missCounter);
         LOGGER.finest("Got a cache miss for SQL: " + sql);
       } else {
         LOGGER.finest("Got a cache hit for SQL: " + sql);
         // Cache hit. Return the cached result
-        hitCounter.inc();
+        incrCounter(hitCounter);
         try {
           result.beforeFirst();
         } catch (final SQLException ex) {
-          if (exceptionClass.isAssignableFrom(ex.getClass())) {
-            throw exceptionClass.cast(ex);
-          }
-          throw new RuntimeException(ex);
+          throw WrapperUtils.wrapExceptionIfNeeded(exceptionClass, ex);
         }
         return resultClass.cast(result);
       }
     }
 
     result = (ResultSet) jdbcMethodFunc.call();
 
+    // We need to cache the query result if we got a cache miss for the query result,
+    // or the query is cacheable and executed inside a transaction.
+    if (isInTransaction && (configuredQueryTtl != null)) {
+      needToCache = true;
+    }
     if (needToCache) {
       try {
         result = cacheResultSet(mainQuery, result, configuredQueryTtl);
       } catch (final SQLException ex) {
-        // ignore exception
-        LOGGER.warning("Encountered SQLException when caching results: " + ex.getMessage());
+        // Log and re-throw exception
+        LOGGER.warning("Encountered SQLException when caching query results: " + ex.getMessage());
+        throw WrapperUtils.wrapExceptionIfNeeded(exceptionClass, ex);
       }
     }
 
     return resultClass.cast(result);
   }
 
+  private void incrCounter(TelemetryCounter counter) {
+    if (counter == null) return;
+    counter.inc();
+  }
+
   protected String getQuery(final Object[] jdbcMethodArgs) {
     // Get query from method argument
     if (jdbcMethodArgs != null && jdbcMethodArgs.length > 0 && jdbcMethodArgs[0] != null) {