Reuse cached plan for explain request

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java

@@ -59,6 +59,7 @@
 import java.sql.Struct;
 import java.util.EnumSet;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -152,13 +153,13 @@
     }
 
     private AstNormalizer(@Nonnull final PreparedParams preparedStatementParameters, boolean caseSensitive,
-                          @Nonnull final PlanHashable.PlanHashMode currentPlanHashMode) {
+                          @Nonnull final PlanHashable.PlanHashMode currentPlanHashMode, boolean forExplain) {
         hashFunction = Hashing.murmur3_32_fixed().newHasher();
         parameterHash = Hashing.murmur3_32_fixed().newHasher().putInt("ParameterHash".hashCode());
         parameterHashSupplier = Suppliers.memoize(() -> parameterHash.hash().asInt())::get;
         sqlCanonicalizer = new StringBuilder();
         // needed to collect information that guide query execution (explain flag, continuation string, offset int, and limit int).
-        queryHasherContextBuilder = NormalizedQueryExecutionContext.newBuilder().setPlanHashMode(currentPlanHashMode);
+        queryHasherContextBuilder = NormalizedQueryExecutionContext.newBuilder().setPlanHashMode(currentPlanHashMode).setForExplain(forExplain);
         this.preparedStatementParameters = preparedStatementParameters;
         allowTokenAddition = true;
         allowLiteralAddition = true;
@@ -237,11 +238,13 @@
     }
 
     @Override
     public Void visitFullDescribeStatement(@Nonnull RelationalParser.FullDescribeStatementContext ctx) {
-        // (yhatem) this is probably not needed, since a cached physical plan _knows_ it is either forExplain or not.
-        //          we should remove this, but ok for now.
-        queryHasherContextBuilder.setForExplain(ctx.EXPLAIN() != null);
-        return visitChildren(ctx);
+        // we are stripping Explain/Describe prefix from the request before passing it to the parser
+        throw Assert.failUnchecked("Explain/Describe statement should not appear at the parser lavel");
+        // // (yhatem) this is probably not needed, since a cached physical plan _knows_ it is either forExplain or not.
+        // //          we should remove this, but ok for now.
+        // queryHasherContextBuilder.setForExplain(ctx.EXPLAIN() != null);
+        // return visitChildren(ctx);
     }
 
     @Override
@@ -569,6 +572,87 @@
         }
     }
 
+    @VisibleForTesting
+    public static class ExplainParser {
+        private ExplainParser() {
+        }
+
+        private static class Word {
+            public int begin;
+            public int end;
+
+            public Word(int b, int e) {
+                begin = b;
+                end = e;
+            }
+        }
+
+        private static int getNotWhitespace(@Nonnull final String query, int pos) {
+            while (pos < query.length() && Character.isWhitespace(query.charAt(pos))) {
+                pos++;
+            }
+            return pos;
+        }
+
+        private static int getWhitespaceEqual(@Nonnull final String query, int pos) {
+            while (pos < query.length() && !Character.isWhitespace(query.charAt(pos)) && query.charAt(pos) != '=') {
+                pos++;
+            }
+            return pos;
+        }
+
+        private static Word getWordFromString(@Nonnull final String query, int start) {
+            int pos = getNotWhitespace(query, start);
+            return new Word(pos, getWhitespaceEqual(query, pos));
+        }
+
+        private static boolean matchWord(@Nonnull final String query, @Nonnull final Word w, @Nonnull final String word) {
+            return w.end - w.begin == word.length() && query.regionMatches(true, w.begin, word, 0, word.length());
+        }
+
+        private static boolean matchWordFromList(@Nonnull final String query, @Nonnull final Word w, @Nonnull final List<String> list) {
+            for (final var word : list) {
+                if (matchWord(query, w, word)) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        public static String truncateExplainStatement(@Nonnull final String query) {
+            final List<String> explainList = List.of("EXPLAIN", "DESCRIBE", "DESC");
+            final String shemaString = "SCHEMA";
+            final List<String> extendedList = List.of("EXTENDED", "PARTITIONS", "FORMAT");
+            final List<String> traditionalList = List.of("TRADITIONAL", "JSON");
+
+
+            final var w0 = getWordFromString(query, 0);
+            if (!matchWordFromList(query, w0, explainList)) {       // EXPLAIN
+                return query;
+            }
+
+            final var w1 = getWordFromString(query, w0.end);
+            if (matchWord(query, w1, shemaString)) {                // EXPLAIN SCHEMA
+                return query;
+            }
+            if (!matchWordFromList(query, w1, extendedList)) {      // EXPLAIN EXTENDED
+                int pos = getNotWhitespace(query, w1.begin);
+                Assert.thatUnchecked(pos != query.length(), ErrorCode.SYNTAX_ERROR, "no statement to explain");
+                return query.substring(pos);
+            }
+
+            int pos = getNotWhitespace(query, w1.end);              // EXPLAIN EXTENDED=
+            Assert.thatUnchecked(query.charAt(pos) == '=', ErrorCode.SYNTAX_ERROR, "equal (=) not found after EXTENDED/PARTITIONS/FORMAT");
+
+            final var w2 = getWordFromString(query, pos + 1);   // EXPLAIN EXTENDED=TRADITIONAL
+            Assert.thatUnchecked(matchWordFromList(query, w2, traditionalList), ErrorCode.SYNTAX_ERROR, "value of EXTENDED/PARTITIONS/FORMAT is not TRADITIONAL/JSON");
+
+            pos = getNotWhitespace(query, w2.end);
+            Assert.thatUnchecked(pos != query.length(), ErrorCode.SYNTAX_ERROR, "no statement to explain");
+            return query.substring(w2.end);
+        }
+    }
+
     /**
      * Normalizes the SQL query using a given planning context.
      * @param context The planning context, captures all the state required to plan and execution query such as prepared parameter values.
@@ -583,9 +667,10 @@
                                                      boolean isCaseSensitive,
                                                      @Nonnull final PlanHashable.PlanHashMode currentPlanHashMode) throws RelationalException {
         // lexing, parsing, and normalization are profiled through the metric collector.
+        final var truncQuery = ExplainParser.truncateExplainStatement(query);
         final var metricCollector = context.getMetricsCollector();
         final var rootContext = metricCollector.clock(RelationalMetric.RelationalEvent.LEX_PARSE,
-                () -> QueryParser.parse(query).getRootContext());
+                () -> QueryParser.parse(truncQuery).getRootContext());
         return metricCollector.clock(RelationalMetric.RelationalEvent.NORMALIZE_QUERY,
                 () -> normalizeAst(
                         context.getSchemaTemplate(), rootContext,
@@ -594,7 +679,8 @@
                         context.getPlannerConfiguration(),
                         isCaseSensitive,
                         currentPlanHashMode,
-                        query
+                        query,
+                        !Objects.equals(query, truncQuery)
                 ));
     }
 
@@ -607,8 +693,9 @@
                                                    @Nonnull final PlannerConfiguration plannerConfiguration,
                                                    boolean caseSensitive,
                                                    @Nonnull final PlanHashable.PlanHashMode currentPlanHashMode,
-                                                   @Nonnull final String query) throws RelationalException {
-        final var astNormalizer = new AstNormalizer(preparedStatementParameters, caseSensitive, currentPlanHashMode);
+                                                   @Nonnull final String query,
+                                                   boolean forExplain) throws RelationalException {
+        final var astNormalizer = new AstNormalizer(preparedStatementParameters, caseSensitive, currentPlanHashMode, forExplain);
         astNormalizer.visit(context);
         final var recordLayerSchemaTemplate = Assert.castUnchecked(schemaTemplate, RecordLayerSchemaTemplate.class);
 
@@ -637,7 +724,7 @@
                         plannerConfiguration,
                         caseSensitive,
                         currentPlanHashMode,
-                        recordLayerRoutine.getDescription());
+                        recordLayerRoutine.getDescription(), forExplain);
                 astNormalizer.queryHasherContextBuilder.getLiteralsBuilder().importLiterals(functionAstResult.queryExecutionContext.getLiterals());
             }
         }

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java

@@ -242,6 +242,7 @@ private Plan<?> generatePhysicalPlanForCompilableStatement(@Nonnull AstNormalize
 
         final var planGenerationContext = new MutablePlanGenerationContext(planContext.getPreparedStatementParameters(),
                 currentPlanHashMode, ast.getQuery(), ast.getQueryCacheKey().getCanonicalQueryString(), parameterHash);
+        planGenerationContext.setForExplain(ast.getQueryExecutionContext().isForExplain());
         final var metadata = Assert.castUnchecked(planContext.getSchemaTemplate(), RecordLayerSchemaTemplate.class);
         try {
             final var maybePlan = planContext.getMetricsCollector().clock(RelationalMetric.RelationalEvent.GENERATE_LOGICAL_PLAN, () ->

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/QueryVisitor.java

@@ -597,13 +597,15 @@
 
     @Nonnull
     @Override
     public QueryPlan.LogicalQueryPlan visitFullDescribeStatement(@Nonnull RelationalParser.FullDescribeStatementContext ctx) {
-        getDelegate().getPlanGenerationContext().setForExplain(ctx.EXPLAIN() != null);
-        final var logicalOperator = Assert.castUnchecked(ctx.describeObjectClause().accept(this), LogicalOperator.class);
-        // Capture semantic type structure as StructType with field names
-        final var semanticStructType = logicalOperator.getOutput().getStructType();
-        return QueryPlan.LogicalQueryPlan.of(logicalOperator.getQuantifier().getRangesOver().get(),
-                getDelegate().getPlanGenerationContext(), getDelegate().getPlanGenerationContext().getQuery(), semanticStructType);
+        // we are stripping Explain/Describe prefix from the request before passing it to the parser
+        throw Assert.failUnchecked("Explain/Describe statement should not appear at the parser lavel");
+        // getDelegate().getPlanGenerationContext().setForExplain(ctx.EXPLAIN() != null);
+        // final var logicalOperator = Assert.castUnchecked(ctx.describeObjectClause().accept(this), LogicalOperator.class);
+        // // Capture semantic type structure as StructType with field names
+        // final var semanticStructType = logicalOperator.getOutput().getStructType();
+        // return QueryPlan.LogicalQueryPlan.of(logicalOperator.getQuantifier().getRangesOver().get(),
+        //        getDelegate().getPlanGenerationContext(), getDelegate().getPlanGenerationContext().getQuery(), semanticStructType);
     }
 
     @Nonnull

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/PlanGenerationStackTest.java

@@ -20,17 +20,26 @@
 
 package com.apple.foundationdb.relational.recordlayer;
 
+import com.apple.foundationdb.record.PlanHashable;
 import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase;
 import com.apple.foundationdb.relational.api.Options;
+import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
 import com.apple.foundationdb.relational.api.exceptions.RelationalException;
+import com.apple.foundationdb.relational.api.exceptions.UncheckedRelationalException;
+import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerSchemaTemplate;
+import com.apple.foundationdb.relational.recordlayer.query.MutablePlanGenerationContext;
 import com.apple.foundationdb.relational.recordlayer.query.Plan;
 import com.apple.foundationdb.relational.recordlayer.query.PlanContext;
 import com.apple.foundationdb.relational.recordlayer.query.PlanGenerator;
+import com.apple.foundationdb.relational.recordlayer.query.QueryParser;
 import com.apple.foundationdb.relational.recordlayer.query.QueryPlan;
+import com.apple.foundationdb.relational.recordlayer.query.visitors.BaseVisitor;
+import com.apple.foundationdb.relational.util.Assert;
 import com.apple.foundationdb.relational.utils.SimpleDatabaseRule;
 import com.apple.foundationdb.relational.utils.TestSchemas;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -46,6 +55,7 @@
 import java.util.stream.Stream;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
  * A series of tests that target the following.
@@ -216,4 +226,46 @@ void queryTestHarness(int ignored, @Nonnull String query, @Nullable String error
             }
         }
     }
+
+    @Test
+    void queryVisitorThrowsWhenExplainReachesParser() throws Exception {
+        // This test covers test gap
+        // QueryVisitor.visitFullDescribeStatement must never be reached because EXPLAIN is stripped
+        // before parsing. This test verifies the defensive assertion fires if that invariant is broken.
+        final var explainQuery = "EXPLAIN SELECT * FROM restaurant";
+        final String schemaName = connection.getSchema();
+        final EmbeddedRelationalConnection embeddedConnection = (EmbeddedRelationalConnection) connection.connection;
+        embeddedConnection.setAutoCommit(false);
+        embeddedConnection.createNewTransaction();
+        try {
+            final AbstractDatabase db = embeddedConnection.getRecordLayerDatabase();
+            final FDBRecordStoreBase<?> store = db.loadSchema(schemaName).loadStore().unwrap(FDBRecordStoreBase.class);
+            final PlanContext planContext = PlanContext.Builder
+                    .create()
+                    .fromDatabase(db)
+                    .fromRecordStore(store, connection.getOptions())
+                    .withSchemaTemplate(embeddedConnection.getSchemaTemplate())
+                    .withMetricsCollector(embeddedConnection.getMetricCollector())
+                    .build();
+            final var parseTree = QueryParser.parse(explainQuery).getRootContext();
+            final var planGenContext = new MutablePlanGenerationContext(
+                    planContext.getPreparedStatementParameters(),
+                    PlanHashable.PlanHashMode.VC0,
+                    explainQuery,
+                    explainQuery,
+                    0);
+            final var metadata = Assert.castUnchecked(planContext.getSchemaTemplate(), RecordLayerSchemaTemplate.class);
+            final var baseVisitor = new BaseVisitor(planGenContext, metadata,
+                    planContext.getDdlQueryFactory(), planContext.getConstantActionFactory(),
+                    planContext.getDbUri(), false);
+            assertThatThrownBy(() -> baseVisitor.generateLogicalPlan(parseTree))
+                    .isInstanceOf(UncheckedRelationalException.class)
+                    .hasMessageContaining("Explain/Describe statement should not appear at the parser")
+                    .extracting(e -> ((UncheckedRelationalException) e).unwrap().getErrorCode())
+                    .isEqualTo(ErrorCode.INTERNAL_ERROR);
+        } finally {
+            embeddedConnection.rollback();
+            embeddedConnection.setAutoCommit(true);
+        }
+    }
 }