(1) Support reading config file 'smart-pgjdbc.config' from user's hom…

…e directory; (2) Support reporting query time to QueryRewriter through the HTTP request;
ISG-ICS · Mar 15, 2023 · bbae070 · bbae070
1 parent a2e9b54
commit bbae070
Show file tree

Hide file tree

Showing 7 changed files with 128 additions and 5 deletions.
diff --git a/pgjdbc/src/main/java/org/postgresql/core/BaseConnection.java b/pgjdbc/src/main/java/org/postgresql/core/BaseConnection.java
@@ -235,4 +235,11 @@ CachedQuery createQuery(String sql, boolean escapeProcessing, boolean isParamete
    * @return true if should be included and passed on to other exceptions
    */
   boolean getLogServerErrorDetail();
+
+  /**
+   * Return the map between query sql and query guid.
+   *
+   * @return querySqlqueryGuid
+   */
+  java.util.HashMap<String, String> getQuerySqlqueryGuid();
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/jdbc/PgConnection.java b/pgjdbc/src/main/java/org/postgresql/jdbc/PgConnection.java
@@ -170,22 +170,37 @@ private enum ReadOnlyBehavior {
   private final @Nullable String xmlFactoryFactoryClass;
   private @Nullable PGXmlFactoryFactory xmlFactoryFactory;
 
+  // ++ SmartJDBC ++ //
+  /* Map between query sql and query guid */
+  private final HashMap<String, String> querySqlqueryGuid;
+
+  public HashMap<String, String> getQuerySqlqueryGuid() {
+    return querySqlqueryGuid;
+  }
+  // -- SmartJDBC -- //
+
   final CachedQuery borrowQuery(String sql) throws SQLException {
     // ++ SmartJDBC ++ //
     StringBuilder logText = new StringBuilder();
     logText.append("==================================================\n");
     logText.append("  PgConnection -> borrowQuery -> [Original SQL]\n");
     logText.append("--------------------------------------------------\n");
     logText.append(sql);
+    logText.append("\n");
     LogUtil.log(logText.toString());
 
     if (sql.toLowerCase().contains("select")) {
-      sql = QueryRewritter.rewriteQuery(sql);
+      String guid = java.util.UUID.randomUUID().toString();
+      sql = QueryRewritter.rewriteQuery(guid, sql).trim();
+      querySqlqueryGuid.put(sql, guid);
       logText = new StringBuilder();
       logText.append("==================================================\n");
       logText.append("  PgConnection -> borrowQuery -> [Rewritten SQL]\n");
       logText.append("--------------------------------------------------\n");
+      logText.append(guid);
+      logText.append("\n");
       logText.append(sql);
+      logText.append("\n");
       LogUtil.log(logText.toString());
     }
     // -- SmartJDBC -- //
@@ -342,6 +357,10 @@ public PgConnection(HostSpec[] hostSpecs,
     replicationConnection = PGProperty.REPLICATION.get(info) != null;
 
     xmlFactoryFactoryClass = PGProperty.XML_FACTORY_FACTORY.get(info);
+
+    // ++ SmartJDBC ++ //
+    querySqlqueryGuid = new HashMap<String, String>();
+    // -- SmartJDBC -- //
   }
 
   private static ReadOnlyBehavior getReadOnlyBehavior(String property) {

diff --git a/pgjdbc/src/main/java/org/postgresql/jdbc/PgStatement.java b/pgjdbc/src/main/java/org/postgresql/jdbc/PgStatement.java
@@ -20,6 +20,8 @@
 import org.postgresql.core.ResultHandlerBase;
 import org.postgresql.core.SqlCommand;
 import org.postgresql.core.Tuple;
+import org.postgresql.smart.LogUtil;
+import org.postgresql.smart.QueryRewritter;
 import org.postgresql.util.GT;
 import org.postgresql.util.PSQLException;
 import org.postgresql.util.PSQLState;
@@ -485,13 +487,48 @@ private void executeInternal(CachedQuery cachedQuery,
     synchronized (this) {
       result = null;
     }
+    // ++ SmartJDBC ++ //
+    long start = System.nanoTime();
+    // -- SmartJDBC -- //
     try {
       startTimer();
       connection.getQueryExecutor().execute(queryToExecute, queryParameters, handler, maxrows,
           fetchSize, flags, adaptiveFetch);
     } finally {
       killTimerTask();
     }
+    // ++ SmartJDBC ++ //
+    long end = System.nanoTime();
+    long queryTimeMs = (end - start) / 1000000;
+    String sql = queryToExecute.toString().trim();
+    String guid = connection.getQuerySqlqueryGuid().get(sql);
+    if (guid != null) {
+      boolean success = QueryRewritter.reportQuery(guid, queryTimeMs);
+      StringBuilder logText = new StringBuilder();
+      logText.append("==================================================\n");
+      logText.append("  PgStatement -> executeInternal\n");
+      logText.append("--------------------------------------------------\n");
+      logText.append(guid);
+      logText.append("\n");
+      logText.append(sql);
+      logText.append("\n");
+      logText.append("queryTimeMs: " + String.valueOf(queryTimeMs));
+      logText.append("\n");
+      logText.append("reportSuccess: " + String.valueOf(success));
+      logText.append("\n");
+      LogUtil.log(logText.toString());
+    } else {
+      StringBuilder logText = new StringBuilder();
+      logText.append("==================================================\n");
+      logText.append("  PgStatement -> executeInternal\n");
+      logText.append("--------------------------------------------------\n");
+      logText.append(sql);
+      logText.append("\n");
+      logText.append("No guid found for this query!!");
+      logText.append("\n");
+      LogUtil.log(logText.toString());
+    }
+    // -- SmartJDBC -- //
     synchronized (this) {
       checkClosed();
 

diff --git a/pgjdbc/src/main/java/org/postgresql/smart/ConfigUtil.java b/pgjdbc/src/main/java/org/postgresql/smart/ConfigUtil.java
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2021, PostgreSQL Global Development Group
+ * See the LICENSE file in the project root for more information.
+ */
+
+package org.postgresql.smart;
+
+import org.postgresql.util.OSUtil;
+
+import java.io.File;
+import java.io.FileReader;
+import java.io.IOException;
+import java.util.Properties;
+
+public class ConfigUtil {
+
+  private static final String CONFIG_PATH = OSUtil.getUserConfigRootDirectory() + "/smart-pgjdbc.config";
+  private static final String SMARTJDBC_SERVER_URL_DEFAULT = "http://localhost:8000";
+  private static final String APP_GUID = "App007";
+
+  public static Properties readConfigProperties() throws IOException {
+    File configFile = new File(CONFIG_PATH);
+    FileReader cofingFileReader = new FileReader(configFile);
+    Properties properties = new Properties();
+    properties.load(cofingFileReader);
+    return properties;
+  }
+
+  public static String getSmartJDBCServerUrl() {
+    try {
+      Properties properties = readConfigProperties();
+      return properties.getProperty("smartjdbc.server.url");
+    } catch (IOException e) {
+      return SMARTJDBC_SERVER_URL_DEFAULT;
+    }
+  }
+
+  public static String getAppGuid() {
+    try {
+      Properties properties = readConfigProperties();
+      return properties.getProperty("app.guid");
+    } catch (IOException e) {
+      return APP_GUID;
+    }
+  }
+}
diff --git a/pgjdbc/src/main/java/org/postgresql/smart/QueryRewritter.java b/pgjdbc/src/main/java/org/postgresql/smart/QueryRewritter.java
@@ -9,11 +9,20 @@
 
 public class QueryRewritter {
 
-  public static String rewriteQuery(String sql) {
+  public static String rewriteQuery(String guid, String sql) {
     try {
-      return SmartJDBCClient.request("{\"cmd\": \"rewrite\", \"db\": \"postgresql\", \"query\": \"" + sql.replace("\"", "\\\"") + "\"}");
+      return SmartJDBCClient.request("{\"cmd\": \"rewrite\", \"db\": \"postgresql\", \"query\": \"" + sql.replace("\"", "\\\"") + "\", \"guid\": \"" + guid + "\", \"appguid\": \"" + ConfigUtil.getAppGuid() + "\"}");
     } catch (IOException e) {
       return sql;
     }
   }
+
+  public static boolean reportQuery(String guid, long queryTimeMs) {
+    try {
+      SmartJDBCClient.request("{\"cmd\": \"report\", \"db\": \"postgresql\", \"queryTimeMs\": " + String.valueOf(queryTimeMs) + ", \"guid\": \"" + guid + "\", \"appguid\": \"" + ConfigUtil.getAppGuid() + "\"}");
+      return true;
+    } catch (IOException e) {
+      return false;
+    }
+  }
 }
diff --git a/pgjdbc/src/main/java/org/postgresql/smart/SmartJDBCClient.java b/pgjdbc/src/main/java/org/postgresql/smart/SmartJDBCClient.java
@@ -13,8 +13,7 @@
 
 public class SmartJDBCClient {
 
-  // TODO - Get the parameter from a configuration file.
-  public static final String SMARTJDBC_SERVER_URL = "http://localhost:8000";
+  public static final String SMARTJDBC_SERVER_URL = ConfigUtil.getSmartJDBCServerUrl();
 
   public static String request(String request) throws IOException {
     URL url = new URL(SMARTJDBC_SERVER_URL);

diff --git a/pgjdbc/src/test/java/org/postgresql/jdbc/AbstractArraysTest.java b/pgjdbc/src/test/java/org/postgresql/jdbc/AbstractArraysTest.java
@@ -42,6 +42,7 @@
 import java.sql.Savepoint;
 import java.sql.Statement;
 import java.sql.Struct;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
 import java.util.TimerTask;
@@ -1014,5 +1015,10 @@ public boolean getAdaptiveFetch() {
     public boolean getLogServerErrorDetail() {
       return false;
     }
+
+    @Override
+    public HashMap<String, String> getQuerySqlqueryGuid() {
+      throw new UnsupportedOperationException();
+    }
   }
 }