Merge pull request #331 from alex268/master

Added transaction closing on stream cancelation
ydb-platform · Oct 15, 2024 · 6161fb8 · 6161fb8
2 parents 7722d88 + 5ded04d
commit 6161fb8
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java
@@ -375,6 +375,14 @@ void handleCompletion(Status status, Throwable th) {
                                         + status, Issue.Severity.ERROR)));
                     }
                 }
+
+                @Override
+                public void cancel() {
+                    super.cancel();
+                    if (txId.compareAndSet(currentId, null)) {
+                        logger.warn("{} transaction with id {} was cancelled", SessionImpl.this, currentId);
+                    }
+                }
             };
         }
 

diff --git a/query/src/test/java/tech/ydb/query/impl/QueryIntegrationTest.java b/query/src/test/java/tech/ydb/query/impl/QueryIntegrationTest.java
@@ -217,6 +217,26 @@ public void testCancelStream() {
 
             try (QuerySession s4 = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
                 Assert.assertNotEquals(id, s4.getId());
+                id = s4.getId();
+
+                QueryTransaction tx = s4.beginTransaction(TxMode.SERIALIZABLE_RW).join().getValue();
+                Assert.assertTrue(tx.isActive());
+
+                final QueryStream query = tx.createQuery("SELECT 2 + 2;");
+                final CompletableFuture<Void> stop = new CompletableFuture<>();
+                CompletableFuture<Result<QueryInfo>> future = query.execute(part -> {
+                    stop.join();
+                    printQuerySetPart(part);
+                });
+                query.cancel();
+                stop.complete(null);
+                Result<QueryInfo> result = future.join();
+                Assert.assertEquals(StatusCode.CLIENT_CANCELLED, result.getStatus().getCode());
+                Assert.assertFalse(tx.isActive());
+            }
+
+            try (QuerySession s5 = client.createSession(Duration.ofSeconds(5)).join().getValue()) {
+                Assert.assertNotEquals(id, s5.getId());
             }
         }
     }