diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java index 0f62f05e4f..762338b5ac 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaData.java @@ -87,6 +87,8 @@ public class RecordMetaData implements RecordMetaDataProvider { @Nonnull private final Map viewMap; @Nonnull + private final Map prepareStatements; + @Nonnull private final Map indexes; @Nonnull private final Map universalIndexes; @@ -118,6 +120,7 @@ protected RecordMetaData(@Nonnull RecordMetaData orig) { Collections.unmodifiableList(orig.formerIndexes), Collections.unmodifiableMap(orig.userDefinedFunctionMap), Collections.unmodifiableMap(orig.viewMap), + Collections.unmodifiableMap(orig.prepareStatements), orig.splitLongRecords, orig.storeRecordVersions, orig.version, @@ -139,6 +142,7 @@ protected RecordMetaData(@Nonnull Descriptors.FileDescriptor recordsDescriptor, @Nonnull List formerIndexes, @Nonnull Map userDefinedFunctionMap, @Nonnull Map viewMap, + @Nonnull Map prepareStatements, boolean splitLongRecords, boolean storeRecordVersions, int version, @@ -157,6 +161,7 @@ protected RecordMetaData(@Nonnull Descriptors.FileDescriptor recordsDescriptor, this.formerIndexes = formerIndexes; this.userDefinedFunctionMap = userDefinedFunctionMap; this.viewMap = viewMap; + this.prepareStatements = prepareStatements; this.splitLongRecords = splitLongRecords; this.storeRecordVersions = storeRecordVersions; this.version = version; @@ -704,6 +709,7 @@ public RecordMetaDataProto.MetaData toProto(@Nullable Descriptors.FileDescriptor builder.addAllUserDefinedFunctions(userDefinedFunctionMap.values().stream().map(UserDefinedFunction::toProto).collect(Collectors.toList())); builder.addAllViews(viewMap.values().stream().map(View::toProto).collect(Collectors.toList())); + builder.putAllPrepareStatements(prepareStatements); builder.setSplitLongRecords(splitLongRecords); builder.setStoreRecordVersions(storeRecordVersions); builder.setVersion(version); @@ -728,6 +734,11 @@ public Map getViewMap() { return viewMap; } + @Nonnull + public Map getPrepareStatements() { + return prepareStatements; + } + @Nonnull public Type.Record getPlannerType(@Nonnull String recordTypeName) { final RecordType recordType = getRecordType(recordTypeName); diff --git a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java index 206e73c01e..9ed0d0bfb0 100644 --- a/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java +++ b/fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/RecordMetaDataBuilder.java @@ -116,6 +116,8 @@ public class RecordMetaDataBuilder implements RecordMetaDataProvider { @Nonnull private final Map viewMap; @Nonnull + private final Map prepareStatements; + @Nonnull private final Map indexes; @Nonnull private final Map universalIndexes; @@ -152,6 +154,7 @@ public class RecordMetaDataBuilder implements RecordMetaDataProvider { syntheticRecordTypes = new HashMap<>(); userDefinedFunctionMap = new HashMap<>(); viewMap = new HashMap<>(); + prepareStatements = new HashMap<>(); } private void processSchemaOptions(boolean processExtensionOptions) { @@ -238,6 +241,7 @@ private void loadProtoExceptRecords(@Nonnull RecordMetaDataProto.MetaData metaDa final View view = View.fromProto(viewProto); viewMap.put(view.getName(), view); } + prepareStatements.putAll(metaDataProto.getPrepareStatementsMap()); if (metaDataProto.hasSplitLongRecords()) { splitLongRecords = metaDataProto.getSplitLongRecords(); } @@ -1215,6 +1219,15 @@ public void addView(@Nonnull View view) { viewMap.put(view.getName(), view); } + @Nonnull + public Map getPrepareStatements() { + return prepareStatements; + } + + public void addPrepareStatement(@Nonnull String name, @Nonnull String prepareStatement) { + prepareStatements.put(name, prepareStatement); + } + public boolean isSplitLongRecords() { return splitLongRecords; } @@ -1456,7 +1469,7 @@ public RecordMetaData build(boolean validate) { Map> recordTypeKeyToSyntheticRecordTypeMap = Maps.newHashMapWithExpectedSize(syntheticRecordTypes.size()); RecordMetaData metaData = new RecordMetaData(recordsDescriptor, getUnionDescriptor(), unionFields, builtRecordTypes, builtSyntheticRecordTypes, recordTypeKeyToSyntheticRecordTypeMap, - indexes, universalIndexes, formerIndexes, userDefinedFunctionMap, viewMap, + indexes, universalIndexes, formerIndexes, userDefinedFunctionMap, viewMap, prepareStatements, splitLongRecords, storeRecordVersions, version, subspaceKeyCounter, usesSubspaceKeyCounter, recordCountKey, localFileDescriptor != null); for (RecordTypeBuilder recordTypeBuilder : recordTypes.values()) { KeyExpression primaryKey = recordTypeBuilder.getPrimaryKey(); diff --git a/fdb-record-layer-core/src/main/proto/record_metadata.proto b/fdb-record-layer-core/src/main/proto/record_metadata.proto index 303dc6d3d8..b9b6da2694 100644 --- a/fdb-record-layer-core/src/main/proto/record_metadata.proto +++ b/fdb-record-layer-core/src/main/proto/record_metadata.proto @@ -208,6 +208,7 @@ message MetaData { repeated UnnestedRecordType unnested_record_types = 13; repeated PUserDefinedFunction user_defined_functions = 14; repeated PView views = 15; + map prepare_statements = 16; extensions 1000 to 2000; } diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/MetaDataProtoEditorUnitTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/MetaDataProtoEditorUnitTest.java index 95d4dfb6ff..bb40904566 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/MetaDataProtoEditorUnitTest.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/MetaDataProtoEditorUnitTest.java @@ -587,6 +587,7 @@ void validateMetaDataCoverage() { assertEquals(Set.of( "split_long_records", "version", "former_indexes", "record_count_key", "store_record_versions", "dependencies", "subspace_key_counter", "uses_subspace_key_counter", + "prepare_statements", // the below reference record types "records", "indexes", "record_types", "joined_record_types", "unnested_record_types", "user_defined_functions", "views"), diff --git a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/SchemaTemplate.java b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/SchemaTemplate.java index 29cd7ab878..0df40f200e 100644 --- a/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/SchemaTemplate.java +++ b/fdb-relational-api/src/main/java/com/apple/foundationdb/relational/api/metadata/SchemaTemplate.java @@ -27,6 +27,7 @@ import javax.annotation.Nonnull; import java.util.BitSet; import java.util.Collection; +import java.util.Map; import java.util.Optional; import java.util.Set; @@ -130,6 +131,14 @@ public interface SchemaTemplate extends Metadata { @Nonnull Collection getTemporaryInvokedRoutines() throws RelationalException; + /** + * Returns the prepare statements defined in this schema template. + * + * @return A map of prepare statement names to their SQL strings. + */ + @Nonnull + Map getPrepareStatements(); + @Nonnull String getTransactionBoundMetadataAsString() throws RelationalException; diff --git a/fdb-relational-core/src/main/antlr/RelationalParser.g4 b/fdb-relational-core/src/main/antlr/RelationalParser.g4 index 604d2c56b9..68576b4728 100644 --- a/fdb-relational-core/src/main/antlr/RelationalParser.g4 +++ b/fdb-relational-core/src/main/antlr/RelationalParser.g4 @@ -92,6 +92,7 @@ utilityStatement templateClause : CREATE ( structDefinition | tableDefinition | enumDefinition | indexDefinition | sqlInvokedFunction | viewDefinition ) + | prepareStatement ; createStatement diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplate.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplate.java index 2a053aabb2..e1de98f4cd 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplate.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplate.java @@ -35,6 +35,8 @@ import javax.annotation.Nonnull; import java.util.BitSet; import java.util.Collection; +import java.util.Collections; +import java.util.Map; import java.util.Optional; import java.util.Set; @@ -139,6 +141,12 @@ public Collection getTemporaryInvokedRoutines() throws Relationa throw new RelationalException("NoOpSchemaTemplate doesn't have temporary invoked routines!", ErrorCode.INVALID_PARAMETER); } + @Nonnull + @Override + public Map getPrepareStatements() { + return Collections.emptyMap(); + } + @Nonnull @Override public String getTransactionBoundMetadataAsString() throws RelationalException { diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java index e9c6d72b0e..3969032863 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/RecordLayerSchemaTemplate.java @@ -80,6 +80,9 @@ public final class RecordLayerSchemaTemplate implements SchemaTemplate { @Nonnull private final Set views; + @Nonnull + private final Map prepareStatements; + private final int version; private final boolean enableLongRows; @@ -107,6 +110,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name, @Nonnull final Set tables, @Nonnull final Set invokedRoutines, @Nonnull final Set views, + @Nonnull final Map prepareStatements, int version, boolean enableLongRows, boolean storeRowVersions, @@ -115,6 +119,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name, this.tables = ImmutableSet.copyOf(tables); this.invokedRoutines = ImmutableSet.copyOf(invokedRoutines); this.views = ImmutableSet.copyOf(views); + this.prepareStatements = ImmutableMap.copyOf(prepareStatements); this.version = version; this.enableLongRows = enableLongRows; this.storeRowVersions = storeRowVersions; @@ -130,6 +135,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name, @Nonnull final Set tables, @Nonnull final Set invokedRoutines, @Nonnull final Set views, + @Nonnull final Map prepareStatements, int version, boolean enableLongRows, boolean storeRowVersions, @@ -140,6 +146,7 @@ private RecordLayerSchemaTemplate(@Nonnull final String name, this.tables = ImmutableSet.copyOf(tables); this.invokedRoutines = ImmutableSet.copyOf(invokedRoutines); this.views = ImmutableSet.copyOf(views); + this.prepareStatements = ImmutableMap.copyOf(prepareStatements); this.enableLongRows = enableLongRows; this.storeRowVersions = storeRowVersions; this.intermingleTables = intermingleTables; @@ -338,6 +345,12 @@ public Set getViews() { return views; } + @Nonnull + @Override + public Map getPrepareStatements() { + return prepareStatements; + } + @Nonnull @Override public Optional findViewByName(@Nonnull final String viewName) { @@ -414,6 +427,9 @@ public static final class Builder { @Nonnull private final Map views; + @Nonnull + private final Map prepareStatements; + private RecordMetaData cachedMetadata; @@ -422,6 +438,7 @@ private Builder() { auxiliaryTypes = new LinkedHashMap<>(); invokedRoutines = new LinkedHashMap<>(); views = new LinkedHashMap<>(); + prepareStatements = new LinkedHashMap<>(); // enable long rows is TRUE by default enableLongRows = true; } @@ -540,6 +557,18 @@ public Builder addViews(@Nonnull final Collection views) { return this; } + @Nonnull + public Builder addPrepareStatement(@Nonnull final String name, @Nonnull final String prepareStatement) { + prepareStatements.put(name, prepareStatement); + return this; + } + + @Nonnull + public Builder addPrepareStatements(@Nonnull final Map prepareStatements) { + this.prepareStatements.putAll(prepareStatements); + return this; + } + /** * Adds an auxiliary type, an auxiliary type is a type that is merely created, so it can be referenced later on * in a table definition. Any {@link DataType.Named} data type can be added as an auxiliary type such as {@code enum}s @@ -632,10 +661,10 @@ public RecordLayerSchemaTemplate build() { if (cachedMetadata != null) { return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), - new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), version, enableLongRows, storeRowVersions, intermingleTables, cachedMetadata); + new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), prepareStatements, version, enableLongRows, storeRowVersions, intermingleTables, cachedMetadata); } else { return new RecordLayerSchemaTemplate(name, new LinkedHashSet<>(tables.values()), - new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), version, enableLongRows, storeRowVersions, intermingleTables); + new LinkedHashSet<>(invokedRoutines.values()), new LinkedHashSet<>(views.values()), prepareStatements, version, enableLongRows, storeRowVersions, intermingleTables); } } @@ -763,6 +792,7 @@ public Builder toBuilder() { .setIntermingleTables(intermingleTables) .addTables(getTables()) .addInvokedRoutines(getInvokedRoutines()) - .addViews(getViews()); + .addViews(getViews()) + .addPrepareStatements(getPrepareStatements()); } } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataDeserializer.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataDeserializer.java index 68b4c15bdc..c563c6ced0 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataDeserializer.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataDeserializer.java @@ -125,6 +125,7 @@ private static RecordLayerSchemaTemplate.Builder deserializeRecordMetaData(@Nonn schemaTemplateBuilder.addView(generateViewBuilder(metadataProvider, view.getKey(), view.getValue().getDefinition()).build()); } } + schemaTemplateBuilder.addPrepareStatements(recordMetaData.getPrepareStatements()); schemaTemplateBuilder.setCachedMetadata(recordMetaData); return schemaTemplateBuilder; } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataSerializer.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataSerializer.java index 11b80a071c..63a2c0a4eb 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataSerializer.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/metadata/serde/RecordMetadataSerializer.java @@ -107,9 +107,13 @@ public void visit(@Nonnull final View view) { @Override public void visit(@Nonnull SchemaTemplate schemaTemplate) { Assert.thatUnchecked(schemaTemplate instanceof RecordLayerSchemaTemplate); + final var recLayerSchemaTemplate = (RecordLayerSchemaTemplate) schemaTemplate; getBuilder().setSplitLongRecords(schemaTemplate.isEnableLongRows()); getBuilder().setStoreRecordVersions(schemaTemplate.isStoreRowVersions()); getBuilder().setVersion(schemaTemplate.getVersion()); + for (final var entry : recLayerSchemaTemplate.getPrepareStatements().entrySet()) { + getBuilder().addPrepareStatement(entry.getKey(), entry.getValue()); + } } @Nonnull diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java index 6066385e16..41bab939e5 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java @@ -26,6 +26,7 @@ import com.apple.foundationdb.record.RecordMetaData; import com.apple.foundationdb.record.RecordStoreState; import com.apple.foundationdb.record.logging.KeyValueLogMessage; +import com.apple.foundationdb.record.logging.LogMessageKeys; import com.apple.foundationdb.record.metadata.MetaDataException; import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase; import com.apple.foundationdb.record.provider.foundationdb.IndexMatchCandidateRegistry; @@ -140,6 +141,35 @@ public Plan getPlan(@Nonnull final String query) throws RelationalException { return plan; } + /** + * Pre-generates and caches plans for the prepare statements defined in the schema template. + * This method is idempotent per template name and version — subsequent calls for the same + * template are no-ops. + */ + public void prepareStatements() { + if (cache.isEmpty()) { + return; + } + final var schemaTemplate = planContext.getSchemaTemplate(); + if (schemaTemplate.getPrepareStatements().isEmpty()) { + return; + } + final var templateKey = schemaTemplate.getName() + ":" + schemaTemplate.getVersion(); + if (cache.get().isPrepared(templateKey)) { + return; + } + for (final var entry : schemaTemplate.getPrepareStatements().entrySet()) { + try { + getPlan(entry.getValue()); + } catch (RelationalException e) { + if (logger.isErrorEnabled()) { + logger.error(KeyValueLogMessage.of("prepare statement", LogMessageKeys.QUERY, entry.getValue()), e); + } + } + } + cache.get().markPrepared(templateKey); + } + private boolean isCaseSensitive() { return options.getOption(Options.Name.CASE_SENSITIVE_IDENTIFIERS); } @@ -480,7 +510,9 @@ public static PlanGenerator create(@Nonnull final Optional @Nonnull final Options options) throws RelationalException { final var planner = new CascadesPlanner(metaData, recordStoreState, matchCandidateRegistry); planner.setConfiguration(planContext.getRecordQueryPlannerConfiguration()); - return new PlanGenerator(cache, planContext, planner, options); + final var planGenerator = new PlanGenerator(cache, planContext, planner, options); + planGenerator.prepareStatements(); + return planGenerator; } /** diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/RelationalPlanCache.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/RelationalPlanCache.java index 13e8192d8a..02f8872c0f 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/RelationalPlanCache.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/RelationalPlanCache.java @@ -27,6 +27,8 @@ import javax.annotation.Nonnull; import javax.annotation.Nullable; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executor; import java.util.concurrent.TimeUnit; @@ -45,6 +47,9 @@ public final class RelationalPlanCache extends MultiStageCache preparedTemplates = ConcurrentHashMap.newKeySet(); + private RelationalPlanCache(int size, int secondarySize, int tertiarySize, @@ -102,4 +107,12 @@ public static RelationalPlanCache buildWithDefaults() { return newRelationalCacheBuilder().build(); } + public boolean isPrepared(@Nonnull String templateKey) { + return preparedTemplates.contains(templateKey); + } + + public void markPrepared(@Nonnull String templateKey) { + preparedTemplates.add(templateKey); + } + } diff --git a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java index 17483c5c64..2bd4908f3c 100644 --- a/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java +++ b/fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/visitors/DdlVisitor.java @@ -414,6 +414,17 @@ public ProceduralPlan visitCreateSchemaTemplateStatement(@Nonnull RelationalPars sqlInvokedFunctionClauses.add(templateClause.sqlInvokedFunction()); } else if (templateClause.viewDefinition() != null) { viewClauses.add(templateClause.viewDefinition()); + } else if (templateClause.prepareStatement() != null) { + final var prepareCtx = templateClause.prepareStatement(); + final var name = visitUid(prepareCtx.uid()).getName(); + final String queryString; + if (prepareCtx.queryString != null) { + final var text = prepareCtx.queryString.getText(); + queryString = text.substring(1, text.length() - 1); + } else { + queryString = prepareCtx.variable.getText(); + } + metadataBuilder.addPrepareStatement(name, queryString); } else { Assert.thatUnchecked(templateClause.indexDefinition() != null); indexClauses.add(templateClause.indexDefinition()); diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplateTests.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplateTests.java index 2698dbb7da..5e59ea451c 100644 --- a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplateTests.java +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/metadata/NoOpSchemaTemplateTests.java @@ -234,4 +234,11 @@ public void testGenerateSchemaWithDifferentParameters() { assertEquals("schema1", schema1.getName()); assertEquals("schema2", schema2.getName()); } + + @Test + public void testGetPrepareStatementsReturnsEmptyMap() { + final NoOpSchemaTemplate template = new NoOpSchemaTemplate("test", 1); + assertNotNull(template.getPrepareStatements()); + assertEquals(0, template.getPrepareStatements().size()); + } } diff --git a/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/SchemaTemplatePrepareTest.java b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/SchemaTemplatePrepareTest.java new file mode 100644 index 0000000000..61802e44df --- /dev/null +++ b/fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/SchemaTemplatePrepareTest.java @@ -0,0 +1,226 @@ +/* + * SchemaTemplatePrepareTest.java + * + * This source file is part of the FoundationDB open source project + * + * Copyright 2021-2026 Apple Inc. and the FoundationDB project authors + * + * 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.apple.foundationdb.relational.recordlayer.query; + +import com.apple.foundationdb.relational.api.RelationalConnection; +import com.apple.foundationdb.relational.api.RelationalResultSet; +import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalConnection; +import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalExtension; +import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerSchemaTemplate; +import com.apple.foundationdb.relational.recordlayer.query.cache.QueryCacheKey; +import com.apple.foundationdb.relational.recordlayer.query.cache.RelationalPlanCache; +import com.apple.foundationdb.relational.utils.Ddl; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.net.URI; +import java.sql.SQLException; + +public class SchemaTemplatePrepareTest { + + private static final String SCHEMA_TEMPLATE = + "CREATE TABLE t1(id bigint, col1 bigint, col2 bigint, PRIMARY KEY(id))" + + " CREATE INDEX i1 AS SELECT col1 FROM t1" + + " PREPARE by_col1 FROM 'select * from t1 where col1 = 10'" + + " PREPARE by_id FROM 'select * from t1 where id = 1'"; + + @RegisterExtension + @Order(0) + public final EmbeddedRelationalExtension relationalExtension = new EmbeddedRelationalExtension(); + + private long countCachedPlans(RelationalConnection connection, String templateName) throws SQLException { + final var embeddedConnection = connection.unwrap(EmbeddedRelationalConnection.class); + final RelationalPlanCache cache = embeddedConnection.getRecordLayerDatabase().getPlanCache(); + if (cache == null) { + return 0; + } + final Long count = cache.getStats().numSecondaryEntries(templateName); + return count != null ? count : 0; + } + + private void showCache(RelationalConnection connection) throws SQLException { + final var embeddedConnection = connection.unwrap(EmbeddedRelationalConnection.class); + final RelationalPlanCache cache = embeddedConnection.getRecordLayerDatabase().getPlanCache(); + if (cache == null) { + System.out.println("[CACHE] no plan cache"); + return; + } + for (String key : cache.getStats().getAllKeys()) { + System.out.println("[CACHE] template: " + key); + for (QueryCacheKey secondaryKey : cache.getStats().getAllSecondaryKeys(key)) { + System.out.println("[CACHE] query: " + secondaryKey.getCanonicalQueryString() + + " (version=" + secondaryKey.getSchemaTemplateVersion() + + ", userVersion=" + secondaryKey.getUserVersion() + ")"); + var tertiaryMappings = cache.getStats().getAllTertiaryMappings(key, secondaryKey); + for (var entry : tertiaryMappings.entrySet()) { + System.out.println("[CACHE] plan: " + entry.getValue().explain()); + } + } + } + } + + @Test + void prepareStatementsStoredInTemplate() throws Exception { + try (var ddl = Ddl.builder() + .database(URI.create("/TEST/PREPARE_DB")) + .relationalExtension(relationalExtension) + .schemaTemplate(SCHEMA_TEMPLATE) + .build()) { + final var connection = ddl.setSchemaAndGetConnection(); + final var embeddedConnection = connection.unwrap(EmbeddedRelationalConnection.class); + embeddedConnection.setAutoCommit(false); + embeddedConnection.createNewTransaction(); + final var schemaTemplate = embeddedConnection.getSchemaTemplate().unwrap(RecordLayerSchemaTemplate.class); + embeddedConnection.rollback(); + embeddedConnection.setAutoCommit(true); + final var prepareStatements = schemaTemplate.getPrepareStatements(); + Assertions.assertEquals(2, prepareStatements.size()); + Assertions.assertEquals("select * from t1 where col1 = 10", prepareStatements.get("BY_COL1")); + Assertions.assertEquals("select * from t1 where id = 1", prepareStatements.get("BY_ID")); + Assertions.assertEquals(0, countCachedPlans(connection, ddl.getSchemaTemplateName())); // this is not good + } + } + + @Test + void prepareStatementsAfterFirstQuery() throws Exception { + try (var ddl = Ddl.builder() + .database(URI.create("/TEST/PREPARE_DB3")) + .relationalExtension(relationalExtension) + .schemaTemplate(SCHEMA_TEMPLATE) + .build()) { + final var connection = ddl.setSchemaAndGetConnection(); + Assertions.assertEquals(0, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement()) { + stmt.execute("INSERT INTO T1 VALUES (1, 10, 1)"); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + } + } + + @Test + void prepareStatementsUsage() throws Exception { + try (var ddl = Ddl.builder() + .database(URI.create("/TEST/PREPARE_DB2")) + .relationalExtension(relationalExtension) + .schemaTemplate(SCHEMA_TEMPLATE) + .build()) { + final var connection = ddl.setSchemaAndGetConnection(); + Assertions.assertEquals(0, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement()) { + stmt.execute("INSERT INTO T1 VALUES (1, 10, 1)"); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement(); RelationalResultSet rs = stmt.executeQuery("select * from t1 where col1 = 10")) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(1, rs.getLong("ID")); + Assertions.assertFalse(rs.next()); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement(); RelationalResultSet rs = stmt.executeQuery("select * from t1 where id = 1")) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(1, rs.getLong("ID")); + Assertions.assertEquals(10, rs.getLong("COL1")); + Assertions.assertEquals(1, rs.getLong("COL2")); + Assertions.assertFalse(rs.next()); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + } + } + + @Test + void prepareStatementsUsageParams() throws Exception { + try (var ddl = Ddl.builder() + .database(URI.create("/TEST/PREPARE_DB2")) + .relationalExtension(relationalExtension) + .schemaTemplate(SCHEMA_TEMPLATE) + .build()) { + final var connection = ddl.setSchemaAndGetConnection(); + Assertions.assertEquals(0, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement()) { + stmt.execute("INSERT INTO T1 VALUES (1, 10, 1)"); + stmt.execute("INSERT INTO T1 VALUES (2, 20, 2)"); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement(); RelationalResultSet rs = stmt.executeQuery("select * from t1 where col1 = 20")) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(2, rs.getLong("ID")); + Assertions.assertFalse(rs.next()); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var stmt = connection.createStatement(); RelationalResultSet rs = stmt.executeQuery("select * from t1 where id = 2")) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(2, rs.getLong("ID")); + Assertions.assertEquals(20, rs.getLong("COL1")); + Assertions.assertEquals(2, rs.getLong("COL2")); + Assertions.assertFalse(rs.next()); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + } + } + + @Test + void prepareStatementsUsageJdbcPrepare() throws Exception { + try (var ddl = Ddl.builder() + .database(URI.create("/TEST/PREPARE_DB4")) + .relationalExtension(relationalExtension) + .schemaTemplate(SCHEMA_TEMPLATE) + .build()) { + final var connection = ddl.setSchemaAndGetConnection(); + + try (var stmt = connection.createStatement()) { + stmt.execute("INSERT INTO T1 VALUES (1, 10, 1)"); + stmt.execute("INSERT INTO T1 VALUES (2, 20, 2)"); + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var ps = connection.prepareStatement("select * from t1 where col1 = ?")) { + ps.setLong(1, 20); + try (RelationalResultSet rs = ps.executeQuery()) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(2, rs.getLong("ID")); + Assertions.assertFalse(rs.next()); + } + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + + try (var ps = connection.prepareStatement("select * from t1 where id = ?")) { + ps.setLong(1, 2); + try (RelationalResultSet rs = ps.executeQuery()) { + Assertions.assertTrue(rs.next()); + Assertions.assertEquals(2, rs.getLong("ID")); + Assertions.assertEquals(20, rs.getLong("COL1")); + Assertions.assertEquals(2, rs.getLong("COL2")); + Assertions.assertFalse(rs.next()); + } + } + Assertions.assertEquals(2, countCachedPlans(connection, ddl.getSchemaTemplateName())); + } + } +} diff --git a/yaml-tests/src/test/java/YamlIntegrationTests.java b/yaml-tests/src/test/java/YamlIntegrationTests.java index 1b4042d060..79699872ae 100644 --- a/yaml-tests/src/test/java/YamlIntegrationTests.java +++ b/yaml-tests/src/test/java/YamlIntegrationTests.java @@ -362,6 +362,11 @@ public void standardTests(YamlTest.Runner runner) throws Exception { runner.runYamsql("standard-tests.yamsql"); } + @TestTemplate + public void schemaTemplatePrepare(YamlTest.Runner runner) throws Exception { + runner.runYamsql("schema-template-prepare.yamsql"); + } + @TestTemplate public void standardTestsWithMetaData(YamlTest.Runner runner) throws Exception { runner.runYamsql("standard-tests-metadata.yamsql"); diff --git a/yaml-tests/src/test/resources/schema-template-prepare.yamsql b/yaml-tests/src/test/resources/schema-template-prepare.yamsql new file mode 100644 index 0000000000..49af5348b7 --- /dev/null +++ b/yaml-tests/src/test/resources/schema-template-prepare.yamsql @@ -0,0 +1,48 @@ +# +# schema-template-prepare.yamsql +# +# This source file is part of the FoundationDB open source project +# +# Copyright 2021-2026 Apple Inc. and the FoundationDB project authors +# +# 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. + +--- +options: + supported_version: !current_version +--- +schema_template: + create table t1(id bigint, col1 bigint, col2 bigint, primary key(id)) + create index i1 as select col1 from t1 + prepare by_col1 from 'select * from t1 where col1 = 1' + prepare by_id from 'select * from t1 where id = 1' +--- +setup: + steps: + - query: INSERT INTO T1 + VALUES (1, 10, 1), + (2, 10, 2), + (3, 20, 3), + (4, 20, 4), + (5, 30, 5) +--- +test_block: + name: schema-template-prepare-tests + tests: + - + - query: select * from T1 where col1 = 10 + - result: [{ID: 1, COL1: 10, COL2: 1}, + {ID: 2, COL1: 10, COL2: 2}] + - + - query: select * from T1 where id = 3 + - result: [{ID: 3, COL1: 20, COL2: 3}] \ No newline at end of file