diff --git a/.kokoro/nightly/integration.cfg b/.kokoro/nightly/integration.cfg index 9640e74d453..5a95c68284c 100644 --- a/.kokoro/nightly/integration.cfg +++ b/.kokoro/nightly/integration.cfg @@ -13,12 +13,12 @@ env_vars: { # TODO: remove this after we've migrated all tests and scripts env_vars: { key: "GCLOUD_PROJECT" - value: "cloud-java-ci-sample" + value: "java-docs-samples-testing" } env_vars: { key: "GOOGLE_CLOUD_PROJECT" - value: "cloud-java-ci-sample" + value: "java-docs-samples-testing" } env_vars: { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java index d5c95d0a5a5..091f26941a5 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java @@ -333,6 +333,20 @@ static final class TagOption extends InternalOption implements ReadQueryUpdateTr void appendToOptions(Options options) { options.tag = tag; } + + @Override + public int hashCode() { + return Objects.hash(this.tag); + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof TagOption)) { + return false; + } + TagOption other = (TagOption) o; + return Objects.equals(this.tag, other.tag); + } } static final class EtagOption extends InternalOption implements DeleteAdminApiOption { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java index 9793a50c636..ca8081a516e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/AbstractStatementParser.java @@ -19,6 +19,8 @@ import com.google.api.core.InternalApi; import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.Options; +import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.Statement; @@ -620,6 +622,29 @@ public String removeCommentsAndTrim(String sql) { /** Removes any statement hints at the beginning of the statement. */ abstract String removeStatementHint(String sql); + @VisibleForTesting + static final ReadQueryUpdateTransactionOption[] EMPTY_OPTIONS = + new ReadQueryUpdateTransactionOption[0]; + + /** + * Extracts any query/update options from the SQL string. Currently, this only supports extracting + * a statement tag, and the statement tag must be given as a statement hint in a comment at the + * start of the query string. + */ + ReadQueryUpdateTransactionOption[] extractOptions(ParsedStatement statement) { + final String statementTagPrefix = "/*@{STATEMENT_TAG="; + + String sql = statement.getStatement().getSql(); + if (sql.startsWith(statementTagPrefix)) { + int endIndex = sql.indexOf("}*/", statementTagPrefix.length()); + if (endIndex > -1) { + String tag = sql.substring(statementTagPrefix.length(), endIndex); + return new ReadQueryUpdateTransactionOption[] {Options.tag(tag)}; + } + } + return EMPTY_OPTIONS; + } + /** Parameter information with positional parameters translated to named parameters. */ @InternalApi public static class ParametersInfo { 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 8ff367b2486..a8f807dbb0c 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 @@ -32,6 +32,7 @@ import com.google.cloud.spanner.Mutation; import com.google.cloud.spanner.Options; import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption; import com.google.cloud.spanner.Options.RpcPriority; import com.google.cloud.spanner.Options.UpdateOption; import com.google.cloud.spanner.PartitionOptions; @@ -1138,6 +1139,8 @@ public ResultSet partitionQuery( "Only queries can be partitioned. Invalid statement: " + query.getSql()); } + QueryOption[] combinedOptions = + concat(getStatementParser().extractOptions(parsedStatement), options); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); return get( transaction.partitionQueryAsync( @@ -1145,7 +1148,8 @@ public ResultSet partitionQuery( parsedStatement, getEffectivePartitionOptions(partitionOptions), mergeDataBoost( - mergeQueryRequestOptions(parsedStatement, mergeQueryStatementTag(options))))); + mergeQueryRequestOptions( + parsedStatement, mergeQueryStatementTag(combinedOptions))))); } private PartitionOptions getEffectivePartitionOptions( @@ -1439,6 +1443,34 @@ private List parseUpdateStatements(Iterable updates) return parsedStatements; } + private UpdateOption[] concat( + ReadQueryUpdateTransactionOption[] statementOptions, UpdateOption[] argumentOptions) { + if (statementOptions == null || statementOptions.length == 0) { + return argumentOptions; + } + if (argumentOptions == null || argumentOptions.length == 0) { + return statementOptions; + } + UpdateOption[] result = + Arrays.copyOf(statementOptions, statementOptions.length + argumentOptions.length); + System.arraycopy(argumentOptions, 0, result, statementOptions.length, argumentOptions.length); + return result; + } + + private QueryOption[] concat( + ReadQueryUpdateTransactionOption[] statementOptions, QueryOption[] argumentOptions) { + if (statementOptions == null || statementOptions.length == 0) { + return argumentOptions; + } + if (argumentOptions == null || argumentOptions.length == 0) { + return statementOptions; + } + QueryOption[] result = + Arrays.copyOf(statementOptions, statementOptions.length + argumentOptions.length); + System.arraycopy(argumentOptions, 0, result, statementOptions.length, argumentOptions.length); + return result; + } + private QueryOption[] mergeDataBoost(QueryOption... options) { if (this.dataBoostEnabled) { options = appendQueryOption(options, Options.dataBoostEnabled(true)); @@ -1515,19 +1547,20 @@ private ResultSet internalExecuteQuery( && (analyzeMode != AnalyzeMode.NONE || statement.hasReturningClause())), "Statement must either be a query or a DML mode with analyzeMode!=NONE or returning clause"); boolean isInternalMetadataQuery = isInternalMetadataQuery(options); + QueryOption[] combinedOptions = concat(getStatementParser().extractOptions(statement), options); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery); if (autoPartitionMode && statement.getType() == StatementType.QUERY && !isInternalMetadataQuery) { return runPartitionedQuery( - statement.getStatement(), PartitionOptions.getDefaultInstance(), options); + statement.getStatement(), PartitionOptions.getDefaultInstance(), combinedOptions); } return get( transaction.executeQueryAsync( callType, statement, analyzeMode, - mergeQueryRequestOptions(statement, mergeQueryStatementTag(options)))); + mergeQueryRequestOptions(statement, mergeQueryStatementTag(combinedOptions)))); } private AsyncResultSet internalExecuteQueryAsync( @@ -1542,44 +1575,54 @@ private AsyncResultSet internalExecuteQueryAsync( ConnectionPreconditions.checkState( !(autoPartitionMode && statement.getType() == StatementType.QUERY), "Partitioned queries cannot be executed asynchronously"); - UnitOfWork transaction = - getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery(options)); + boolean isInternalMetadataQuery = isInternalMetadataQuery(options); + QueryOption[] combinedOptions = concat(getStatementParser().extractOptions(statement), options); + UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(isInternalMetadataQuery); return ResultSets.toAsyncResultSet( transaction.executeQueryAsync( callType, statement, analyzeMode, - mergeQueryRequestOptions(statement, mergeQueryStatementTag(options))), + mergeQueryRequestOptions(statement, mergeQueryStatementTag(combinedOptions))), spanner.getAsyncExecutorProvider(), - options); + combinedOptions); } private ApiFuture internalExecuteUpdateAsync( - final CallType callType, final ParsedStatement update, UpdateOption... options) { + final CallType callType, final ParsedStatement update, final UpdateOption... options) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); + UpdateOption[] combinedOptions = concat(getStatementParser().extractOptions(update), options); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); return transaction.executeUpdateAsync( - callType, update, mergeUpdateRequestOptions(mergeUpdateStatementTag(options))); + callType, update, mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions))); } private ApiFuture internalAnalyzeUpdateAsync( final CallType callType, final ParsedStatement update, - AnalyzeMode analyzeMode, - UpdateOption... options) { + final AnalyzeMode analyzeMode, + final UpdateOption... options) { Preconditions.checkArgument( update.getType() == StatementType.UPDATE, "Statement must be an update"); + UpdateOption[] combinedOptions = concat(getStatementParser().extractOptions(update), options); UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); return transaction.analyzeUpdateAsync( - callType, update, analyzeMode, mergeUpdateRequestOptions(mergeUpdateStatementTag(options))); + callType, + update, + analyzeMode, + mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions))); } private ApiFuture internalExecuteBatchUpdateAsync( - CallType callType, List updates, UpdateOption... options) { + final CallType callType, final List updates, final UpdateOption... options) { UnitOfWork transaction = getCurrentUnitOfWorkOrStartNewUnitOfWork(); + UpdateOption[] combinedOptions = + concat( + getStatementParser().extractOptions(updates.isEmpty() ? null : updates.get(0)), + options); return transaction.executeBatchUpdateAsync( - callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(options))); + callType, updates, mergeUpdateRequestOptions(mergeUpdateStatementTag(combinedOptions))); } private UnitOfWork getCurrentUnitOfWorkOrStartNewUnitOfWork() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java index 3739aa11064..f537e4921a3 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/StatementParserTest.java @@ -16,8 +16,10 @@ package com.google.cloud.spanner.connection; +import static com.google.cloud.spanner.connection.AbstractStatementParser.EMPTY_OPTIONS; import static com.google.common.truth.Truth.assertThat; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotSame; @@ -28,6 +30,8 @@ import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.ErrorCode; +import com.google.cloud.spanner.Options; +import com.google.cloud.spanner.Options.ReadQueryUpdateTransactionOption; import com.google.cloud.spanner.SpannerException; import com.google.cloud.spanner.Statement; import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement; @@ -1690,6 +1694,30 @@ public void testStatementCache_ParameterizedStatement() { assertEquals(1, stats.hitCount()); } + @Test + public void testExtractOptions() { + assertArrayEquals(EMPTY_OPTIONS, parser.extractOptions(parser.parse(Statement.of("select 1")))); + + assertArrayEquals( + new ReadQueryUpdateTransactionOption[] {Options.tag("foo")}, + parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=foo}*/ select 1")))); + assertArrayEquals( + new ReadQueryUpdateTransactionOption[] {Options.tag("tag with space")}, + parser.extractOptions( + parser.parse(Statement.of("/*@{STATEMENT_TAG=tag with space}*/ select 1")))); + assertArrayEquals( + new ReadQueryUpdateTransactionOption[] {Options.tag("foo}")}, + parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=foo}}*/ select 1")))); + + assertArrayEquals( + EMPTY_OPTIONS, + parser.extractOptions(parser.parse(Statement.of("/*@{STATEMENT_TAG=not_a_tag*/select 1")))); + assertArrayEquals( + EMPTY_OPTIONS, + parser.extractOptions( + parser.parse(Statement.of("/*@{STATEMENT_TAG=not_a_tag} */select 1")))); + } + static void assertUnclosedLiteral(AbstractStatementParser parser, String sql) { SpannerException exception = assertThrows(