feat: Add support for Explain feature (#1852)

* Added the code for handling the explain * Update ConnectionStatementExecutorImpl.java * Added tests * Update PG_ClientSideStatements.json * resolved the comments * Update google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json Co-authored-by: Knut Olav Løite <koloite@gmail.com> * Update ConnectionStatementExecutorTest.java * Update PG_ClientSideStatements.json * Update PG_ClientSideStatements.json * Formatted the files for ci lint * Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java Co-authored-by: Knut Olav Løite <koloite@gmail.com> * Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java Co-authored-by: Knut Olav Løite <koloite@gmail.com> * resolved some comments * formatted the code * added some code * added support for "explain (format ) foo" kind of statements * resolved some comments * Update ConnectionStatementExecutorImpl.java * fixed a small bug * Added the code for formatting query plan for export * Update ConnectionStatementExecutorImpl.java * Update ConnectionStatementExecutorImpl.java * Update ConnectionStatementExecutorImpl.java * Update ConnectionStatementExecutorImpl.java * Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java Co-authored-by: Knut Olav Løite <koloite@gmail.com> * Changed assertThat to assertEquals * removed unnecessary lines * Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java Co-authored-by: Knut Olav Løite <koloite@gmail.com> * Changed assertThat to assertEquals * Update ConnectionStatementExecutorImpl.java * Update ConnectionStatementExecutorImpl.java * Added tests * format * Update PartitionedDmlTransaction.java * generated sql script * resolved comments * resolved comments * Update PG_ClientSideStatements.json * Update PG_ClientSideStatements.json * Update PG_ClientSideStatements.json * Create ITExplainTest.java * added Integration tests * reformatted * changed region * Revert "changed region" This reverts commit 10f06e8. * Update ITExplainTest.java * Update ITExplainTest.java * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Knut Olav Løite <koloite@gmail.com> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · May 19, 2022 · 01f460e · 01f460e
1 parent 6ab8ed2
commit 01f460e
Show file tree

Hide file tree

Showing 13 changed files with 11,036 additions and 9,815 deletions.
diff --git a/...src/main/java/com/google/cloud/spanner/connection/ClientSideStatementExplainExecutor.java b/...src/main/java/com/google/cloud/spanner/connection/ClientSideStatementExplainExecutor.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.connection;
+
+import com.google.cloud.spanner.ErrorCode;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.connection.ClientSideStatementImpl.CompileException;
+import com.google.cloud.spanner.connection.ClientSideStatementValueConverters.ExplainCommandConverter;
+import com.google.common.collect.ImmutableSet;
+import java.lang.reflect.Method;
+import java.util.Set;
+import java.util.regex.Matcher;
+
+/** Specific executor for the EXPLAIN statement for PostgreSQL. */
+class ClientSideStatementExplainExecutor implements ClientSideStatementExecutor {
+  private final ClientSideStatementImpl statement;
+  private final Method method;
+  private final ExplainCommandConverter converter;
+  public static final Set<String> EXPLAIN_OPTIONS =
+      ImmutableSet.of(
+          "verbose", "costs", "settings", "buffers", "wal", "timing", "summary", "format");
+
+  ClientSideStatementExplainExecutor(ClientSideStatementImpl statement) throws CompileException {
+    try {
+      this.statement = statement;
+      this.converter = new ExplainCommandConverter();
+      this.method =
+          ConnectionStatementExecutor.class.getDeclaredMethod(
+              statement.getMethodName(), converter.getParameterClass());
+    } catch (Exception e) {
+      throw new CompileException(e, statement);
+    }
+  }
+
+  @Override
+  public StatementResult execute(ConnectionStatementExecutor connection, String sql)
+      throws Exception {
+    return (StatementResult) method.invoke(connection, getParameterValue(sql));
+  }
+
+  String getParameterValue(String sql) {
+    Matcher matcher = statement.getPattern().matcher(sql);
+    if (matcher.find() && matcher.groupCount() >= 1) {
+      String value = matcher.group(0);
+      if (value != null) {
+        String res = converter.convert(value.trim());
+        if (res != null) {
+          return res;
+        }
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT, String.format("Invalid argument for EXPLAIN: %s", value));
+      }
+    }
+    return null;
+  }
+}
diff --git a/...src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java b/...src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java
@@ -342,4 +342,23 @@ public Priority convert(String value) {
       return values.get("PRIORITY_" + value);
     }
   }
+
+  static class ExplainCommandConverter implements ClientSideStatementValueConverter<String> {
+    @Override
+    public Class<String> getParameterClass() {
+      return String.class;
+    }
+
+    @Override
+    public String convert(String value) {
+      /* The first word in the string should be "explain"
+       *  So, if the size of the string <= 7 (number of letters in the word "explain"), its an invalid statement
+       *  If the size is greater than 7, we'll consider everything after explain as the query.
+       */
+      if (value.length() <= 7) {
+        return null;
+      }
+      return value.substring(7).trim();
+    }
+  }
 }
diff --git a/...panner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java b/...panner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java
@@ -111,4 +111,6 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
   StatementResult statementShowRPCPriority();
 
   StatementResult statementShowTransactionIsolationLevel();
+
+  StatementResult statementExplain(String sql);
 }
diff --git a/...er/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java b/...er/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java
@@ -59,9 +59,13 @@
 import com.google.cloud.spanner.CommitResponse;
 import com.google.cloud.spanner.CommitStats;
 import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.ResultSets;
+import com.google.cloud.spanner.SpannerExceptionFactory;
+import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Struct;
 import com.google.cloud.spanner.TimestampBound;
 import com.google.cloud.spanner.Type;
@@ -71,8 +75,11 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.Duration;
+import com.google.spanner.v1.PlanNode;
+import com.google.spanner.v1.QueryPlan;
 import com.google.spanner.v1.RequestOptions;
 import com.google.spanner.v1.RequestOptions.Priority;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -83,6 +90,7 @@
  * calls are then forwarded into a {@link Connection}.
  */
 class ConnectionStatementExecutorImpl implements ConnectionStatementExecutor {
+
   static final class StatementTimeoutGetter implements DurationValueGetter {
     private final Connection connection;
 
@@ -442,4 +450,196 @@ public StatementResult statementShowRPCPriority() {
   public StatementResult statementShowTransactionIsolationLevel() {
     return resultSet("transaction_isolation", "serializable", SHOW_TRANSACTION_ISOLATION_LEVEL);
   }
+
+  private String processQueryPlan(PlanNode planNode) {
+    StringBuilder planNodeDescription = new StringBuilder(" : { ");
+    com.google.protobuf.Struct metadata = planNode.getMetadata();
+
+    for (String key : metadata.getFieldsMap().keySet()) {
+      planNodeDescription
+          .append(key)
+          .append(" : ")
+          .append(metadata.getFieldsMap().get(key).getStringValue())
+          .append(" , ");
+    }
+    String substring = planNodeDescription.substring(0, planNodeDescription.length() - 3);
+    planNodeDescription.setLength(0);
+    planNodeDescription.append(substring).append(" }");
+
+    return planNodeDescription.toString();
+  }
+
+  private String processExecutionStats(PlanNode planNode) {
+    StringBuilder executionStats = new StringBuilder("");
+    for (String key : planNode.getExecutionStats().getFieldsMap().keySet()) {
+      executionStats.append(key).append(" : { ");
+      com.google.protobuf.Struct value =
+          planNode.getExecutionStats().getFieldsMap().get(key).getStructValue();
+      for (String newKey : value.getFieldsMap().keySet()) {
+        String newValue = value.getFieldsMap().get(newKey).getStringValue();
+        executionStats.append(newKey).append(" : ").append(newValue).append(" , ");
+      }
+      String substring = executionStats.substring(0, executionStats.length() - 3);
+      executionStats.setLength(0);
+      executionStats.append(substring).append(" } , ");
+    }
+    String substring = executionStats.substring(0, executionStats.length() - 3);
+    executionStats.setLength(0);
+    executionStats.append(substring);
+    return executionStats.toString();
+  }
+
+  private StatementResult getStatementResultFromQueryPlan(QueryPlan queryPlan, boolean isAnalyze) {
+    ArrayList<Struct> list = new ArrayList<>();
+    for (PlanNode planNode : queryPlan.getPlanNodesList()) {
+      String planNodeDescription = planNode.getDisplayName();
+      String executionStats = "";
+
+      if (!planNode.getMetadata().toString().equalsIgnoreCase("")) {
+        planNodeDescription += processQueryPlan(planNode);
+      }
+
+      if (!planNode.getShortRepresentation().toString().equalsIgnoreCase("")) {
+        planNodeDescription += " : " + planNode.getShortRepresentation().getDescription();
+      }
+
+      if (isAnalyze && !planNode.getExecutionStats().toString().equals("")) {
+        executionStats = processExecutionStats(planNode);
+      }
+      Struct.Builder builder = Struct.newBuilder().set("QUERY PLAN").to(planNodeDescription);
+
+      if (isAnalyze) {
+        builder.set("EXECUTION STATS").to(executionStats);
+      }
+      list.add(builder.build());
+    }
+
+    ResultSet resultSet;
+    if (isAnalyze) {
+      resultSet =
+          ResultSets.forRows(
+              Type.struct(
+                  StructField.of("QUERY PLAN", Type.string()),
+                  StructField.of("EXECUTION STATS", Type.string())),
+              list);
+    } else {
+      resultSet =
+          ResultSets.forRows(Type.struct(StructField.of("QUERY PLAN", Type.string())), list);
+    }
+    return StatementResultImpl.of(resultSet);
+  }
+
+  private StatementResult executeStatement(String sql, QueryAnalyzeMode queryAnalyzeMode) {
+    Statement statement = Statement.newBuilder(sql).build();
+    try (ResultSet resultSet = getConnection().analyzeQuery(statement, queryAnalyzeMode)) {
+      while (resultSet.next()) {
+        // ResultSet.next() should return false in order to access the ResultSet.Stats
+      }
+
+      if (resultSet.getStats() != null && resultSet.getStats().getQueryPlan() != null) {
+        return getStatementResultFromQueryPlan(
+            resultSet.getStats().getQueryPlan(), queryAnalyzeMode.equals(QueryAnalyzeMode.PROFILE));
+      }
+    }
+    throw SpannerExceptionFactory.newSpannerException(
+        ErrorCode.FAILED_PRECONDITION, String.format("Couldn't fetch stats for %s", sql));
+  }
+
+  // This method removes parenthesis from the sql string assuming it is ending with the closing
+  // parenthesis
+  private String removeParenthesisAndTrim(String sql) {
+    sql = sql.trim();
+    if (sql.charAt(0) == '(') {
+      sql = sql.substring(1, sql.length() - 1);
+    }
+    return sql.trim();
+  }
+
+  /*
+   * This method executes the given SQL string in either PLAN or PROFILE mode and returns
+   * the query plan and execution stats (if necessary).
+   *
+   * The only additional option that is supported is ANALYZE. The method will throw a SpannerException
+   * if it is invoked with a statement that includes any other options.
+   *
+   * If the SQL string has ANALYZE option, it will be executed in PROFILE mode and will return a resultset
+   * with two String columns namely QUERY PLAN and EXECUTION STATS.
+   *
+   * If the sql string doesn't have any option, it will be executed in PLAN mode and will return a resultset
+   * with one string column namely QUERY PLAN.
+   */
+  @Override
+  public StatementResult statementExplain(String sql) {
+    if (sql == null) {
+      throw SpannerExceptionFactory.newSpannerException(
+          ErrorCode.INVALID_ARGUMENT, String.format("Invalid String with Explain"));
+    }
+
+    if (sql.charAt(0) == '(') {
+      int index = sql.indexOf(')');
+      if (index == -1) {
+        throw SpannerExceptionFactory.newSpannerException(
+            ErrorCode.INVALID_ARGUMENT,
+            String.format("Missing closing parenthesis in the query: %s", sql));
+      }
+      String options[] = sql.substring(1, index).split("\\s*,\\s*");
+      boolean isAnalyze = false, startAfterIndex = false;
+      for (String option : options) {
+        String optionExpression[] = option.trim().split("\\s+");
+        if (optionExpression.length >= 3) {
+          isAnalyze = false;
+          break;
+        } else if (ClientSideStatementExplainExecutor.EXPLAIN_OPTIONS.contains(
+            optionExpression[0].toLowerCase())) {
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.UNIMPLEMENTED,
+              String.format("%s is not implemented yet", optionExpression[0]));
+        } else if (optionExpression[0].equalsIgnoreCase("analyse")
+            || optionExpression[0].equalsIgnoreCase("analyze")) {
+          isAnalyze = true;
+        } else {
+          isAnalyze = false;
+          break;
+        }
+
+        if (optionExpression.length == 2) {
+          if (optionExpression[1].equalsIgnoreCase("false")
+              || optionExpression[1].equalsIgnoreCase("0")
+              || optionExpression[1].equalsIgnoreCase("off")) {
+            isAnalyze = false;
+            startAfterIndex = true;
+          } else if (!(optionExpression[1].equalsIgnoreCase("true")
+              || optionExpression[1].equalsIgnoreCase("1")
+              || optionExpression[1].equalsIgnoreCase("on"))) {
+            isAnalyze = false;
+            break;
+          }
+        }
+      }
+      if (isAnalyze) {
+        String newSql = removeParenthesisAndTrim(sql.substring(index + 1));
+        return executeStatement(newSql, QueryAnalyzeMode.PROFILE);
+      } else if (startAfterIndex) {
+        String newSql = removeParenthesisAndTrim(sql.substring(index + 1));
+        return executeStatement(newSql, QueryAnalyzeMode.PLAN);
+      } else {
+        return executeStatement(removeParenthesisAndTrim(sql), QueryAnalyzeMode.PLAN);
+      }
+    } else {
+      String[] arr = sql.split("\\s+", 2);
+      if (arr.length >= 2) {
+        String option = arr[0].toLowerCase();
+        String statementToBeExplained = arr[1];
+
+        if (ClientSideStatementExplainExecutor.EXPLAIN_OPTIONS.contains(option)) {
+          throw SpannerExceptionFactory.newSpannerException(
+              ErrorCode.UNIMPLEMENTED, String.format("%s is not implemented yet", option));
+        } else if (option.equals("analyze") || option.equals("analyse")) {
+          return executeStatement(
+              removeParenthesisAndTrim(statementToBeExplained), QueryAnalyzeMode.PROFILE);
+        }
+      }
+      return executeStatement(sql, QueryAnalyzeMode.PLAN);
+    }
+  }
 }
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java
@@ -83,7 +83,8 @@ enum ClientSideStatementType {
     ABORT_BATCH,
     SET_RPC_PRIORITY,
     SHOW_RPC_PRIORITY,
-    SHOW_TRANSACTION_ISOLATION_LEVEL
+    SHOW_TRANSACTION_ISOLATION_LEVEL,
+    EXPLAIN
   }
 
   /**

diff --git a/...anner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json b/...anner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json
@@ -149,6 +149,15 @@
       "method": "statementShowTransactionIsolationLevel",
       "exampleStatements": ["show transaction isolation level","show variable transaction isolation level"]
     },
+    {
+      "name": "EXPLAIN <sql>",
+      "executorName": "ClientSideStatementExplainExecutor",
+      "resultType": "RESULT_SET",
+      "statementType": "EXPLAIN",
+      "regex": "(?is)\\A\\s*explain(\\s+|\\()(.*)\\z",
+      "method": "statementExplain",
+      "exampleStatements": []
+    },
     {
       "name": "{START | BEGIN} [TRANSACTION | WORK] [{ (READ ONLY|READ WRITE) | (ISOLATION LEVEL (DEFAULT|SERIALIZABLE)) }]",
       "executorName": "ClientSideStatementPgBeginExecutor",