Repository: googleapis/java-spanner-jdbc
Pull request #1579 (2024-04-21 06:43:11)

feat: support statement tags in hints
The JDBC driver now supports setting statement tags and RPC priority as a statement hint.

Usage:
* `@{STATEMENT_TAG='my-tag'} SELECT 1;` (GoogleSQL)
* `/*@STATEMENT_TAG='my-tag'*/ SELECT 1;` (PostgreSQL)

Also, RPC priority can be set using a statement hint:
* `@{RPC_PRIORITY=PRIORITY_LOW} SELECT 1;` (GoogleSQL)
* `/*@RPC_PRIORITY=PRIORITY_LOW*/ SELECT 1;` (PostgreSQL)



Patch:
diff --git a/pom.xml b/pom.xml
index 48ce6da..644108c 100644
--- a/pom.xml
+++ b/pom.xml
@@ -171,6 +171,12 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.testcontainers</groupId>
+      <artifactId>testcontainers</artifactId>
+      <version>1.19.7</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.google.truth</groupId>
       <artifactId>truth</artifactId>
diff --git a/samples/install-without-bom/pom.xml b/samples/install-without-bom/pom.xml
index a2537a7..0be1c77 100644
--- a/samples/install-without-bom/pom.xml
+++ b/samples/install-without-bom/pom.xml
@@ -29,7 +29,7 @@
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-spanner-jdbc</artifactId>
-      <version>2.16.1</version>
+      <version>2.16.3</version>
     </dependency>
     <!-- [END spanner-jdbc_install_without_bom] -->
 


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/ClientSideStatementHintsTest.java b/src/test/java/com/google/cloud/spanner/jdbc/ClientSideStatementHintsTest.java
new file mode 100644
index 0000000..9a69ba2
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ClientSideStatementHintsTest.java
@@ -0,0 +1,114 @@
+/*
+ * Copyright 2024 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.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.RequestOptions.Priority;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public class ClientSideStatementHintsTest extends AbstractMockServerTest {
+
+  @Parameter public Dialect dialect;
+
+  private Dialect currentDialect;
+
+  @Parameters(name = "dialect = {0}")
+  public static Object[] data() {
+    return Dialect.values();
+  }
+
+  @Before
+  public void setupDialect() {
+    if (this.dialect != currentDialect) {
+      mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+      this.currentDialect = dialect;
+    }
+  }
+
+  @After
+  public void clearRequests() {
+    mockSpanner.clearRequests();
+  }
+
+  private String createUrl() {
+    return String.format(
+        "jdbc:cloudspanner://localhost:%d/projects/%s/instances/%s/databases/%s?usePlainText=true",
+        getPort(), "proj", "inst", "db" + (dialect == Dialect.POSTGRESQL ? "pg" : ""));
+  }
+
+  private Connection createConnection() throws SQLException {
+    return DriverManager.getConnection(createUrl());
+  }
+
+  @Test
+  public void testStatementTagInHint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      try (ResultSet resultSet =
+          connection
+              .createStatement()
+              .executeQuery(
+                  dialect == Dialect.POSTGRESQL
+                      ? "/*@statement_tag='test-tag'*/SELECT 1"
+                      : "@{statement_tag='test-tag'}SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals("test-tag", request.getRequestOptions().getRequestTag());
+  }
+
+  @Test
+  public void testRpcPriorityInHint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      try (ResultSet resultSet =
+          connection
+              .createStatement()
+              .executeQuery(
+                  dialect == Dialect.POSTGRESQL
+                      ? "/*@rpc_priority=PRIORITY_LOW*/SELECT 1"
+                      : "@{rpc_priority=PRIORITY_LOW}SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals(Priority.PRIORITY_LOW, request.getRequestOptions().getPriority());
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/ConcurrentTransactionOnEmulatorTest.java b/src/test/java/com/google/cloud/spanner/jdbc/ConcurrentTransactionOnEmulatorTest.java
new file mode 100644
index 0000000..e9931e6
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/jdbc/ConcurrentTransactionOnEmulatorTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2024 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.jdbc;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
+
+import com.google.cloud.spanner.connection.SpannerPool;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Properties;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.testcontainers.DockerClientFactory;
+import org.testcontainers.containers.GenericContainer;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.utility.DockerImageName;
+
+@RunWith(JUnit4.class)
+public class ConcurrentTransactionOnEmulatorTest {
+  private static GenericContainer<?> emulator;
+
+  private static Properties properties;
+
+  @BeforeClass
+  public static void startEmulator() {
+    assumeTrue(DockerClientFactory.instance().isDockerAvailable());
+
+    emulator =
+        new GenericContainer<>(
+                DockerImageName.parse("gcr.io/cloud-spanner-emulator/emulator:latest"))
+            .withExposedPorts(9010)
+            .waitingFor(Wait.forListeningPort());
+    emulator.start();
+    properties = new Properties();
+    properties.setProperty("autoConfigEmulator", "true");
+    properties.setProperty(
+        "endpoint", String.format("%s:%d", emulator.getHost(), emulator.getMappedPort(9010)));
+  }
+
+  @AfterClass
+  public static void cleanup() {
+    SpannerPool.closeSpannerPool();
+    if (emulator != null) {
+      emulator.stop();
+    }
+  }
+
+  @Test
+  public void testRunConcurrentTransactions() throws SQLException {
+    String connectionUrl =
+        String.format(
+            "jdbc:cloudspanner:/projects/%s/instances/%s/databases/%s",
+            "emulator-project", "test-instance", "test-database");
+    try (Connection connection1 = DriverManager.getConnection(connectionUrl, properties);
+        Connection connection2 = DriverManager.getConnection(connectionUrl, properties)) {
+      // Create a test table.
+      connection1
+          .createStatement()
+          .execute("create table test (id int64, value string(max)) primary key (id)");
+
+      // Put both connections into autoCommit=false mode and verify that both connections can run
+      // a transaction using a single thread.
+      connection1.setAutoCommit(false);
+      connection2.setAutoCommit(false);
+
+      connection1.createStatement().executeUpdate("insert into test (id, value) values (1, 'One')");
+      connection2.createStatement().executeUpdate("insert into test (id, value) values (2, 'Two')");
+      connection1.commit();
+      connection2.commit();
+
+      // Verify that both transactions succeeded.
+      connection1.setAutoCommit(true);
+      try (ResultSet resultSet =
+          connection1.createStatement().executeQuery("select * from test order by id")) {
+        assertTrue(resultSet.next());
+        assertEquals("One", resultSet.getString("value"));
+        assertTrue(resultSet.next());
+        assertEquals("Two", resultSet.getString("value"));
+        assertFalse(resultSet.next());
+      }
+    }
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/TagMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/TagMockServerTest.java
index 5c626c0..8401297 100644
--- a/src/test/java/com/google/cloud/spanner/jdbc/TagMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/jdbc/TagMockServerTest.java
@@ -18,6 +18,7 @@ package com.google.cloud.spanner.jdbc;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.spanner.Dialect;
@@ -40,8 +41,14 @@ import org.junit.runners.Parameterized.Parameters;
 
 @RunWith(Parameterized.class)
 public class TagMockServerTest extends AbstractMockServerTest {
+  private static final String SELECT_RANDOM_SQL = SELECT_RANDOM_STATEMENT.getSql();
+
+  private static final String INSERT_SQL = INSERT_STATEMENT.getSql();
+
   @Parameter public Dialect dialect;
 
+  private Dialect currentDialect;
+
   @Parameters(name = "dialect = {0}")
   public static Object[] data() {
     return Dialect.values();
@@ -49,7 +56,10 @@ public class TagMockServerTest extends AbstractMockServerTest {
 
   @Before
   public void setupDialect() {
-    mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+    if (this.dialect != currentDialect) {
+      mockSpanner.putStatementResult(StatementResult.detectDialectResult(this.dialect));
+      this.currentDialect = dialect;
+    }
   }
 
   @After
@@ -77,8 +87,7 @@ public class TagMockServerTest extends AbstractMockServerTest {
       connection
           .createStatement()
           .execute(String.format("set %sstatement_tag='my-tag'", getVariablePrefix()));
-      try (ResultSet resultSet =
-          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
         assertTrue(resultSet.next());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -87,8 +96,7 @@ public class TagMockServerTest extends AbstractMockServerTest {
 
       // Verify that the tag is cleared after having been used.
       mockSpanner.clearRequests();
-      try (ResultSet resultSet =
-          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
         assertTrue(resultSet.next());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -103,8 +111,7 @@ public class TagMockServerTest extends AbstractMockServerTest {
       connection
           .createStatement()
           .execute(String.format("set %stransaction_tag='my-tag'", getVariablePrefix()));
-      try (ResultSet resultSet =
-          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
         assertTrue(resultSet.next());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -113,8 +120,7 @@ public class TagMockServerTest extends AbstractMockServerTest {
 
       // Verify that the tag is used for the entire transaction.
       mockSpanner.clearRequests();
-      try (ResultSet resultSet =
-          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
         assertTrue(resultSet.next());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -125,8 +131,7 @@ public class TagMockServerTest extends AbstractMockServerTest {
       connection.commit();
 
       mockSpanner.clearRequests();
-      try (ResultSet resultSet =
-          connection.createStatement().executeQuery(SELECT_RANDOM_STATEMENT.getSql())) {
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT_RANDOM_SQL)) {
         assertTrue(resultSet.next());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
@@ -143,8 +148,8 @@ public class TagMockServerTest extends AbstractMockServerTest {
           .execute(String.format("set %sstatement_tag='my-tag'", getVariablePrefix()));
 
       try (Statement statement = connection.createStatement()) {
-        statement.addBatch(INSERT_STATEMENT.getSql());
-        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_SQL);
+        statement.addBatch(INSERT_SQL);
         assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
@@ -155,8 +160,8 @@ public class TagMockServerTest extends AbstractMockServerTest {
       // Verify that the tag is cleared after having been used.
       mockSpanner.clearRequests();
       try (Statement statement = connection.createStatement()) {
-        statement.addBatch(INSERT_STATEMENT.getSql());
-        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_SQL);
+        statement.addBatch(INSERT_SQL);
         assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
@@ -173,8 +178,8 @@ public class TagMockServerTest extends AbstractMockServerTest {
           .execute(String.format("set %stransaction_tag='my-tag'", getVariablePrefix()));
 
       try (Statement statement = connection.createStatement()) {
-        statement.addBatch(INSERT_STATEMENT.getSql());
-        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_SQL);
+        statement.addBatch(INSERT_SQL);
         assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
@@ -185,8 +190,8 @@ public class TagMockServerTest extends AbstractMockServerTest {
       // Verify that the tag is used for the entire transaction.
       mockSpanner.clearRequests();
       try (Statement statement = connection.createStatement()) {
-        statement.addBatch(INSERT_STATEMENT.getSql());
-        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_SQL);
+        statement.addBatch(INSERT_SQL);
         assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
@@ -197,8 +202,8 @@ public class TagMockServerTest extends AbstractMockServerTest {
       connection.commit();
       mockSpanner.clearRequests();
       try (Statement statement = connection.createStatement()) {
-        statement.addBatch(INSERT_STATEMENT.getSql());
-        statement.addBatch(INSERT_STATEMENT.getSql());
+        statement.addBatch(INSERT_SQL);
+        statement.addBatch(INSERT_SQL);
         assertArrayEquals(new int[] {1, 1}, statement.executeBatch());
       }
       assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
@@ -206,4 +211,24 @@ public class TagMockServerTest extends AbstractMockServerTest {
       assertEquals("", request.getRequestOptions().getTransactionTag());
     }
   }
+
+  @Test
+  public void testStatementTagInHint() throws SQLException {
+    try (Connection connection = createConnection()) {
+      try (ResultSet resultSet =
+          connection
+              .createStatement()
+              .executeQuery(
+                  dialect == Dialect.POSTGRESQL
+                      ? "/*@statement_tag='test-tag'*/SELECT 1"
+                      : "@{statement_tag='test-tag'}SELECT 1")) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+    }
+    assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+    assertEquals("test-tag", request.getRequestOptions().getRequestTag());
+  }
 }