From ebb17fc8aeac5fc75e4f135f33dba970f2480585 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Fri, 11 Aug 2023 19:07:05 +0200 Subject: [PATCH] fix: always allow metadata queries (#2580) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: always allow metadata queries Internal metadata queries should always be allowed, regardless of the type of transaction that is currently running or any special mode that has been set on the connection. Internal metadata queries are system queries that are generated by a connection to return metadata to an application, for example methods like getTables() in JDBC. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * chore: run code formatter * chore: make if condition more readable Co-authored-by: Rajat Bhatta <93644539+rajatbhatta@users.noreply.github.com> --------- Co-authored-by: Owl Bot Co-authored-by: Rajat Bhatta <93644539+rajatbhatta@users.noreply.github.com> --- .../spanner/connection/ConnectionImpl.java | 54 +++--- .../cloud/spanner/connection/DdlBatch.java | 26 --- .../connection/SingleUseTransaction.java | 13 +- .../connection/ConnectionImplTest.java | 6 +- .../spanner/connection/DdlBatchTest.java | 22 --- .../InternalMetadataQueryMockServerTest.java | 156 ++++++++++++++++++ 6 files changed, 205 insertions(+), 72 deletions(-) create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/InternalMetadataQueryMockServerTest.java 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)); + } +}