diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java index f2230e59ba..b5d4cdf8a2 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionImpl.java @@ -133,17 +133,6 @@ enum BatchMode { DML } - /** - * This query option is used internally to indicate that a query is executed by the library itself - * to fetch metadata. These queries are specifically allowed to be executed even when a DDL batch - * is active. - */ - static final class InternalMetadataQuery implements QueryOption { - static final InternalMetadataQuery INSTANCE = new InternalMetadataQuery(); - - private InternalMetadataQuery() {} - } - /** The combination of all transaction modes and batch modes. */ enum UnitOfWorkType { READ_ONLY_TRANSACTION { @@ -1219,6 +1208,18 @@ private AsyncResultSet parseAndExecuteQueryAsync( + parsedStatement.getSqlWithoutComments()); } + private boolean isInternalMetadataQuery(QueryOption... options) { + if (options == null) { + return false; + } + for (QueryOption option : options) { + if (option instanceof InternalMetadataQuery) { + return true; + } + } + return false; + } + @Override public long executeUpdate(Statement update) { Preconditions.checkNotNull(update); @@ -1450,8 +1451,11 @@ private ResultSet internalExecuteQuery( || (statement.getType() == StatementType.UPDATE && (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())), "Statement must either be a query or a DML mode with analyzeMode!=NONE or returning clause"); - UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); - if (autoPartitionMode && statement.getType() == StatementType.QUERY) { + boolean isInternalMetadataQuery = isInternalMetadataQuery(options); + UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery); + if (autoPartitionMode + && statement.getType() == StatementType.QUERY + && !isInternalMetadataQuery) { return runPartitionedQuery( statement.getStatement(), PartitionOptions.getDefaultInstance(), options); } @@ -1475,7 +1479,8 @@ private AsyncResultSet internalExecuteQueryAsync( ConnectionPreconditions.checkState( !(autoPartitionMode && statement.getType() == StatementType.QUERY), "Partitioned queries cannot be executed asynchronously"); - UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); + UnitOfWork transaction = + getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery(options)); return ResultSets.toAsyncResultSet( transaction.executeQueryAsync( callType, @@ -1514,22 +1519,31 @@ private ApiFuture internalExecuteBatchUpdateAsync( callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(options))); } + private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { + return getCurrentUnitOfWorkOrStartNewUnitOfWork(false); + } + /** * Returns the current {@link UnitOfWork} of this connection, or creates a new one based on the * current transaction settings of the connection and returns that. */ @VisibleForTesting - UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { + UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) { + if (isInternalMetadataQuery) { + // Just return a temporary single-use transaction. + return createNewUnitOfWork(true); + } if (this.currentUnitOfWork == null || !this.currentUnitOfWork.isActive()) { - this.currentUnitOfWork = createNewUnitOfWork(); + this.currentUnitOfWork = createNewUnitOfWork(false); } return this.currentUnitOfWork; } @VisibleForTesting - UnitOfWork createNewUnitOfWork() { - if (isAutocommit() && !isInTransaction() && !isInBatch()) { + UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery) { + if (isInternalMetadataQuery || (isAutocommit() && !isInTransaction() && !isInBatch())) { return SingleUseTransaction.newBuilder() + .setInternalMetadataQuery(isInternalMetadataQuery) .setDdlClient(ddlClient) .setDatabaseClient(dbClient) .setBatchClient(batchClient) @@ -1660,7 +1674,7 @@ public void startBatchDdl() { !transactionBeginMarked, "Cannot start a DDL batch when a transaction has begun"); this.batchMode = BatchMode.DDL; this.unitOfWorkType = UnitOfWorkType.DDL_BATCH; - this.currentUnitOfWork = createNewUnitOfWork(); + this.currentUnitOfWork = createNewUnitOfWork(false); } @Override @@ -1678,7 +1692,7 @@ public void startBatchDml() { // Then create the DML batch. this.batchMode = BatchMode.DML; this.unitOfWorkType = UnitOfWorkType.DML_BATCH; - this.currentUnitOfWork = createNewUnitOfWork(); + this.currentUnitOfWork = createNewUnitOfWork(false); } @Override diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java index ebfd2bc454..55b780c571 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java @@ -33,15 +33,12 @@ import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType; -import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.spanner.admin.database.v1.DatabaseAdminGrpc; import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; -import com.google.spanner.v1.SpannerGrpc; import java.util.ArrayList; import java.util.Arrays; -import java.util.Collections; import java.util.List; import java.util.concurrent.Callable; @@ -121,29 +118,6 @@ public ApiFuture executeQueryAsync( final ParsedStatement statement, AnalyzeMode analyzeMode, QueryOption... options) { - if (options != null) { - for (int i = 0; i < options.length; i++) { - if (options[i] instanceof InternalMetadataQuery) { - Preconditions.checkNotNull(statement); - Preconditions.checkArgument(statement.isQuery(), "Statement is not a query"); - Preconditions.checkArgument( - analyzeMode == AnalyzeMode.NONE, "Analyze is not allowed for DDL batch"); - // Queries marked with internal metadata queries are allowed during a DDL batch. - // These can only be generated by library internal methods and may be used to check - // whether a database object such as table or an index exists. - List temp = new ArrayList<>(); - Collections.addAll(temp, options); - temp.remove(i); - final QueryOption[] internalOptions = temp.toArray(new QueryOption[0]); - Callable callable = - () -> - DirectExecuteResultSet.ofResultSet( - dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions)); - return executeStatementAsync( - callType, statement, callable, SpannerGrpc.getExecuteStreamingSqlMethod()); - } - } - } // Queries are by default not allowed on DDL batches. throw SpannerExceptionFactory.newSpannerException( ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DDL batches."); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java index ea2159b00a..52486ed43e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SingleUseTransaction.java @@ -78,6 +78,7 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork { private final TimestampBound readOnlyStaleness; private final AutocommitDmlMode autocommitDmlMode; private final boolean returnCommitStats; + private final boolean internalMetdataQuery; private volatile SettableApiFuture readTimestamp = null; private volatile TransactionRunner writeTransaction; private boolean used = false; @@ -91,6 +92,7 @@ static class Builder extends AbstractBaseUnitOfWork.Builder executeQueryAsync( return executeDmlReturningAsync(callType, statement, options); } + // Do not use a read-only staleness for internal metadata queries. final ReadOnlyTransaction currentTransaction = - dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); + internalMetdataQuery + ? dbClient.singleUseReadOnlyTransaction() + : dbClient.singleUseReadOnlyTransaction(readOnlyStaleness); Callable callable = () -> { try { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java index d767c9f517..09d2f92a6d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java @@ -1389,7 +1389,7 @@ public void testMergeQueryOptions() { new ConnectionImpl( connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) { @Override - UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { + UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) { return unitOfWork; } }) { @@ -1498,7 +1498,7 @@ public void testStatementTagAlwaysAllowed() { new ConnectionImpl( connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) { @Override - UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { + UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) { return unitOfWork; } }) { @@ -1609,7 +1609,7 @@ public void testTransactionTagNotAllowedAfterTransactionStarted() { new ConnectionImpl( connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) { @Override - UnitOfWork createNewUnitOfWork() { + UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery) { return unitOfWork; } }) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java index 699ab23f3a..5ef4c5291d 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java @@ -40,15 +40,12 @@ import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.ErrorCode; import com.google.cloud.spanner.Mutation; -import com.google.cloud.spanner.ReadContext; -import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.SpannerBatchUpdateException; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType; -import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery; import com.google.cloud.spanner.connection.UnitOfWork.CallType; import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState; import com.google.protobuf.Timestamp; @@ -157,25 +154,6 @@ public void testExecuteCreateDatabase() { .parse(Statement.of("CREATE DATABASE foo")))); } - @Test - public void testExecuteMetadataQuery() { - Statement statement = Statement.of("SELECT * FROM INFORMATION_SCHEMA.TABLES"); - ParsedStatement parsedStatement = mock(ParsedStatement.class); - when(parsedStatement.isQuery()).thenReturn(true); - when(parsedStatement.getStatement()).thenReturn(statement); - DatabaseClient dbClient = mock(DatabaseClient.class); - ReadContext singleUse = mock(ReadContext.class); - ResultSet resultSet = mock(ResultSet.class); - when(singleUse.executeQuery(statement)).thenReturn(resultSet); - when(dbClient.singleUse()).thenReturn(singleUse); - DdlBatch batch = createSubject(createDefaultMockDdlClient(), dbClient); - assertThat( - get(batch.executeQueryAsync( - CallType.SYNC, parsedStatement, AnalyzeMode.NONE, InternalMetadataQuery.INSTANCE)) - .hashCode(), - is(equalTo(resultSet.hashCode()))); - } - @Test public void testExecuteUpdate() { DdlBatch batch = createSubject(); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/InternalMetadataQueryMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/InternalMetadataQueryMockServerTest.java new file mode 100644 index 0000000000..89c4dd0b03 --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/InternalMetadataQueryMockServerTest.java @@ -0,0 +1,156 @@ +/* + * Copyright 2023 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spanner.TimestampBound; +import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery; +import com.google.protobuf.ByteString; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.PartitionQueryRequest; +import java.util.concurrent.TimeUnit; +import org.junit.After; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +@RunWith(Parameterized.class) +public class InternalMetadataQueryMockServerTest extends AbstractMockServerTest { + private static final Statement STATEMENT = + Statement.of("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES"); + + @Parameters(name = "dialect = {0}") + public static Object[] data() { + return Dialect.values(); + } + + @Parameter public Dialect dialect; + + @BeforeClass + public static void setupInternalMetadataQueryResults() { + mockSpanner.putStatementResult( + StatementResult.query(STATEMENT, SELECT_COUNT_RESULTSET_BEFORE_INSERT)); + } + + @Before + public void setupDialect() { + mockSpanner.putStatementResult(StatementResult.detectDialectResult(dialect)); + } + + @After + public void clearRequests() { + mockSpanner.clearRequests(); + } + + @Test + public void testInternalMetadataQueryInAutocommit() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + verifyInternalMetadataQuery(connection); + } + } + + @Test + public void testInternalMetadataQueryWithStaleness() { + try (Connection connection = createConnection()) { + connection.setAutocommit(true); + connection.setReadOnlyStaleness(TimestampBound.ofMaxStaleness(10L, TimeUnit.SECONDS)); + verifyInternalMetadataQuery(connection); + } + } + + @Test + public void testInternalMetadataQueryReadOnlyTransaction() { + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + connection.setReadOnly(true); + verifyInternalMetadataQuery(connection); + } + } + + @Test + public void testInternalMetadataQueryReadOnlyTransactionWithStaleness() { + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + connection.setReadOnly(true); + connection.setReadOnlyStaleness(TimestampBound.ofExactStaleness(10L, TimeUnit.SECONDS)); + verifyInternalMetadataQuery(connection); + } + } + + @Test + public void testInternalMetadataQueryReadWriteTransaction() { + try (Connection connection = createConnection()) { + connection.setAutocommit(false); + connection.setReadOnly(false); + verifyInternalMetadataQuery(connection); + } + } + + @Test + public void testInternalMetadataQueryInDmlBatch() { + try (Connection connection = createConnection()) { + connection.startBatchDml(); + verifyInternalMetadataQuery(connection); + connection.runBatch(); + } + } + + @Test + public void testInternalMetadataQueryInDdlBatch() { + try (Connection connection = createConnection()) { + connection.startBatchDdl(); + verifyInternalMetadataQuery(connection); + connection.runBatch(); + } + } + + @Test + public void testInternalMetadataQueryInAutoPartitionMode() { + try (Connection connection = createConnection()) { + connection.setAutoPartitionMode(true); + verifyInternalMetadataQuery(connection); + } + } + + private void verifyInternalMetadataQuery(Connection connection) { + try (ResultSet resultSet = connection.executeQuery(STATEMENT, InternalMetadataQuery.INSTANCE)) { + assertTrue(resultSet.next()); + assertEquals(0L, resultSet.getLong(0)); + assertFalse(resultSet.next()); + } + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertTrue(request.hasTransaction()); + assertTrue(request.getTransaction().hasSingleUse()); + assertTrue(request.getTransaction().getSingleUse().hasReadOnly()); + assertTrue(request.getTransaction().getSingleUse().getReadOnly().hasStrong()); + assertEquals(ByteString.EMPTY, request.getPartitionToken()); + assertEquals(0, mockSpanner.countRequestsOfType(PartitionQueryRequest.class)); + } +}