Address some review comments. Fix perf testing program.

Author: QuChen88

benchmarks/src/jmh/java/software/amazon/jdbc/benchmarks/PgCacheBenchmarks.java

@@ -1,6 +1,7 @@
 package software.amazon.jdbc.benchmarks;
 
 import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
 import java.sql.*;
 import java.util.*;
 import java.util.concurrent.TimeUnit;
@@ -89,36 +90,54 @@ public void warmUpDataSet() throws SQLException {
     }
   }
 
-  private void validateResultSet(ResultSet rs) throws SQLException {
+  private void validateResultSet(ResultSet rs, Blackhole b) throws SQLException {
     while (rs.next()) {
-      assert rs.getInt(1) >= 0;
-      assert rs.getInt(2) >= 0;
-      assert rs.getString(3) != null;
-      assert rs.getString(4) != null;
-      assert rs.getDouble(5) >= 0.0;
-      assert rs.getDate(6) != null;
-      assert rs.getTime(7) != null;
-      assert rs.getTime(8) != null;
-      assert rs.getTimestamp(9) != null;
-      assert rs.getTimestamp(10) != null;
-      assert !rs.wasNull();
+      b.consume(rs.getInt(1));
+      b.consume(rs.getInt(2));
+      b.consume(rs.getString(3));
+      b.consume(rs.getString(4));
+      b.consume(rs.getDouble(5));
+      b.consume(rs.getDate(6));
+      b.consume(rs.getTime(7));
+      b.consume(rs.getTime(8));
+      b.consume(rs.getTimestamp(9));
+      b.consume(rs.getTimestamp(10));
+      b.consume(rs.wasNull());
     }
   }
 
   @Benchmark
-  public void runBenchmarkPrimaryKeyLookup() throws SQLException {
+  public void runBenchmarkPrimaryKeyLookupNoCaching(Blackhole b) throws SQLException {
+    try (Statement stmt = connection.createStatement();
+         ResultSet rs = stmt.executeQuery("SELECT * FROM test where id = " + counter)) {
+      validateResultSet(rs, b);
+    }
+    counter++;
+  }
+
+  @Benchmark
+  public void runBenchmarkNonIndexedLookupNoCaching(Blackhole b) throws SQLException {
+    try (Statement stmt = connection.createStatement();
+         ResultSet rs = stmt.executeQuery("SELECT * FROM test where int_col = " + counter*10)) {
+      validateResultSet(rs, b);
+    }
+    counter++;
+  }
+
+  @Benchmark
+  public void runBenchmarkPrimaryKeyLookupWithCaching(Blackhole b) throws SQLException {
     try (Statement stmt = connection.createStatement();
          ResultSet rs = stmt.executeQuery("/*+ CACHE_PARAM(ttl=172800s) */ SELECT * FROM test where id = " + counter)) {
-         validateResultSet(rs);
+         validateResultSet(rs, b);
     }
     counter++;
   }
 
   @Benchmark
-  public void runBenchmarkNonIndexedLookup() throws SQLException {
+  public void runBenchmarkNonIndexedLookupWithCaching(Blackhole b) throws SQLException {
     try (Statement stmt = connection.createStatement();
          ResultSet rs = stmt.executeQuery("/*+ CACHE_PARAM(ttl=172800s) */ SELECT * FROM test where int_col = " + counter*10)) {
-      validateResultSet(rs);
+      validateResultSet(rs, b);
     }
     counter++;
   }

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

@@ -116,7 +116,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).withLibraryName("aws-jdbc-lettuce").build();
+          .withSsl(useSSL).withVerifyPeer(false).withLibraryName("aws-sql-jdbc-lettuce").build();
 
       RedisClient client = RedisClient.create(resources, redisUriCluster);
       GenericObjectPool<StatefulRedisConnection<byte[], byte[]>> pool = new GenericObjectPool<>(

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

@@ -26,9 +26,11 @@
 import java.util.Properties;
 import java.util.Set;
 import java.util.logging.Logger;
+import software.amazon.jdbc.AwsWrapperProperty;
 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.states.SessionStateService;
 import software.amazon.jdbc.util.Messages;
@@ -41,7 +43,6 @@
 
 public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
   private static final Logger LOGGER = Logger.getLogger(DataRemoteCachePlugin.class.getName());
-  private static final int MAX_CACHEABLE_QUERY_SIZE = 16000;
   private static final String QUERY_HINT_START_PATTERN = "/*+";
   private static final String QUERY_HINT_END_PATTERN = "*/";
   private static final String CACHE_PARAM_PATTERN = "CACHE_PARAM(";
@@ -55,6 +56,7 @@ public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
           JdbcMethod.CALLABLESTATEMENT_EXECUTE.methodName,
           JdbcMethod.CALLABLESTATEMENT_EXECUTEQUERY.methodName)));
 
+  private int maxCacheableQuerySize;
   private PluginService pluginService;
   private TelemetryFactory telemetryFactory;
   private TelemetryCounter cacheHitCounter;
@@ -63,6 +65,17 @@ public class DataRemoteCachePlugin extends AbstractConnectionPlugin {
   private TelemetryCounter malformedHintCounter;
   private TelemetryCounter cacheBypassCounter;
   private CacheConnection cacheConnection;
+  private String dbUserName;
+
+  private static final AwsWrapperProperty CACHE_MAX_QUERY_SIZE =
+      new AwsWrapperProperty(
+          "cacheMaxQuerySize",
+          "16384",
+          "The max query size for remote caching");
+
+  static {
+    PropertyDefinition.registerPluginProperties(DataRemoteCachePlugin.class);
+  }
 
   public DataRemoteCachePlugin(final PluginService pluginService, final Properties properties) {
     try {
@@ -78,7 +91,9 @@ public DataRemoteCachePlugin(final PluginService pluginService, final Properties
     this.totalQueryCounter = telemetryFactory.createCounter("JdbcCacheTotalQueryCount");
     this.malformedHintCounter = telemetryFactory.createCounter("JdbcCacheMalformedQueryHint");
     this.cacheBypassCounter = telemetryFactory.createCounter("JdbcCacheBypassCount");
+    this.maxCacheableQuerySize = CACHE_MAX_QUERY_SIZE.getInteger(properties);
     this.cacheConnection = new CacheConnection(properties);
+    this.dbUserName = PropertyDefinition.USER.getString(properties);
   }
 
   // Used for unit testing purposes only
@@ -93,25 +108,36 @@ public Set<String> getSubscribedMethods() {
 
   private String getCacheQueryKey(String query) {
     // Check some basic session states. The important ones for caching include (but not limited to):
-    //   schema name, username which can affect the query result from the DB in addition to the query string
+    //  schema name, username which can affect the query result from the DB in addition to the query string
     try {
       Connection currentConn = pluginService.getCurrentConnection();
       DatabaseMetaData metadata = currentConn.getMetaData();
       // Fetch and record the schema name if the session state doesn't currently have it
       SessionStateService sessionStateService = pluginService.getSessionStateService();
+      String catalog = sessionStateService.getCatalog().orElse(null);
       String schema = sessionStateService.getSchema().orElse(null);
-      if (schema == null) {
+      if (catalog == null && schema == null) {
         // Fetch the current schema name and store it in sessionStateService
+        catalog = currentConn.getCatalog();
         schema = currentConn.getSchema();
-        sessionStateService.setSchema(schema);
+        if (catalog != null) sessionStateService.setCatalog(catalog);
+        if (schema != null) sessionStateService.setSchema(schema);
       }
 
+      if (dbUserName == null) {
+        // For MySQL, metadata username is actually <UserName>@<ip>. We just need the part before '@'.
+        dbUserName = metadata.getUserName();
+        int nameIndexEnd = dbUserName.indexOf('@');
+        if (nameIndexEnd > 0) {
+          dbUserName = dbUserName.substring(0, nameIndexEnd);
+        }
+      }
       LOGGER.finest("DB driver protocol " + pluginService.getDriverProtocol()
           + ", database product: " + metadata.getDatabaseProductName() + " " + metadata.getDatabaseProductVersion()
-          + ", schema: " + schema + ", user: " + metadata.getUserName()
+          + ", catalog: " + catalog + ", schema: " + schema + ", user: " + dbUserName
           + ", driver: " + metadata.getDriverName() + " " + metadata.getDriverVersion());
       // The cache key contains the schema name, user name, and the query string
-      String[] words = {schema, metadata.getUserName(), query};
+      String[] words = {catalog, schema, dbUserName, query};
       return String.join("_", words);
     } catch (SQLException e) {
       LOGGER.warning("Error getting session state: " + e.getMessage());
@@ -239,11 +265,11 @@ public <T, E extends Exception> T execute(
     int endOfQueryHint = 0;
     Integer configuredQueryTtl = null;
     // Queries longer than 16KB is not cacheable
-    if ((sql.length() < MAX_CACHEABLE_QUERY_SIZE) && sql.startsWith(QUERY_HINT_START_PATTERN)) {
+    if ((sql.length() < maxCacheableQuerySize) && sql.startsWith(QUERY_HINT_START_PATTERN)) {
       endOfQueryHint = sql.indexOf(QUERY_HINT_END_PATTERN);
       if (endOfQueryHint > 0) {
-        configuredQueryTtl = getTtlForQuery(sql.substring(2, endOfQueryHint).trim());
-        mainQuery = sql.substring(endOfQueryHint + 2).trim();
+        configuredQueryTtl = getTtlForQuery(sql.substring(QUERY_HINT_START_PATTERN.length(), endOfQueryHint).trim());
+        mainQuery = sql.substring(endOfQueryHint + QUERY_HINT_END_PATTERN.length()).trim();
       }
     }