Skip to content

Commit ebb17fc

Browse files
olavloitegcf-owl-bot[bot]rajatbhatta
authored
fix: always allow metadata queries (#2580)
* 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 <[email protected]> --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com> Co-authored-by: Rajat Bhatta <[email protected]>
1 parent 850435b commit ebb17fc

File tree

6 files changed

+205
-72
lines changed

6 files changed

+205
-72
lines changed

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

Lines changed: 34 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -133,17 +133,6 @@ enum BatchMode {
133133
DML
134134
}
135135

136-
/**
137-
* This query option is used internally to indicate that a query is executed by the library itself
138-
* to fetch metadata. These queries are specifically allowed to be executed even when a DDL batch
139-
* is active.
140-
*/
141-
static final class InternalMetadataQuery implements QueryOption {
142-
static final InternalMetadataQuery INSTANCE = new InternalMetadataQuery();
143-
144-
private InternalMetadataQuery() {}
145-
}
146-
147136
/** The combination of all transaction modes and batch modes. */
148137
enum UnitOfWorkType {
149138
READ_ONLY_TRANSACTION {
@@ -1219,6 +1208,18 @@ private AsyncResultSet parseAndExecuteQueryAsync(
12191208
+ parsedStatement.getSqlWithoutComments());
12201209
}
12211210

1211+
private boolean isInternalMetadataQuery(QueryOption... options) {
1212+
if (options == null) {
1213+
return false;
1214+
}
1215+
for (QueryOption option : options) {
1216+
if (option instanceof InternalMetadataQuery) {
1217+
return true;
1218+
}
1219+
}
1220+
return false;
1221+
}
1222+
12221223
@Override
12231224
public long executeUpdate(Statement update) {
12241225
Preconditions.checkNotNull(update);
@@ -1450,8 +1451,11 @@ private ResultSet internalExecuteQuery(
14501451
|| (statement.getType() == StatementType.UPDATE
14511452
&& (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())),
14521453
"Statement must either be a query or a DML mode with analyzeMode!=NONE or returning clause");
1453-
UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
1454-
if (autoPartitionMode && statement.getType() == StatementType.QUERY) {
1454+
boolean isInternalMetadataQuery = isInternalMetadataQuery(options);
1455+
UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery);
1456+
if (autoPartitionMode
1457+
&& statement.getType() == StatementType.QUERY
1458+
&& !isInternalMetadataQuery) {
14551459
return runPartitionedQuery(
14561460
statement.getStatement(), PartitionOptions.getDefaultInstance(), options);
14571461
}
@@ -1475,7 +1479,8 @@ private AsyncResultSet internalExecuteQueryAsync(
14751479
ConnectionPreconditions.checkState(
14761480
!(autoPartitionMode && statement.getType() == StatementType.QUERY),
14771481
"Partitioned queries cannot be executed asynchronously");
1478-
UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork();
1482+
UnitOfWork transaction =
1483+
getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery(options));
14791484
return ResultSets.toAsyncResultSet(
14801485
transaction.executeQueryAsync(
14811486
callType,
@@ -1514,22 +1519,31 @@ private ApiFuture<long[]> internalExecuteBatchUpdateAsync(
15141519
callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(options)));
15151520
}
15161521

1522+
private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
1523+
return getCurrentUnitOfWorkOrStartNewUnitOfWork(false);
1524+
}
1525+
15171526
/**
15181527
* Returns the current {@link UnitOfWork} of this connection, or creates a new one based on the
15191528
* current transaction settings of the connection and returns that.
15201529
*/
15211530
@VisibleForTesting
1522-
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
1531+
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) {
1532+
if (isInternalMetadataQuery) {
1533+
// Just return a temporary single-use transaction.
1534+
return createNewUnitOfWork(true);
1535+
}
15231536
if (this.currentUnitOfWork == null || !this.currentUnitOfWork.isActive()) {
1524-
this.currentUnitOfWork = createNewUnitOfWork();
1537+
this.currentUnitOfWork = createNewUnitOfWork(false);
15251538
}
15261539
return this.currentUnitOfWork;
15271540
}
15281541

15291542
@VisibleForTesting
1530-
UnitOfWork createNewUnitOfWork() {
1531-
if (isAutocommit() && !isInTransaction() && !isInBatch()) {
1543+
UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery) {
1544+
if (isInternalMetadataQuery || (isAutocommit() && !isInTransaction() && !isInBatch())) {
15321545
return SingleUseTransaction.newBuilder()
1546+
.setInternalMetadataQuery(isInternalMetadataQuery)
15331547
.setDdlClient(ddlClient)
15341548
.setDatabaseClient(dbClient)
15351549
.setBatchClient(batchClient)
@@ -1660,7 +1674,7 @@ public void startBatchDdl() {
16601674
!transactionBeginMarked, "Cannot start a DDL batch when a transaction has begun");
16611675
this.batchMode = BatchMode.DDL;
16621676
this.unitOfWorkType = UnitOfWorkType.DDL_BATCH;
1663-
this.currentUnitOfWork = createNewUnitOfWork();
1677+
this.currentUnitOfWork = createNewUnitOfWork(false);
16641678
}
16651679

16661680
@Override
@@ -1678,7 +1692,7 @@ public void startBatchDml() {
16781692
// Then create the DML batch.
16791693
this.batchMode = BatchMode.DML;
16801694
this.unitOfWorkType = UnitOfWorkType.DML_BATCH;
1681-
this.currentUnitOfWork = createNewUnitOfWork();
1695+
this.currentUnitOfWork = createNewUnitOfWork(false);
16821696
}
16831697

16841698
@Override

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

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -33,15 +33,12 @@
3333
import com.google.cloud.spanner.SpannerExceptionFactory;
3434
import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
3535
import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
36-
import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery;
3736
import com.google.common.annotations.VisibleForTesting;
3837
import com.google.common.base.Preconditions;
3938
import com.google.spanner.admin.database.v1.DatabaseAdminGrpc;
4039
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
41-
import com.google.spanner.v1.SpannerGrpc;
4240
import java.util.ArrayList;
4341
import java.util.Arrays;
44-
import java.util.Collections;
4542
import java.util.List;
4643
import java.util.concurrent.Callable;
4744

@@ -121,29 +118,6 @@ public ApiFuture<ResultSet> executeQueryAsync(
121118
final ParsedStatement statement,
122119
AnalyzeMode analyzeMode,
123120
QueryOption... options) {
124-
if (options != null) {
125-
for (int i = 0; i < options.length; i++) {
126-
if (options[i] instanceof InternalMetadataQuery) {
127-
Preconditions.checkNotNull(statement);
128-
Preconditions.checkArgument(statement.isQuery(), "Statement is not a query");
129-
Preconditions.checkArgument(
130-
analyzeMode == AnalyzeMode.NONE, "Analyze is not allowed for DDL batch");
131-
// Queries marked with internal metadata queries are allowed during a DDL batch.
132-
// These can only be generated by library internal methods and may be used to check
133-
// whether a database object such as table or an index exists.
134-
List<QueryOption> temp = new ArrayList<>();
135-
Collections.addAll(temp, options);
136-
temp.remove(i);
137-
final QueryOption[] internalOptions = temp.toArray(new QueryOption[0]);
138-
Callable<ResultSet> callable =
139-
() ->
140-
DirectExecuteResultSet.ofResultSet(
141-
dbClient.singleUse().executeQuery(statement.getStatement(), internalOptions));
142-
return executeStatementAsync(
143-
callType, statement, callable, SpannerGrpc.getExecuteStreamingSqlMethod());
144-
}
145-
}
146-
}
147121
// Queries are by default not allowed on DDL batches.
148122
throw SpannerExceptionFactory.newSpannerException(
149123
ErrorCode.FAILED_PRECONDITION, "Executing queries is not allowed for DDL batches.");

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ class SingleUseTransaction extends AbstractBaseUnitOfWork {
7878
private final TimestampBound readOnlyStaleness;
7979
private final AutocommitDmlMode autocommitDmlMode;
8080
private final boolean returnCommitStats;
81+
private final boolean internalMetdataQuery;
8182
private volatile SettableApiFuture<Timestamp> readTimestamp = null;
8283
private volatile TransactionRunner writeTransaction;
8384
private boolean used = false;
@@ -91,6 +92,7 @@ static class Builder extends AbstractBaseUnitOfWork.Builder<Builder, SingleUseTr
9192
private TimestampBound readOnlyStaleness;
9293
private AutocommitDmlMode autocommitDmlMode;
9394
private boolean returnCommitStats;
95+
private boolean internalMetadataQuery;
9496

9597
private Builder() {}
9698

@@ -133,6 +135,11 @@ Builder setReturnCommitStats(boolean returnCommitStats) {
133135
return this;
134136
}
135137

138+
Builder setInternalMetadataQuery(boolean internalMetadataQuery) {
139+
this.internalMetadataQuery = internalMetadataQuery;
140+
return this;
141+
}
142+
136143
@Override
137144
SingleUseTransaction build() {
138145
Preconditions.checkState(ddlClient != null, "No DDL client specified");
@@ -157,6 +164,7 @@ private SingleUseTransaction(Builder builder) {
157164
this.readOnlyStaleness = builder.readOnlyStaleness;
158165
this.autocommitDmlMode = builder.autocommitDmlMode;
159166
this.returnCommitStats = builder.returnCommitStats;
167+
this.internalMetdataQuery = builder.internalMetadataQuery;
160168
}
161169

162170
@Override
@@ -207,8 +215,11 @@ public ApiFuture<ResultSet> executeQueryAsync(
207215
return executeDmlReturningAsync(callType, statement, options);
208216
}
209217

218+
// Do not use a read-only staleness for internal metadata queries.
210219
final ReadOnlyTransaction currentTransaction =
211-
dbClient.singleUseReadOnlyTransaction(readOnlyStaleness);
220+
internalMetdataQuery
221+
? dbClient.singleUseReadOnlyTransaction()
222+
: dbClient.singleUseReadOnlyTransaction(readOnlyStaleness);
212223
Callable<ResultSet> callable =
213224
() -> {
214225
try {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionImplTest.java

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,7 +1389,7 @@ public void testMergeQueryOptions() {
13891389
new ConnectionImpl(
13901390
connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) {
13911391
@Override
1392-
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
1392+
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) {
13931393
return unitOfWork;
13941394
}
13951395
}) {
@@ -1498,7 +1498,7 @@ public void testStatementTagAlwaysAllowed() {
14981498
new ConnectionImpl(
14991499
connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) {
15001500
@Override
1501-
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() {
1501+
UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork(boolean isInternalMetadataQuery) {
15021502
return unitOfWork;
15031503
}
15041504
}) {
@@ -1609,7 +1609,7 @@ public void testTransactionTagNotAllowedAfterTransactionStarted() {
16091609
new ConnectionImpl(
16101610
connectionOptions, spannerPool, ddlClient, dbClient, mock(BatchClient.class)) {
16111611
@Override
1612-
UnitOfWork createNewUnitOfWork() {
1612+
UnitOfWork createNewUnitOfWork(boolean isInternalMetadataQuery) {
16131613
return unitOfWork;
16141614
}
16151615
}) {

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/DdlBatchTest.java

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -40,15 +40,12 @@
4040
import com.google.cloud.spanner.Dialect;
4141
import com.google.cloud.spanner.ErrorCode;
4242
import com.google.cloud.spanner.Mutation;
43-
import com.google.cloud.spanner.ReadContext;
44-
import com.google.cloud.spanner.ResultSet;
4543
import com.google.cloud.spanner.SpannerBatchUpdateException;
4644
import com.google.cloud.spanner.SpannerException;
4745
import com.google.cloud.spanner.SpannerExceptionFactory;
4846
import com.google.cloud.spanner.Statement;
4947
import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
5048
import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
51-
import com.google.cloud.spanner.connection.Connection.InternalMetadataQuery;
5249
import com.google.cloud.spanner.connection.UnitOfWork.CallType;
5350
import com.google.cloud.spanner.connection.UnitOfWork.UnitOfWorkState;
5451
import com.google.protobuf.Timestamp;
@@ -157,25 +154,6 @@ public void testExecuteCreateDatabase() {
157154
.parse(Statement.of("CREATE DATABASE foo"))));
158155
}
159156

160-
@Test
161-
public void testExecuteMetadataQuery() {
162-
Statement statement = Statement.of("SELECT * FROM INFORMATION_SCHEMA.TABLES");
163-
ParsedStatement parsedStatement = mock(ParsedStatement.class);
164-
when(parsedStatement.isQuery()).thenReturn(true);
165-
when(parsedStatement.getStatement()).thenReturn(statement);
166-
DatabaseClient dbClient = mock(DatabaseClient.class);
167-
ReadContext singleUse = mock(ReadContext.class);
168-
ResultSet resultSet = mock(ResultSet.class);
169-
when(singleUse.executeQuery(statement)).thenReturn(resultSet);
170-
when(dbClient.singleUse()).thenReturn(singleUse);
171-
DdlBatch batch = createSubject(createDefaultMockDdlClient(), dbClient);
172-
assertThat(
173-
get(batch.executeQueryAsync(
174-
CallType.SYNC, parsedStatement, AnalyzeMode.NONE, InternalMetadataQuery.INSTANCE))
175-
.hashCode(),
176-
is(equalTo(resultSet.hashCode())));
177-
}
178-
179157
@Test
180158
public void testExecuteUpdate() {
181159
DdlBatch batch = createSubject();

0 commit comments

Comments
 (0)