From a4bd82c8e4ce8b7179b943ac06b049598276f1b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Wed, 14 Jul 2021 10:44:28 +0200
Subject: [PATCH] feat: add support for tagging to JDBC connection (#270)

* feat: add support for tagging to JDBC connection

* fix: add ignored differences + default implementations
---
 clirr-ignored-differences.xml                 | 22 ++++++++++
 .../jdbc/CloudSpannerJdbcConnection.java      | 41 +++++++++++++++++++
 .../cloud/spanner/jdbc/JdbcConnection.java    | 32 +++++++++++++++
 3 files changed, 95 insertions(+)

diff --git a/clirr-ignored-differences.xml b/clirr-ignored-differences.xml
index ac9c147d..0962d5e1 100644
--- a/clirr-ignored-differences.xml
+++ b/clirr-ignored-differences.xml
@@ -111,4 +111,26 @@
     <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
     <method>void setReturnCommitStats(boolean)</method>
   </difference>
+
+  <!-- Tagging - These are not really breaking changes, as we provide default implementations. -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>java.lang.String getTransactionTag()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setTransactionTag(java.lang.String)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>java.lang.String getStatementTag()</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection</className>
+    <method>void setStatementTag(java.lang.String)</method>
+  </difference>
 </differences>
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
index 4a5f001a..50ae71c5 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/CloudSpannerJdbcConnection.java
@@ -39,6 +39,47 @@
  */
 public interface CloudSpannerJdbcConnection extends Connection {
 
+  /**
+   * Sets the transaction tag to use for the current transaction. This method may only be called
+   * when in a transaction, and before the transaction is actually started, i.e. before any
+   * statements have been executed in the transaction.
+   *
+   * <p>The tag will be set as the transaction tag of all statements during the transaction, and as
+   * the transaction tag of the commit.
+   *
+   * <p>The transaction tag will automatically be cleared after the transaction has ended.
+   *
+   * @param tag The tag to use.
+   */
+  default void setTransactionTag(String tag) throws SQLException {
+    throw new UnsupportedOperationException();
+  };
+
+  /** @return The transaction tag of the current transaction. */
+  default String getTransactionTag() throws SQLException {
+    throw new UnsupportedOperationException();
+  };
+
+  /**
+   * Sets the statement tag to use for the next statement that will be executed. The tag is
+   * automatically cleared after the statement is executed. Statement tags can be used both with
+   * autocommit=true and autocommit=false, and can be used for partitioned DML.
+   *
+   * @param tag The statement tag to use with the next statement that will be executed on this
+   *     connection.
+   */
+  default void setStatementTag(String tag) throws SQLException {
+    throw new UnsupportedOperationException();
+  };
+
+  /**
+   * @return The statement tag that will be used with the next statement that is executed on this
+   *     connection.
+   */
+  default String getStatementTag() throws SQLException {
+    throw new UnsupportedOperationException();
+  };
+
   /**
    * Sets the transaction mode to use for current transaction. This method may only be called when
    * in a transaction, and before the transaction is actually started, i.e. before any statements
diff --git a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
index bf9120f5..53c1dc09 100644
--- a/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
+++ b/src/main/java/com/google/cloud/spanner/jdbc/JdbcConnection.java
@@ -78,6 +78,38 @@ public String nativeSQL(String sql) throws SQLException {
         .sqlWithNamedParameters;
   }
 
+  @Override
+  public String getStatementTag() throws SQLException {
+    checkClosed();
+    return getSpannerConnection().getStatementTag();
+  }
+
+  @Override
+  public void setStatementTag(String tag) throws SQLException {
+    checkClosed();
+    try {
+      getSpannerConnection().setStatementTag(tag);
+    } catch (SpannerException e) {
+      throw JdbcSqlExceptionFactory.of(e);
+    }
+  }
+
+  @Override
+  public String getTransactionTag() throws SQLException {
+    checkClosed();
+    return getSpannerConnection().getTransactionTag();
+  }
+
+  @Override
+  public void setTransactionTag(String tag) throws SQLException {
+    checkClosed();
+    try {
+      getSpannerConnection().setTransactionTag(tag);
+    } catch (SpannerException e) {
+      throw JdbcSqlExceptionFactory.of(e);
+    }
+  }
+
   @Override
   public void setTransactionMode(TransactionMode mode) throws SQLException {
     checkClosed();