samples: add samples for tagging feature

samples/snippets/src/main/java/com/example/spanner/TransactionWithTagSample.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2021 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.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.Options;
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.TransactionContext;
+import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
+
+/** Sample showing how to add and query tags to Cloud Spanner operations. */
+public class TransactionWithTagSample {
+
+  // [START spanner_set_transaction_tag]
+  static void setTransactionTag(DatabaseClient databaseClient) {
+    // Sets the transaction tag to "app=concert,env=dev".
+    // This transaction tag will be applied to all the individual operations inside this
+    // transaction.
+    databaseClient
+        .readWriteTransaction(Options.tag("app=concert,env=dev"))
+        .run(
+            new TransactionCallable<Void>() {
+              @Override
+              public Void run(TransactionContext transaction) throws Exception {
+                Statement selectStatement =
+                    Statement.newBuilder(
+                        "SELECT VenueId, VenueName, Capacity FROM Venues "
+                            + "WHERE OutdoorVenue = @outdoorVenue")
+                        .bind("outdoorVenue")
+                        .to(false)
+                        .build();
+
+                // Sets the request tag to "app=concert,env=dev,action=select".
+                // This request tag will only be set on this request.
+                ResultSet resultSet =
+                    transaction.executeQuery(
+                        selectStatement, Options.tag("app=concert,env=dev,action=select"));
+
+                while (resultSet.next()) {
+                  long newCapacity = resultSet.getLong(2) / 2;
+                  Statement updateStatement =
+                      Statement.newBuilder(
+                          "UPDATE Venues SET Capacity = @capacity WHERE VenueId = @venueId")
+                          .bind("capacity")
+                          .to(newCapacity)
+                          .bind("venueId")
+                          .to(resultSet.getLong(0))
+                          .build();
+
+                  // Sets the request tag to "app=concert,env=dev,action=update".
+                  // This request tag will only be set on this request.
+                  transaction.executeUpdate(
+                      updateStatement, Options.tag("app=concert,env=dev,action=update"));
+                  System.out.printf(
+                      "Capacity of %s updated to %d\n", resultSet.getString(1), newCapacity);
+                }
+                return null;
+              }
+            });
+  }
+  // [END spanner_set_transaction_tag]
+
+  // [START spanner_set_request_tag]
+  static void setRequestTag(DatabaseClient databaseClient) {
+    // Sets the request tag to "app=concert,env=dev,action=select".
+    // This request tag will only be set on this request.
+    try (ResultSet resultSet =
+        databaseClient
+            .singleUse()
+            .executeQuery(
+                Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"),
+                Options.tag("app=concert,env=dev,action=select"))) {
+      while (resultSet.next()) {
+        System.out.printf(
+            "%d %d %s\n", resultSet.getLong(0), resultSet.getLong(1), resultSet.getString(2));
+      }
+    }
+  }
+  // [END spanner_set_request_tag]
+}

samples/snippets/src/test/java/com/example/spanner/TransactionWithTagSampleIT.java

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2021 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 static com.example.spanner.SampleRunner.runSample;
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.KeySet;
+import com.google.cloud.spanner.Mutation;
+import com.google.cloud.spanner.ResultSet;
+import com.google.cloud.spanner.Statement;
+import com.google.common.collect.ImmutableList;
+import java.util.Arrays;
+import java.util.Collections;
+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.JUnit4;
+
+/** Integration tests for {@link TransactionWithTagSample} */
+@RunWith(JUnit4.class)
+public class TransactionWithTagSampleIT extends SampleTestBase {
+
+  private static DatabaseId databaseId;
+
+  @BeforeClass
+  public static void createTestDatabase() throws Exception {
+    final String database = idGenerator.generateDatabaseId();
+    databaseAdminClient
+        .createDatabase(
+            instanceId,
+            database,
+            ImmutableList.of(
+                "CREATE TABLE Albums ("
+                    + "  SingerId   INT64 NOT NULL,"
+                    + "  AlbumId  INT64,"
+                    + "  AlbumTitle   STRING(1024)"
+                    + ") PRIMARY KEY (SingerId)",
+                "CREATE TABLE Venues ("
+                    + "  VenueId        INT64 NOT NULL,"
+                    + "  VenueName      STRING(MAX),"
+                    + "  Capacity       INT64,"
+                    + "  OutdoorVenue   BOOL"
+                    + ") PRIMARY KEY (VenueId)"))
+        .get();
+    databaseId = DatabaseId.of(projectId, instanceId, database);
+  }
+
+  @Before
+  public void insertTestData() {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+    client.write(
+        Arrays.asList(
+            Mutation.newInsertOrUpdateBuilder("Albums")
+                .set("SingerId")
+                .to(1L)
+                .set("AlbumId")
+                .to(1L)
+                .set("AlbumTitle")
+                .to("title 1")
+                .build(),
+            Mutation.newInsertOrUpdateBuilder("Venues")
+                .set("VenueId")
+                .to(4L)
+                .set("VenueName")
+                .to("name")
+                .set("Capacity")
+                .to(4000000)
+                .set("OutdoorVenue")
+                .to(false)
+                .build()));
+  }
+
+  @After
+  public void removeTestData() {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+    client.write(Collections.singletonList(Mutation.delete("Albums", KeySet.all())));
+    client.write(Collections.singleton(Mutation.delete("Venues", KeySet.all())));
+  }
+
+  @Test
+  public void testSetRequestTag() throws Exception {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+
+    // Execute query multiple times to give a better chance of being part of the TopN tables.
+    for (int i = 0; i < 50; i++) {
+      final String out = runSample(() -> TransactionWithTagSample.setRequestTag(client));
+      assertThat(out).contains("1 1 title 1");
+    }
+
+    boolean queryStatsFound = false;
+    // Query statistics will be collected within 1 minute intervals.
+    for (int i = 0; i < 50; i++) {
+      try (ResultSet resultSet =
+          client
+              .singleUse()
+              .executeQuery(
+                  Statement.of("SELECT text FROM SPANNER_SYS.QUERY_STATS_TOP_MINUTE "
+                      + "Where request_tag=\"app=concert,env=dev,action=select\" limit 1"))) {
+        while (resultSet.next()) {
+          assertThat(resultSet.getString(0))
+              .isEqualTo("SELECT SingerId, AlbumId, AlbumTitle FROM Albums");
+          queryStatsFound = true;
+        }
+        if (queryStatsFound) {
+          break;
+        }
+        Thread.sleep(5000);
+      }
+    }
+    assertTrue(queryStatsFound);
+  }
+
+  @Test
+  public void testSetTransactionTag() throws Exception {
+    final DatabaseClient client = spanner.getDatabaseClient(databaseId);
+    // Execute transaction multiple times to give a better chance of being part of the TopN tables.
+    for (int i = 0; i < 50; i++) {
+      final String out = runSample(() -> TransactionWithTagSample.setTransactionTag(client));
+      assertThat(out).contains("Capacity of name updated to");
+    }
+
+    boolean txnStatsFound = false;
+    // Transaction statistics will be collected within 1 minute intervals.
+    for (int i = 0; i < 50; i++) {
+      try (ResultSet resultSet =
+          client
+              .singleUse()
+              .executeQuery(
+                  Statement.of(
+                      "SELECT READ_COLUMNS FROM SPANNER_SYS.TXN_STATS_TOP_MINUTE where "
+                          + "transaction_tag=\"app=concert,env=dev\"limit 1"))) {
+        while (resultSet.next()) {
+          assertThat(resultSet.getStringList(0)).contains("Venues.Capacity");
+          assertThat(resultSet.getStringList(0)).contains("Venues.VenueName");
+          txnStatsFound = true;
+        }
+        if (txnStatsFound) {
+          break;
+        }
+        Thread.sleep(5000);
+      }
+    }
+    assertTrue(txnStatsFound);
+  }
+}