googleapis · olavloite · Mar 31, 2025 · Mar 31, 2025
@@ -859,6 +859,27 @@ public void setTransactionMode(TransactionMode transactionMode) {
     this.unitOfWorkType = UnitOfWorkType.of(transactionMode);
   }
 
+  IsolationLevel getTransactionIsolationLevel() {
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(!isDdlBatchActive(), "This connection is in a DDL batch");
+    ConnectionPreconditions.checkState(isInTransaction(), "This connection has no transaction");
+    return this.transactionIsolationLevel;
+  }
+
+  void setTransactionIsolationLevel(IsolationLevel isolationLevel) {
+    Preconditions.checkNotNull(isolationLevel);
+    ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);
+    ConnectionPreconditions.checkState(
+        !isBatchActive(), "Cannot set transaction isolation level while in a batch");
+    ConnectionPreconditions.checkState(isInTransaction(), "This connection has no transaction");
+    ConnectionPreconditions.checkState(
+        !isTransactionStarted(),
+        "The transaction isolation level cannot be set after the transaction has started");
+
+    this.transactionBeginMarked = true;
+    this.transactionIsolationLevel = isolationLevel;
+  }
+
   @Override
   public String getTransactionTag() {
     ConnectionPreconditions.checkState(!isClosed(), CLOSED_ERROR_MSG);

@@ -490,6 +490,11 @@ public StatementResult statementSetTransactionMode(TransactionMode mode) {
 
   @Override
   public StatementResult statementSetPgTransactionMode(PgTransactionMode transactionMode) {
+    if (transactionMode.getIsolationLevel() != null) {
+      getConnection()
+          .setTransactionIsolationLevel(
+              transactionMode.getIsolationLevel().getSpannerIsolationLevel());
+    }
     if (transactionMode.getAccessMode() != null) {
       switch (transactionMode.getAccessMode()) {
         case READ_ONLY_TRANSACTION:

@@ -122,12 +122,13 @@ public void testBatchDml() {
   }
 
   @Test
-  public void testTransactionIsolationLevel() {
+  public void testBeginTransactionIsolationLevel() {
+    SpannerPool.closeSpannerPool();
     for (Dialect dialect : new Dialect[] {Dialect.POSTGRESQL, Dialect.GOOGLE_STANDARD_SQL}) {
       mockSpanner.putStatementResult(
           MockSpannerServiceImpl.StatementResult.detectDialectResult(dialect));
 
-      try (Connection connection = createConnection()) {
+      try (Connection connection = super.createConnection()) {
         for (IsolationLevel isolationLevel :
             new IsolationLevel[] {IsolationLevel.REPEATABLE_READ, IsolationLevel.SERIALIZABLE}) {
           for (boolean useSql : new boolean[] {true, false}) {
@@ -158,4 +159,41 @@ public void testTransactionIsolationLevel() {
       SpannerPool.closeSpannerPool();
     }
   }
+
+  @Test
+  public void testSetTransactionIsolationLevel() {
+    SpannerPool.closeSpannerPool();
+    mockSpanner.putStatementResult(
+        MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.POSTGRESQL));
+
+    try (Connection connection = super.createConnection()) {
+      for (boolean autocommit : new boolean[] {true, false}) {
+        connection.setAutocommit(autocommit);
+
+        for (IsolationLevel isolationLevel :
+            new IsolationLevel[] {IsolationLevel.REPEATABLE_READ, IsolationLevel.SERIALIZABLE}) {
+          // Manually start a transaction if autocommit is enabled.
+          if (autocommit) {
+            connection.execute(Statement.of("begin"));
+          }
+          connection.execute(
+              Statement.of(
+                  "set transaction isolation level " + isolationLevel.name().replace("_", " ")));
+          connection.executeUpdate(INSERT_STATEMENT);
+          connection.commit();
+
+          assertEquals(1, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+          ExecuteSqlRequest request = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class).get(0);
+          assertTrue(request.getTransaction().hasBegin());
+          assertTrue(request.getTransaction().getBegin().hasReadWrite());
+          assertEquals(isolationLevel, request.getTransaction().getBegin().getIsolationLevel());
+          assertFalse(request.getLastStatement());
+          assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+
+          mockSpanner.clearRequests();
+        }
+      }
+    }
+    SpannerPool.closeSpannerPool();
+  }
 }