feat: support statement tags (#1157)

* feat: support statement tags

The Spanner JDBC driver now supports adding statement tags
as a query hint. This can be used in the Hibernate dialect
to include statement tags with generated queries.

* docs: add statement tag documentation to sample

README.adoc

@@ -171,7 +171,22 @@ adding them either as a Hibernate query hint, or by adding them as specifically
 These specifically formatted comments are processed by this Hibernate dialect, which then modifies
 the generated query before it is sent to the JDBC driver.
 
-Usage with Hibernate directly using a query hint:
+Simple statement hints that only need to be prepended to a query can be added as if they were a
+comment:
+
+[source,java]
+----
+/** Get all singers that have a last name that starts with the given prefix. */
+@Query("SELECT s FROM Singer s WHERE starts_with(s.lastName, :lastName)=true")
+@QueryHints(
+  @QueryHint(
+      name = AvailableHints.HINT_COMMENT,
+      value = "@{STATEMENT_TAG=search_singers_by_last_name_starts_with}"))
+Stream<Singer> searchByLastNameStartsWith(@Param("lastName") String lastName);
+----
+
+More complex hints that need to be added somewhere in the middle of the statement, such as index
+hints, can be added like this:
 
 [source,java]
 ----
@@ -190,7 +205,7 @@ Query<Singer> query = session.createQuery(cr)
 List<Singer> singers = query.getResultList().size();
 ----
 
-You can also add hints as comments to queries that are generated by JPA:
+You can also add more complex hints as comments to queries that are generated by JPA:
 
 [source,java]
 ----
@@ -209,8 +224,8 @@ You can also add hints as comments to queries that are generated by JPA:
 List<Singer> findByActive(boolean active);
 ----
 
-This https://github.com/GoogleCloudPlatform/google-cloud-spanner-hibernate/blob/-/google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample[working sample application]
-shows how to use the above hint.
+This https://github.com/GoogleCloudPlatform/google-cloud-spanner-hibernate/blob/-/google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/main/java/com/google/cloud/spanner/sample/repository/SingerRepository.java[working sample application]
+shows how to use the above hints.
 
 ==== Transaction Tags
 
@@ -273,6 +288,26 @@ public class VenueService {
 This https://github.com/GoogleCloudPlatform/google-cloud-spanner-hibernate/blob/-/google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample[working sample application]
 shows how to use transaction tags.
 
+==== Statement Tags
+
+NOTE: This feature requires that you use Spanner JDBC driver version 2.16.3 or higher.
+
+Spanner supports adding
+https://cloud.google.com/spanner/docs/introspection/troubleshooting-with-tags[statement tags]
+for troubleshooting queries and transactions. You can add statement tags to your Hibernate or
+Spring Data JPA application by adding a hint to a query.
+
+[source,java]
+----
+/** Get all singers that have a last name that starts with the given prefix. */
+@Query("SELECT s FROM Singer s WHERE starts_with(s.lastName, :lastName)=true")
+@QueryHints(
+  @QueryHint(
+      name = AvailableHints.HINT_COMMENT,
+      value = "@{STATEMENT_TAG=search_singers_by_last_name_starts_with}"))
+Stream<Singer> searchByLastNameStartsWith(@Param("lastName") String lastName);
+----
+
 ==== Custom Spanner Column Types
 
 This project offers the following Hibernate type mappings for specific Spanner column types:

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/SpannerDialect.java

@@ -350,11 +350,15 @@ public SqmMultiTableInsertStrategy getFallbackSqmInsertStrategy(
   @Override
   public String addSqlHintOrComment(
       String sql, QueryOptions queryOptions, boolean commentsEnabled) {
-    if (hasCommentHint(queryOptions)) {
-      sql = applyHint(sql, queryOptions.getComment());
-    }
-    if (queryOptions.getDatabaseHints() != null && !queryOptions.getDatabaseHints().isEmpty()) {
-      sql = applyQueryHints(sql, queryOptions);
+    if (hasStatementHint(queryOptions)) {
+      sql = queryOptions.getComment() + sql;
+    } else {
+      if (hasCommentHint(queryOptions)) {
+        sql = applyHint(sql, queryOptions.getComment());
+      }
+      if (queryOptions.getDatabaseHints() != null && !queryOptions.getDatabaseHints().isEmpty()) {
+        sql = applyQueryHints(sql, queryOptions);
+      }
     }
     return super.addSqlHintOrComment(sql, queryOptions, commentsEnabled);
   }
@@ -390,4 +394,12 @@ private static boolean stringCouldContainReplacementHint(String hint) {
         && hint.contains("}")
         && hint.contains(ReplaceQueryPartsHint.SPANNER_REPLACEMENTS_FIELD_NAME);
   }
+
+  private static boolean hasStatementHint(QueryOptions queryOptions) {
+    return hasStatementHint(queryOptions.getComment());
+  }
+
+  private static boolean hasStatementHint(String hint) {
+    return !Strings.isNullOrEmpty(hint) && hint.startsWith("@{") && hint.endsWith("}");
+  }
 }

google-cloud-spanner-hibernate-dialect/src/main/java/com/google/cloud/spanner/hibernate/hints/Hints.java

@@ -121,6 +121,16 @@ private static ReplaceQueryPartsHint statementHint(String hint, Object value) {
         SELECT_OR_DML, String.format("@{%s=%s}$1", hint, value), ReplaceMode.FIRST);
   }
 
+  /**
+   * Creates a hint that adds @{STATEMENT_TAG=value} to the statement.
+   *
+   * @param tag the statement tag to add to the statement
+   * @return a hint that can be added as a comment or query hint to a Hibernate statement
+   */
+  public static ReplaceQueryPartsHint statementTag(String tag) {
+    return statementHint("STATEMENT_TAG", tag);
+  }
+
   /**
    * Creates a hint that adds @{USE_ADDITIONAL_PARALLELISM=value} to the statement.
    *

google-cloud-spanner-hibernate-dialect/src/test/java/com/google/cloud/spanner/hibernate/hints/HintsTest.java

@@ -171,6 +171,22 @@ public void testForceJoinOrder() {
         Hints.forceJoinOrder("singers", true, ReplaceMode.ALL).toComment());
   }
 
+  @Test
+  public void testStatementTag() {
+    assertEquals(
+        "@{STATEMENT_TAG=my_tag}select * from singers",
+        Hints.statementTag("my_tag").replace("select * from singers"));
+    assertEquals(
+        "@{STATEMENT_TAG=other_tag}SELECT * from singers",
+        Hints.statementTag("other_tag").replace("SELECT * from singers"));
+    assertEquals(
+        "@{STATEMENT_TAG=insert_singer}" + "insert into singers (id, value) SELECT * from singers",
+        Hints.statementTag("insert_singer")
+            .replace("insert into singers (id, value) SELECT * from singers"));
+
+    System.out.println(Hints.statementTag("get_albums_by_title").toComment());
+  }
+
   @Test
   public void testUseAdditionalParallelism() {
     assertEquals(

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/README.adoc

@@ -59,6 +59,10 @@ This sample application showcases how to use the following features:
    `TransactionTagInterceptor` to your Hibernate configuration.
    See link:src/main/java/com/google/cloud/spanner/sample/TaggingHibernatePropertiesCustomizer.java[TaggingHibernatePropertiesCustomizer.java] for how this is done in this sample application.
 
-8. Add statement and query hints to generated queries. See the `findByActive` method in the
+8. Add statement tags to generated queries. See the `searchByLastNameStartsWith` method in
+   link:src/main/java/com/google/cloud/spanner/sample/repository/SingerRepository.java[SingerRepository.java]
+   for an example of how to add a statement tag.
+
+9. Add statement and query hints to generated queries. See the `findByActive` method in the
    link:src/main/java/com/google/cloud/spanner/sample/repository/SingerRepository.java[SingerRepository.java]
    file for an example.

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/main/java/com/google/cloud/spanner/sample/repository/AlbumRepository.java

@@ -19,6 +19,19 @@
 package com.google.cloud.spanner.sample.repository;
 
 import com.google.cloud.spanner.sample.entities.Album;
+import jakarta.persistence.QueryHint;
+import java.util.List;
+import org.hibernate.jpa.AvailableHints;
 import org.springframework.data.jpa.repository.JpaRepository;
+import org.springframework.data.jpa.repository.QueryHints;
 
-public interface AlbumRepository extends JpaRepository<Album, String> {}
+public interface AlbumRepository extends JpaRepository<Album, String> {
+
+  // Statement hints, for example hints for adding a statement tag, can be added directly as a
+  // simple comment.
+  @QueryHints(
+      @QueryHint(
+          name = AvailableHints.HINT_COMMENT,
+          value = "@{STATEMENT_TAG=get_albums_by_title}"))
+  List<Album> getAlbumsByTitle(String title);
+}

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/main/java/com/google/cloud/spanner/sample/repository/SingerRepository.java

@@ -32,6 +32,10 @@ public interface SingerRepository extends JpaRepository<Singer, String> {
 
   /** Get all singers that have a last name that starts with the given prefix. */
   @Query("SELECT s FROM Singer s WHERE starts_with(s.lastName, :lastName)=true")
+  @QueryHints(
+      @QueryHint(
+          name = AvailableHints.HINT_COMMENT,
+          value = "@{STATEMENT_TAG=search_singers_by_last_name_starts_with}"))
   Stream<Singer> searchByLastNameStartsWith(@Param("lastName") String lastName);
 
   // The hint value used here is generated by calling the method:

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/main/java/com/google/cloud/spanner/sample/service/AlbumService.java

@@ -18,6 +18,7 @@
 
 package com.google.cloud.spanner.sample.service;
 
+import com.google.cloud.spanner.hibernate.TransactionTag;
 import com.google.cloud.spanner.sample.entities.Album;
 import com.google.cloud.spanner.sample.entities.Singer;
 import com.google.cloud.spanner.sample.repository.AlbumRepository;
@@ -52,6 +53,10 @@ public AlbumService(
     this.singerRepository = singerRepository;
   }
 
+  public List<Album> getAlbums(String title) {
+    return this.albumRepository.getAlbumsByTitle(title);
+  }
+
   /** Deletes all Album records in the database. */
   @Transactional
   public void deleteAllAlbums() {
@@ -60,6 +65,7 @@ public void deleteAllAlbums() {
 
   /** Generates the specified number of random Album records. */
   @Transactional
+  @TransactionTag("generate_random_albums")
   public List<Album> generateRandomAlbums(int count) {
     Random random = new Random();
 

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/main/java/com/google/cloud/spanner/sample/service/SingerService.java

@@ -18,6 +18,7 @@
 
 package com.google.cloud.spanner.sample.service;
 
+import com.google.cloud.spanner.hibernate.TransactionTag;
 import com.google.cloud.spanner.sample.entities.Album;
 import com.google.cloud.spanner.sample.entities.Concert;
 import com.google.cloud.spanner.sample.entities.Singer;
@@ -93,6 +94,7 @@ public void deleteAllSingers() {
 
   /** Generates the specified number of random singer records. */
   @Transactional
+  @TransactionTag("generate_random_singers")
   public List<Singer> generateRandomSingers(int count) {
     List<Singer> singers = new ArrayList<>(count);
     for (int i = 0; i < count; i++) {

google-cloud-spanner-hibernate-samples/spring-data-jpa-full-sample/src/test/java/com/google/cloud/spanner/sample/SampleApplicationMockServerTest.java

@@ -670,6 +670,13 @@ public static void setupQueryResults() {
                                 .build())
                         .build())
                 .build()));
+
+    // Add results for selecting albums.
+    String selectAlbumsSql =
+        "select a1_0.id,a1_0.cover_picture,a1_0.created_at,a1_0.marketing_budget,a1_0.release_date,a1_0.singer_id,a1_0.title,a1_0.updated_at from album a1_0 where a1_0.title=@p1";
+    mockSpanner.putStatementResult(
+        StatementResult.query(
+            Statement.newBuilder(selectAlbumsSql).bind("p1").to("Foo").build(), empty()));
   }
 
   @Test
@@ -750,7 +757,27 @@ public void testRunApplication() {
                     request
                         .getRequestOptions()
                         .getTransactionTag()
-                        .equals("service_SingerService_generateRandomSingers"))
+                        .equals("generate_random_singers"))
+            .count());
+    assertEquals(
+        1,
+        mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).stream()
+            .filter(
+                request ->
+                    request
+                        .getRequestOptions()
+                        .getTransactionTag()
+                        .equals("generate_random_albums"))
+            .count());
+    assertEquals(
+        3,
+        mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).stream()
+            .filter(
+                request ->
+                    request
+                        .getRequestOptions()
+                        .getRequestTag()
+                        .equals("search_singers_by_last_name_starts_with"))
             .count());
   }