Skip to content

Commit 01f460e

Browse files
gauravsnjolavloitegcf-owl-bot[bot]
authored
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 <[email protected]> * 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 <[email protected]> * Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java Co-authored-by: Knut Olav Løite <[email protected]> * 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 <[email protected]> * 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 <[email protected]> * 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 <[email protected]> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 6ab8ed2 commit 01f460e

File tree

13 files changed

+11036
-9815
lines changed

13 files changed

+11036
-9815
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementExplainExecutor.java

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* Copyright 2022 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.spanner.connection;
18+
19+
import com.google.cloud.spanner.ErrorCode;
20+
import com.google.cloud.spanner.SpannerExceptionFactory;
21+
import com.google.cloud.spanner.connection.ClientSideStatementImpl.CompileException;
22+
import com.google.cloud.spanner.connection.ClientSideStatementValueConverters.ExplainCommandConverter;
23+
import com.google.common.collect.ImmutableSet;
24+
import java.lang.reflect.Method;
25+
import java.util.Set;
26+
import java.util.regex.Matcher;
27+
28+
/** Specific executor for the EXPLAIN statement for PostgreSQL. */
29+
class ClientSideStatementExplainExecutor implements ClientSideStatementExecutor {
30+
private final ClientSideStatementImpl statement;
31+
private final Method method;
32+
private final ExplainCommandConverter converter;
33+
public static final Set<String> EXPLAIN_OPTIONS =
34+
ImmutableSet.of(
35+
"verbose", "costs", "settings", "buffers", "wal", "timing", "summary", "format");
36+
37+
ClientSideStatementExplainExecutor(ClientSideStatementImpl statement) throws CompileException {
38+
try {
39+
this.statement = statement;
40+
this.converter = new ExplainCommandConverter();
41+
this.method =
42+
ConnectionStatementExecutor.class.getDeclaredMethod(
43+
statement.getMethodName(), converter.getParameterClass());
44+
} catch (Exception e) {
45+
throw new CompileException(e, statement);
46+
}
47+
}
48+
49+
@Override
50+
public StatementResult execute(ConnectionStatementExecutor connection, String sql)
51+
throws Exception {
52+
return (StatementResult) method.invoke(connection, getParameterValue(sql));
53+
}
54+
55+
String getParameterValue(String sql) {
56+
Matcher matcher = statement.getPattern().matcher(sql);
57+
if (matcher.find() && matcher.groupCount() >= 1) {
58+
String value = matcher.group(0);
59+
if (value != null) {
60+
String res = converter.convert(value.trim());
61+
if (res != null) {
62+
return res;
63+
}
64+
throw SpannerExceptionFactory.newSpannerException(
65+
ErrorCode.INVALID_ARGUMENT, String.format("Invalid argument for EXPLAIN: %s", value));
66+
}
67+
}
68+
return null;
69+
}
70+
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ClientSideStatementValueConverters.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,4 +342,23 @@ public Priority convert(String value) {
342342
return values.get("PRIORITY_" + value);
343343
}
344344
}
345+
346+
static class ExplainCommandConverter implements ClientSideStatementValueConverter<String> {
347+
@Override
348+
public Class<String> getParameterClass() {
349+
return String.class;
350+
}
351+
352+
@Override
353+
public String convert(String value) {
354+
/* The first word in the string should be "explain"
355+
* So, if the size of the string <= 7 (number of letters in the word "explain"), its an invalid statement
356+
* If the size is greater than 7, we'll consider everything after explain as the query.
357+
*/
358+
if (value.length() <= 7) {
359+
return null;
360+
}
361+
return value.substring(7).trim();
362+
}
363+
}
345364
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutor.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,4 +111,6 @@ StatementResult statementSetPgSessionCharacteristicsTransactionMode(
111111
StatementResult statementShowRPCPriority();
112112

113113
StatementResult statementShowTransactionIsolationLevel();
114+
115+
StatementResult statementExplain(String sql);
114116
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionStatementExecutorImpl.java

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,13 @@
5959
import com.google.cloud.spanner.CommitResponse;
6060
import com.google.cloud.spanner.CommitStats;
6161
import com.google.cloud.spanner.Dialect;
62+
import com.google.cloud.spanner.ErrorCode;
6263
import com.google.cloud.spanner.Options.RpcPriority;
64+
import com.google.cloud.spanner.ReadContext.QueryAnalyzeMode;
6365
import com.google.cloud.spanner.ResultSet;
6466
import com.google.cloud.spanner.ResultSets;
67+
import com.google.cloud.spanner.SpannerExceptionFactory;
68+
import com.google.cloud.spanner.Statement;
6569
import com.google.cloud.spanner.Struct;
6670
import com.google.cloud.spanner.TimestampBound;
6771
import com.google.cloud.spanner.Type;
@@ -71,8 +75,11 @@
7175
import com.google.common.base.Preconditions;
7276
import com.google.common.collect.ImmutableMap;
7377
import com.google.protobuf.Duration;
78+
import com.google.spanner.v1.PlanNode;
79+
import com.google.spanner.v1.QueryPlan;
7480
import com.google.spanner.v1.RequestOptions;
7581
import com.google.spanner.v1.RequestOptions.Priority;
82+
import java.util.ArrayList;
7683
import java.util.Collections;
7784
import java.util.Map;
7885
import java.util.concurrent.TimeUnit;
@@ -83,6 +90,7 @@
8390
* calls are then forwarded into a {@link Connection}.
8491
*/
8592
class ConnectionStatementExecutorImpl implements ConnectionStatementExecutor {
93+
8694
static final class StatementTimeoutGetter implements DurationValueGetter {
8795
private final Connection connection;
8896

@@ -442,4 +450,196 @@ public StatementResult statementShowRPCPriority() {
442450
public StatementResult statementShowTransactionIsolationLevel() {
443451
return resultSet("transaction_isolation", "serializable", SHOW_TRANSACTION_ISOLATION_LEVEL);
444452
}
453+
454+
private String processQueryPlan(PlanNode planNode) {
455+
StringBuilder planNodeDescription = new StringBuilder(" : { ");
456+
com.google.protobuf.Struct metadata = planNode.getMetadata();
457+
458+
for (String key : metadata.getFieldsMap().keySet()) {
459+
planNodeDescription
460+
.append(key)
461+
.append(" : ")
462+
.append(metadata.getFieldsMap().get(key).getStringValue())
463+
.append(" , ");
464+
}
465+
String substring = planNodeDescription.substring(0, planNodeDescription.length() - 3);
466+
planNodeDescription.setLength(0);
467+
planNodeDescription.append(substring).append(" }");
468+
469+
return planNodeDescription.toString();
470+
}
471+
472+
private String processExecutionStats(PlanNode planNode) {
473+
StringBuilder executionStats = new StringBuilder("");
474+
for (String key : planNode.getExecutionStats().getFieldsMap().keySet()) {
475+
executionStats.append(key).append(" : { ");
476+
com.google.protobuf.Struct value =
477+
planNode.getExecutionStats().getFieldsMap().get(key).getStructValue();
478+
for (String newKey : value.getFieldsMap().keySet()) {
479+
String newValue = value.getFieldsMap().get(newKey).getStringValue();
480+
executionStats.append(newKey).append(" : ").append(newValue).append(" , ");
481+
}
482+
String substring = executionStats.substring(0, executionStats.length() - 3);
483+
executionStats.setLength(0);
484+
executionStats.append(substring).append(" } , ");
485+
}
486+
String substring = executionStats.substring(0, executionStats.length() - 3);
487+
executionStats.setLength(0);
488+
executionStats.append(substring);
489+
return executionStats.toString();
490+
}
491+
492+
private StatementResult getStatementResultFromQueryPlan(QueryPlan queryPlan, boolean isAnalyze) {
493+
ArrayList<Struct> list = new ArrayList<>();
494+
for (PlanNode planNode : queryPlan.getPlanNodesList()) {
495+
String planNodeDescription = planNode.getDisplayName();
496+
String executionStats = "";
497+
498+
if (!planNode.getMetadata().toString().equalsIgnoreCase("")) {
499+
planNodeDescription += processQueryPlan(planNode);
500+
}
501+
502+
if (!planNode.getShortRepresentation().toString().equalsIgnoreCase("")) {
503+
planNodeDescription += " : " + planNode.getShortRepresentation().getDescription();
504+
}
505+
506+
if (isAnalyze && !planNode.getExecutionStats().toString().equals("")) {
507+
executionStats = processExecutionStats(planNode);
508+
}
509+
Struct.Builder builder = Struct.newBuilder().set("QUERY PLAN").to(planNodeDescription);
510+
511+
if (isAnalyze) {
512+
builder.set("EXECUTION STATS").to(executionStats);
513+
}
514+
list.add(builder.build());
515+
}
516+
517+
ResultSet resultSet;
518+
if (isAnalyze) {
519+
resultSet =
520+
ResultSets.forRows(
521+
Type.struct(
522+
StructField.of("QUERY PLAN", Type.string()),
523+
StructField.of("EXECUTION STATS", Type.string())),
524+
list);
525+
} else {
526+
resultSet =
527+
ResultSets.forRows(Type.struct(StructField.of("QUERY PLAN", Type.string())), list);
528+
}
529+
return StatementResultImpl.of(resultSet);
530+
}
531+
532+
private StatementResult executeStatement(String sql, QueryAnalyzeMode queryAnalyzeMode) {
533+
Statement statement = Statement.newBuilder(sql).build();
534+
try (ResultSet resultSet = getConnection().analyzeQuery(statement, queryAnalyzeMode)) {
535+
while (resultSet.next()) {
536+
// ResultSet.next() should return false in order to access the ResultSet.Stats
537+
}
538+
539+
if (resultSet.getStats() != null && resultSet.getStats().getQueryPlan() != null) {
540+
return getStatementResultFromQueryPlan(
541+
resultSet.getStats().getQueryPlan(), queryAnalyzeMode.equals(QueryAnalyzeMode.PROFILE));
542+
}
543+
}
544+
throw SpannerExceptionFactory.newSpannerException(
545+
ErrorCode.FAILED_PRECONDITION, String.format("Couldn't fetch stats for %s", sql));
546+
}
547+
548+
// This method removes parenthesis from the sql string assuming it is ending with the closing
549+
// parenthesis
550+
private String removeParenthesisAndTrim(String sql) {
551+
sql = sql.trim();
552+
if (sql.charAt(0) == '(') {
553+
sql = sql.substring(1, sql.length() - 1);
554+
}
555+
return sql.trim();
556+
}
557+
558+
/*
559+
* This method executes the given SQL string in either PLAN or PROFILE mode and returns
560+
* the query plan and execution stats (if necessary).
561+
*
562+
* The only additional option that is supported is ANALYZE. The method will throw a SpannerException
563+
* if it is invoked with a statement that includes any other options.
564+
*
565+
* If the SQL string has ANALYZE option, it will be executed in PROFILE mode and will return a resultset
566+
* with two String columns namely QUERY PLAN and EXECUTION STATS.
567+
*
568+
* If the sql string doesn't have any option, it will be executed in PLAN mode and will return a resultset
569+
* with one string column namely QUERY PLAN.
570+
*/
571+
@Override
572+
public StatementResult statementExplain(String sql) {
573+
if (sql == null) {
574+
throw SpannerExceptionFactory.newSpannerException(
575+
ErrorCode.INVALID_ARGUMENT, String.format("Invalid String with Explain"));
576+
}
577+
578+
if (sql.charAt(0) == '(') {
579+
int index = sql.indexOf(')');
580+
if (index == -1) {
581+
throw SpannerExceptionFactory.newSpannerException(
582+
ErrorCode.INVALID_ARGUMENT,
583+
String.format("Missing closing parenthesis in the query: %s", sql));
584+
}
585+
String options[] = sql.substring(1, index).split("\\s*,\\s*");
586+
boolean isAnalyze = false, startAfterIndex = false;
587+
for (String option : options) {
588+
String optionExpression[] = option.trim().split("\\s+");
589+
if (optionExpression.length >= 3) {
590+
isAnalyze = false;
591+
break;
592+
} else if (ClientSideStatementExplainExecutor.EXPLAIN_OPTIONS.contains(
593+
optionExpression[0].toLowerCase())) {
594+
throw SpannerExceptionFactory.newSpannerException(
595+
ErrorCode.UNIMPLEMENTED,
596+
String.format("%s is not implemented yet", optionExpression[0]));
597+
} else if (optionExpression[0].equalsIgnoreCase("analyse")
598+
|| optionExpression[0].equalsIgnoreCase("analyze")) {
599+
isAnalyze = true;
600+
} else {
601+
isAnalyze = false;
602+
break;
603+
}
604+
605+
if (optionExpression.length == 2) {
606+
if (optionExpression[1].equalsIgnoreCase("false")
607+
|| optionExpression[1].equalsIgnoreCase("0")
608+
|| optionExpression[1].equalsIgnoreCase("off")) {
609+
isAnalyze = false;
610+
startAfterIndex = true;
611+
} else if (!(optionExpression[1].equalsIgnoreCase("true")
612+
|| optionExpression[1].equalsIgnoreCase("1")
613+
|| optionExpression[1].equalsIgnoreCase("on"))) {
614+
isAnalyze = false;
615+
break;
616+
}
617+
}
618+
}
619+
if (isAnalyze) {
620+
String newSql = removeParenthesisAndTrim(sql.substring(index + 1));
621+
return executeStatement(newSql, QueryAnalyzeMode.PROFILE);
622+
} else if (startAfterIndex) {
623+
String newSql = removeParenthesisAndTrim(sql.substring(index + 1));
624+
return executeStatement(newSql, QueryAnalyzeMode.PLAN);
625+
} else {
626+
return executeStatement(removeParenthesisAndTrim(sql), QueryAnalyzeMode.PLAN);
627+
}
628+
} else {
629+
String[] arr = sql.split("\\s+", 2);
630+
if (arr.length >= 2) {
631+
String option = arr[0].toLowerCase();
632+
String statementToBeExplained = arr[1];
633+
634+
if (ClientSideStatementExplainExecutor.EXPLAIN_OPTIONS.contains(option)) {
635+
throw SpannerExceptionFactory.newSpannerException(
636+
ErrorCode.UNIMPLEMENTED, String.format("%s is not implemented yet", option));
637+
} else if (option.equals("analyze") || option.equals("analyse")) {
638+
return executeStatement(
639+
removeParenthesisAndTrim(statementToBeExplained), QueryAnalyzeMode.PROFILE);
640+
}
641+
}
642+
return executeStatement(sql, QueryAnalyzeMode.PLAN);
643+
}
644+
}
445645
}

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/StatementResult.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ enum ClientSideStatementType {
8383
ABORT_BATCH,
8484
SET_RPC_PRIORITY,
8585
SHOW_RPC_PRIORITY,
86-
SHOW_TRANSACTION_ISOLATION_LEVEL
86+
SHOW_TRANSACTION_ISOLATION_LEVEL,
87+
EXPLAIN
8788
}
8889

8990
/**

‎google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,15 @@
149149
"method": "statementShowTransactionIsolationLevel",
150150
"exampleStatements": ["show transaction isolation level","show variable transaction isolation level"]
151151
},
152+
{
153+
"name": "EXPLAIN <sql>",
154+
"executorName": "ClientSideStatementExplainExecutor",
155+
"resultType": "RESULT_SET",
156+
"statementType": "EXPLAIN",
157+
"regex": "(?is)\\A\\s*explain(\\s+|\\()(.*)\\z",
158+
"method": "statementExplain",
159+
"exampleStatements": []
160+
},
152161
{
153162
"name": "{START | BEGIN} [TRANSACTION | WORK] [{ (READ ONLY|READ WRITE) | (ISOLATION LEVEL (DEFAULT|SERIALIZABLE)) }]",
154163
"executorName": "ClientSideStatementPgBeginExecutor",

0 commit comments

Comments
 (0)