diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index eee779685e87..a199de202d57 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -644,6 +644,7 @@ private void initTransactionInternal(BeginTransactionRequest request) { private final DirectedReadOptions defaultDirectedReadOptions; private final DecodeMode defaultDecodeMode; private final Clock clock; + private volatile String cachedRequestTag; @GuardedBy("lock") private boolean isValid = true; @@ -841,7 +842,22 @@ RequestOptions buildRequestOptions(Options options) { builder.setClientContext(clientContextBuilder.build()); } if (getTransactionTag() != null) { - builder.setTransactionTag(getTransactionTag()); + // Read-write transactions support transaction-level tags only. We populate the + // transaction tag on the builder if it is non-empty. + if (!getTransactionTag().isEmpty()) { + builder.setTransactionTag(getTransactionTag()); + } + } else if (session.getSpanner().getOptions().isAutoTaggingEnabled() + && builder.getRequestTag().isEmpty()) { + // Read-only contexts (both single-use and multi-use) do not support transaction-level tags. + // We lazily resolve and populate the request tag instead. + if (this.cachedRequestTag == null) { + String autoTag = AutoTagHelper.getAutoTag(session.getSpanner().getOptions()); + this.cachedRequestTag = autoTag == null ? "" : autoTag; + } + if (!this.cachedRequestTag.isEmpty()) { + builder.setRequestTag(this.cachedRequestTag); + } } return builder.build(); } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AutoTagHelper.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AutoTagHelper.java new file mode 100644 index 000000000000..9a9c2447f851 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AutoTagHelper.java @@ -0,0 +1,95 @@ +/* + * Copyright 2026 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; + +import java.util.List; + +/** Helper for Spanner transaction tags. */ +final class AutoTagHelper { + + /** Maximum allowed character length for resolved tags. */ + private static final int MAX_TAG_LENGTH = 50; + + /** Ignored packages. */ + private static final String[] INTERNAL_PACKAGES; + + static { + INTERNAL_PACKAGES = + new String[] { + "java.", + "javax.", + "jdk.", + "sun.", + "io.grpc.", + "com.google.cloud.spanner.", + "com.google.api." + }; + } + + private AutoTagHelper() { + // prevent instantiation + } + + static String getAutoTag(final SpannerOptions options) { + StackTraceElement[] stackTrace = new Throwable().getStackTrace(); + int tracerLimit = options.getAutoTaggingTracerLimit(); + int limit = Math.min(stackTrace.length, tracerLimit); + List targetPackages = options.getAutoTaggingPackages(); + boolean hasTarget = targetPackages != null && !targetPackages.isEmpty(); + + for (int i = 0; i < limit; i++) { + StackTraceElement element = stackTrace[i]; + String className = element.getClassName(); + if (hasTarget) { + for (String targetPackage : targetPackages) { + if (className.startsWith(targetPackage)) { + return formatTag(className, element.getMethodName()); + } + } + } else if (isInternalPackage(className)) { + continue; + } else { + return formatTag(className, element.getMethodName()); + } + } + return null; + } + + private static boolean isInternalPackage(final String cls) { + for (String internalPackage : INTERNAL_PACKAGES) { + if (cls.startsWith(internalPackage)) { + return true; + } + } + return false; + } + + private static String formatTag(final String cls, final String method) { + int lastDot = cls.lastIndexOf('.'); + String simpleClassName; + if (lastDot == -1) { + simpleClassName = cls; + } else { + simpleClassName = cls.substring(lastDot + 1); + } + String tag = simpleClassName + "." + method; + if (tag.length() > MAX_TAG_LENGTH) { + tag = tag.substring(tag.length() - MAX_TAG_LENGTH); + } + return tag; + } +} diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java index 9c542bc52365..3b4190a237c7 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java @@ -96,6 +96,7 @@ import java.time.Duration; import java.util.ArrayList; import java.util.Base64; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -308,6 +309,9 @@ static GcpChannelPoolOptions mergeWithDefaultChannelPoolOptions( private final String monitoringHost; private final TransactionOptions defaultTransactionOptions; private final RequestOptions.ClientContext clientContext; + private final boolean autoTaggingEnabled; + private final List autoTaggingPackages; + private final int autoTaggingTracerLimit; enum TracingFramework { OPEN_CENSUS, @@ -993,6 +997,9 @@ protected SpannerOptions(Builder builder) { monitoringHost = builder.monitoringHost; defaultTransactionOptions = builder.defaultTransactionOptions; clientContext = builder.clientContext; + autoTaggingEnabled = builder.autoTaggingEnabled; + autoTaggingPackages = builder.autoTaggingPackages; + autoTaggingTracerLimit = builder.autoTaggingTracerLimit; } private String getResolvedUniverseDomain() { @@ -1064,6 +1071,14 @@ default boolean isEnableLocationApi() { return false; } + default boolean isAutoTaggingDisabled() { + return false; + } + + default boolean isAutoTaggingEnabled() { + return false; + } + @Deprecated @ObsoleteApi( "This will be removed in an upcoming version without a major version bump. You should use" @@ -1168,6 +1183,18 @@ public boolean isEnableLocationApi() { return Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR)); } + @Override + public boolean isAutoTaggingDisabled() { + return Boolean.parseBoolean(System.getenv("SPANNER_DISABLE_AUTO_TAGGING")) + || Boolean.parseBoolean(System.getProperty("spanner.disable_auto_tagging")); + } + + @Override + public boolean isAutoTaggingEnabled() { + return Boolean.parseBoolean(System.getenv("SPANNER_ENABLE_AUTO_TAGGING")) + || Boolean.getBoolean("spanner.enable_auto_tagging"); + } + @Override public String getMonitoringHost() { return System.getenv(SPANNER_MONITORING_HOST); @@ -1256,6 +1283,9 @@ public static class Builder private boolean usePlainText = false; private TransactionOptions defaultTransactionOptions = TransactionOptions.getDefaultInstance(); private RequestOptions.ClientContext clientContext; + private boolean autoTaggingEnabled = false; + private List autoTaggingPackages = Collections.emptyList(); + private int autoTaggingTracerLimit = 50; private static String createCustomClientLibToken(String token) { return token + " " + ServiceOptions.getGoogApiClientLibName(); @@ -1362,6 +1392,9 @@ protected Builder() { this.monitoringHost = options.monitoringHost; this.defaultTransactionOptions = options.defaultTransactionOptions; this.clientContext = options.clientContext; + this.autoTaggingEnabled = options.autoTaggingEnabled; + this.autoTaggingPackages = options.autoTaggingPackages; + this.autoTaggingTracerLimit = options.autoTaggingTracerLimit; } @Override @@ -2120,6 +2153,36 @@ public Builder setDefaultClientContext(RequestOptions.ClientContext clientContex return this; } + public Builder enableAutoTagging() { + this.autoTaggingEnabled = true; + return this; + } + + public Builder disableAutoTagging() { + this.autoTaggingEnabled = false; + return this; + } + + public Builder setAutoTaggingPackages(String... autoTaggingPackages) { + this.autoTaggingPackages = + Collections.unmodifiableList( + new ArrayList<>( + java.util.Arrays.asList(Preconditions.checkNotNull(autoTaggingPackages)))); + return this; + } + + public Builder setAutoTaggingPackages(List autoTaggingPackages) { + this.autoTaggingPackages = + Collections.unmodifiableList( + new ArrayList<>(Preconditions.checkNotNull(autoTaggingPackages))); + return this; + } + + public Builder setAutoTaggingTracerLimit(int autoTaggingTracerLimit) { + this.autoTaggingTracerLimit = autoTaggingTracerLimit; + return this; + } + @SuppressWarnings("rawtypes") @Override public SpannerOptions build() { @@ -2547,6 +2610,25 @@ public TransactionOptions getDefaultTransactionOptions() { return defaultTransactionOptions; } + public boolean isAutoTaggingEnabled() { + if (environment.isAutoTaggingDisabled()) { + return false; + } + return autoTaggingEnabled || environment.isAutoTaggingEnabled(); + } + + public List getAutoTaggingPackages() { + return autoTaggingPackages; + } + + public int getAutoTaggingTracerLimit() { + return autoTaggingTracerLimit; + } + + public boolean isAutoTaggingDisabled() { + return environment.isAutoTaggingDisabled(); + } + @BetaApi public boolean isUseVirtualThreads() { return useVirtualThreads; diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index f1d34cd81cdf..4ffe87104d9b 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -198,6 +198,7 @@ public void removeListener(Runnable listener) { private boolean aborted; private final Options options; + private volatile String cachedTransactionTag; /** Default to -1 to indicate not available. */ @GuardedBy("lock") @@ -780,6 +781,15 @@ String getTransactionTag() { if (this.options.hasTag()) { return this.options.tag(); } + if (session.getSpanner().getOptions().isAutoTaggingEnabled()) { + if (this.cachedTransactionTag == null) { + this.cachedTransactionTag = AutoTagHelper.getAutoTag(session.getSpanner().getOptions()); + if (this.cachedTransactionTag == null) { + this.cachedTransactionTag = ""; + } + } + return this.cachedTransactionTag; + } return null; } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/example/spanner/TagTestHelper.java b/java-spanner/google-cloud-spanner/src/test/java/com/example/spanner/TagTestHelper.java new file mode 100644 index 000000000000..fa122f001286 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/example/spanner/TagTestHelper.java @@ -0,0 +1,105 @@ +/* + * Copyright 2026 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.example.spanner; + +import com.google.api.core.ApiFuture; +import com.google.cloud.spanner.AbortedException; +import com.google.cloud.spanner.AsyncTransactionManager; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.Options.TransactionOption; +import com.google.cloud.spanner.TransactionContext; +import com.google.cloud.spanner.TransactionManager; +import com.google.cloud.spanner.TransactionRunner; +import com.google.common.util.concurrent.MoreExecutors; + +/** Helper class in a non-ignored package to stable-test auto-tagging stack walks. */ +public final class TagTestHelper { + + private TagTestHelper() {} + + public static T run( + DatabaseClient client, TransactionRunner.TransactionCallable callable) { + return client.readWriteTransaction().run(callable); + } + + public static T runWithOptions( + DatabaseClient client, + TransactionRunner.TransactionCallable callable, + TransactionOption option) { + return client.readWriteTransaction(option).run(callable); + } + + public static ApiFuture runAsync( + DatabaseClient client, com.google.cloud.spanner.AsyncRunner.AsyncWork work) { + return client.runAsync().runAsync(work, MoreExecutors.directExecutor()); + } + + public static void singleUseConsume( + DatabaseClient client, com.google.cloud.spanner.Statement stmt) { + try (com.google.cloud.spanner.ResultSet resultSet = client.singleUse().executeQuery(stmt)) { + while (resultSet.next()) {} + } + } + + public static void readOnlyTxnConsume( + com.google.cloud.spanner.ReadOnlyTransaction txn, com.google.cloud.spanner.Statement stmt) { + try (com.google.cloud.spanner.ResultSet resultSet = txn.executeQuery(stmt)) { + while (resultSet.next()) {} + } + } + + public static void runWithManager( + DatabaseClient client, com.google.cloud.spanner.Statement stmt) { + try (TransactionManager manager = client.transactionManager()) { + TransactionContext transaction = manager.begin(); + while (true) { + try { + transaction.executeUpdate(stmt); + manager.commit(); + break; + } catch (AbortedException abortedException) { + transaction = manager.resetForRetry(); + } + } + } + } + + public static ApiFuture runWithAsyncManager( + DatabaseClient client, com.google.cloud.spanner.Statement stmt) { + final AsyncTransactionManager manager = client.transactionManagerAsync(); + AsyncTransactionManager.TransactionContextFuture txnFuture = manager.beginAsync(); + return com.google.api.core.ApiFutures.transformAsync( + txnFuture + .then( + (transaction, ignored) -> transaction.executeUpdateAsync(stmt), + MoreExecutors.directExecutor()) + .commitAsync(), + commitTimestamp -> { + manager.close(); + return com.google.api.core.ApiFutures.immediateFuture(null); + }, + MoreExecutors.directExecutor()); + } + + /** Subclass with an extremely long name to test tag truncation rules. */ + public static final class ExtremelyLongClassNameForTestingTagTruncationSupportUnderCheckstyle { + public static T run( + DatabaseClient client, TransactionRunner.TransactionCallable callable) { + return client.readWriteTransaction().run(callable); + } + } +} diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockAutoTaggingTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockAutoTaggingTest.java new file mode 100644 index 000000000000..1506e5d9a078 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockAutoTaggingTest.java @@ -0,0 +1,341 @@ +/* + * Copyright 2026 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; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import com.example.spanner.TagTestHelper; +import com.google.cloud.NoCredentials; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class MockAutoTaggingTest extends AbstractMockServerTest { + + @Before + public void setUpProperties() { + System.clearProperty("spanner.disable_auto_tagging"); + System.clearProperty("spanner.enable_auto_tagging"); + } + + @After + public void tearDownProperties() { + System.clearProperty("spanner.disable_auto_tagging"); + System.clearProperty("spanner.enable_auto_tagging"); + } + + private Spanner createSpanner(boolean enableAutoTag, String targetPackage) { + SpannerOptions.Builder builder = + SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelProvider(channelProvider) + .setCredentials(NoCredentials.getInstance()); + if (enableAutoTag) { + builder.enableAutoTagging(); + } + if (targetPackage != null) { + builder.setAutoTaggingPackages(targetPackage); + } + builder.setAutoTaggingTracerLimit(200); + builder.setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setWaitForMinSessionsDuration(java.time.Duration.ofSeconds(10)) + .build()); + return builder.build().getService(); + } + + @Test + public void testReadWriteTransactionAutoTagging() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.run( + databaseClient, + transaction -> { + transaction.executeUpdate(Statement.of("UPDATE Venues SET Capacity = 100")); + return null; + }); + + // Verify transaction tag populated on ExecuteSqlRequest + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + String transactionTag = sqlRequest.getRequestOptions().getTransactionTag(); + assertEquals("TagTestHelper.run", transactionTag); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + + // Verify transaction tag matches on CommitRequest + assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class)); + CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).get(0); + assertEquals(transactionTag, commitRequest.getRequestOptions().getTransactionTag()); + } + } + + @Test + public void testReadWriteTransactionAutoTaggingWithExplicitTag() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.runWithOptions( + databaseClient, + transaction -> { + transaction.executeUpdate(Statement.of("UPDATE Venues SET Capacity = 100")); + return null; + }, + Options.tag("my-explicit-tag")); + + // Verify explicit tag is preserved and not overwritten by auto-tag + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals("my-explicit-tag", sqlRequest.getRequestOptions().getTransactionTag()); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } + + @Test + public void testReadWriteTransactionAutoTaggingTruncation() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.ExtremelyLongClassNameForTestingTagTruncationSupportUnderCheckstyle.run( + databaseClient, + transaction -> { + transaction.executeUpdate(Statement.of("UPDATE Venues SET Capacity = 100")); + return null; + }); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + String transactionTag = sqlRequest.getRequestOptions().getTransactionTag(); + assertNotNull(transactionTag); + assertTrue(transactionTag.length() <= 50); + assertTrue(transactionTag.endsWith("uncationSupportUnderCheckstyle.run")); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } + + @Test + public void testGlobalDisableOverride() { + System.setProperty("spanner.disable_auto_tagging", "true"); + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.run( + databaseClient, + transaction -> { + transaction.executeUpdate(Statement.of("UPDATE Venues SET Capacity = 100")); + return null; + }); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals("", sqlRequest.getRequestOptions().getTransactionTag()); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } + + @Test + public void testGlobalEnableOverride() { + System.setProperty("spanner.enable_auto_tagging", "true"); + try (Spanner spannerInstance = createSpanner(false, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.run( + databaseClient, + transaction -> { + transaction.executeUpdate(Statement.of("UPDATE Venues SET Capacity = 100")); + return null; + }); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals("TagTestHelper.run", sqlRequest.getRequestOptions().getTransactionTag()); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } + + @Test + public void testReadWriteTransactionAsyncAutoTagging() throws Exception { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.runAsync( + databaseClient, + transaction -> + transaction.executeUpdateAsync(Statement.of("UPDATE Venues SET Capacity = 100"))) + .get(); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + String transactionTag = sqlRequest.getRequestOptions().getTransactionTag(); + assertEquals("TagTestHelper.runAsync", transactionTag); + } + } + + @Test + public void testSingleUseQueryAutoTagging() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + com.google.spanner.v1.ResultSet select1ResultSet = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + com.google.protobuf.ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata( + com.google.spanner.v1.ResultSetMetadata.newBuilder() + .setRowType( + com.google.spanner.v1.StructType.newBuilder() + .addFields( + com.google.spanner.v1.StructType.Field.newBuilder() + .setName("COL1") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(com.google.spanner.v1.TypeCode.INT64) + .build()) + .build()) + .build()) + .build()) + .build(); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.query( + Statement.of("SELECT * FROM Albums"), select1ResultSet)); + + TagTestHelper.singleUseConsume(databaseClient, Statement.of("SELECT * FROM Albums")); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals( + "TagTestHelper.singleUseConsume", sqlRequest.getRequestOptions().getRequestTag()); + assertEquals("", sqlRequest.getRequestOptions().getTransactionTag()); + } + } + + @Test + public void testMultiUseReadOnlyTransactionAutoTagging() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + com.google.spanner.v1.ResultSet select1ResultSet = + com.google.spanner.v1.ResultSet.newBuilder() + .addRows( + com.google.protobuf.ListValue.newBuilder() + .addValues(com.google.protobuf.Value.newBuilder().setStringValue("1").build()) + .build()) + .setMetadata( + com.google.spanner.v1.ResultSetMetadata.newBuilder() + .setRowType( + com.google.spanner.v1.StructType.newBuilder() + .addFields( + com.google.spanner.v1.StructType.Field.newBuilder() + .setName("COL1") + .setType( + com.google.spanner.v1.Type.newBuilder() + .setCode(com.google.spanner.v1.TypeCode.INT64) + .build()) + .build()) + .build()) + .build()) + .build(); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.query( + Statement.of("SELECT * FROM Albums"), select1ResultSet)); + + try (ReadOnlyTransaction transaction = databaseClient.readOnlyTransaction()) { + TagTestHelper.readOnlyTxnConsume(transaction, Statement.of("SELECT * FROM Albums")); + } + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals( + "TagTestHelper.readOnlyTxnConsume", sqlRequest.getRequestOptions().getRequestTag()); + assertEquals("", sqlRequest.getRequestOptions().getTransactionTag()); + } + } + + @Test + public void testTransactionManagerAutoTagging() { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.runWithManager( + databaseClient, Statement.of("UPDATE Venues SET Capacity = 100")); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + assertEquals( + "TagTestHelper.runWithManager", sqlRequest.getRequestOptions().getTransactionTag()); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } + + @Test + public void testAsyncTransactionManagerAutoTagging() throws Exception { + try (Spanner spannerInstance = createSpanner(true, "com.example.spanner")) { + DatabaseClient databaseClient = + spannerInstance.getDatabaseClient(DatabaseId.of("proj", "inst", "db")); + mockSpanner.putStatementResult( + MockSpannerServiceImpl.StatementResult.update( + Statement.of("UPDATE Venues SET Capacity = 100"), 1L)); + + TagTestHelper.runWithAsyncManager( + databaseClient, Statement.of("UPDATE Venues SET Capacity = 100")) + .get(); + + assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); + ExecuteSqlRequest sqlRequest = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0); + String transactionTag = sqlRequest.getRequestOptions().getTransactionTag(); + assertTrue(transactionTag.contains("runWithAsyncManager")); + assertEquals("", sqlRequest.getRequestOptions().getRequestTag()); + } + } +}