|
29 | 29 | import java.util.Properties; |
30 | 30 | import java.util.UUID; |
31 | 31 | import java.util.concurrent.CountDownLatch; |
| 32 | +import java.util.concurrent.TimeUnit; |
| 33 | +import java.util.concurrent.atomic.AtomicReference; |
32 | 34 |
|
33 | 35 | import static org.testng.Assert.assertEquals; |
34 | 36 | import static org.testng.Assert.assertFalse; |
@@ -711,6 +713,133 @@ public void testConcurrentCancel() throws Exception { |
711 | 713 | } |
712 | 714 | } |
713 | 715 |
|
| 716 | + /** |
| 717 | + * Waits until the given query id appears in {@code system.processes}, so we know the operation has actually |
| 718 | + * started on the server before attempting to cancel it. Uses a dedicated connection (no session) to observe. |
| 719 | + */ |
| 720 | + private boolean waitForQueryToStart(String queryId, int timeoutSeconds) throws Exception { |
| 721 | + if (queryId == null) { |
| 722 | + return false; |
| 723 | + } |
| 724 | + long deadline = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(timeoutSeconds); |
| 725 | + while (System.currentTimeMillis() < deadline) { |
| 726 | + try (Connection conn = getJdbcConnection(); |
| 727 | + Statement stmt = conn.createStatement(); |
| 728 | + ResultSet rs = stmt.executeQuery("SELECT count() FROM system.processes WHERE query_id = '" + queryId + "'")) { |
| 729 | + |
| 730 | + if (rs.next() && rs.getLong(1) > 0) { |
| 731 | + return true; |
| 732 | + } |
| 733 | + } |
| 734 | + Thread.sleep(200); |
| 735 | + } |
| 736 | + return false; |
| 737 | + } |
| 738 | + |
| 739 | + private String waitForQueryId(StatementImpl stmt, int timeoutSeconds) throws Exception { |
| 740 | + long deadline = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(timeoutSeconds); |
| 741 | + while (System.currentTimeMillis() < deadline) { |
| 742 | + String queryId = stmt.getLastQueryId(); |
| 743 | + if (queryId != null && !queryId.isEmpty()) { |
| 744 | + return queryId; |
| 745 | + } |
| 746 | + Thread.sleep(50); |
| 747 | + } |
| 748 | + return null; |
| 749 | + } |
| 750 | + |
| 751 | + @Test(groups = {"integration"}) |
| 752 | + public void testCancelQueryWithSession() throws Exception { |
| 753 | + if (isCloud()) { |
| 754 | + throw new SkipException("Cloud + HTTP doesn't work well. Enough to test locally"); |
| 755 | + } |
| 756 | + |
| 757 | + // Regression test for #2690: cancelling a query that runs inside a session must not fail with |
| 758 | + // "Session is locked by a concurrent client" (SESSION_IS_LOCKED). The KILL QUERY request issued by |
| 759 | + // cancel() must not carry the session id of the query being cancelled. |
| 760 | + String sessionId = "test-session-" + UUID.randomUUID(); |
| 761 | + try (Connection conn = getJdbcConnection()) { |
| 762 | + try (StatementImpl stmt = (StatementImpl) conn.createStatement()) { |
| 763 | + stmt.getLocalSettings().setSessionId(sessionId); |
| 764 | + stmt.setQueryTimeout(30); // safety net so a failed cancel cannot hang the test |
| 765 | + |
| 766 | + final AtomicReference<Throwable> threadError = new AtomicReference<>(); |
| 767 | + final CountDownLatch started = new CountDownLatch(1); |
| 768 | + Thread worker = new Thread(() -> { |
| 769 | + started.countDown(); |
| 770 | + // Long-running query that only completes when killed. |
| 771 | + try (ResultSet rs = stmt.executeQuery("SELECT count() FROM system.numbers_mt")) { |
| 772 | + rs.next(); |
| 773 | + } catch (Throwable t) { |
| 774 | + System.out.println("Error: " + t.getMessage()); |
| 775 | + threadError.set(t); |
| 776 | + } |
| 777 | + }); |
| 778 | + worker.start(); |
| 779 | + started.await(); |
| 780 | + |
| 781 | + String queryId = waitForQueryId(stmt, 15); |
| 782 | + assertNotNull(queryId, "Query id was not assigned in time"); |
| 783 | + assertTrue(waitForQueryToStart(queryId, 15), "Query did not start on the server in time"); |
| 784 | + |
| 785 | + // Cancel from the main thread - must not throw SESSION_IS_LOCKED. |
| 786 | + stmt.cancel(); |
| 787 | + |
| 788 | + worker.join(TimeUnit.SECONDS.toMillis(20)); |
| 789 | + assertFalse(worker.isAlive(), "Query was not cancelled and is still running"); |
| 790 | + } |
| 791 | + } |
| 792 | + } |
| 793 | + |
| 794 | + @Test(groups = {"integration"}) |
| 795 | + public void testCancelInsertWithSession() throws Exception { |
| 796 | + if (isCloud()) { |
| 797 | + throw new SkipException("Cloud + HTTP doesn't work well. Enough to test locally"); |
| 798 | + } |
| 799 | + // Regression test for #2690 covering a long-running INSERT executed inside a session. |
| 800 | + String tableName = getDatabase() + ".cancel_insert_with_session"; |
| 801 | + String sessionId = "test-session-" + UUID.randomUUID(); |
| 802 | + try (Connection conn = getJdbcConnection(Map.of(ASYNC_INSERT_SETTING_KEY, ServerSettings.OFF))) { |
| 803 | + try (Statement setup = conn.createStatement()) { |
| 804 | + setup.execute("DROP TABLE IF EXISTS " + tableName); |
| 805 | + setup.execute("CREATE TABLE " + tableName + " (num UInt64) ENGINE = MergeTree ORDER BY ()"); |
| 806 | + } |
| 807 | + |
| 808 | + try (StatementImpl stmt = (StatementImpl) conn.createStatement()) { |
| 809 | + stmt.getLocalSettings().setSessionId(sessionId); |
| 810 | + stmt.setQueryTimeout(30); // safety net so a failed cancel cannot hang the test |
| 811 | + |
| 812 | + final AtomicReference<Throwable> threadError = new AtomicReference<>(); |
| 813 | + final CountDownLatch started = new CountDownLatch(1); |
| 814 | + Thread worker = new Thread(() -> { |
| 815 | + started.countDown(); |
| 816 | + // Long-running insert that only completes when killed. |
| 817 | + try { |
| 818 | + stmt.executeUpdate("INSERT INTO " + tableName + " SELECT number FROM system.numbers_mt"); |
| 819 | + } catch (Throwable t) { |
| 820 | + threadError.set(t); |
| 821 | + } |
| 822 | + }); |
| 823 | + worker.start(); |
| 824 | + started.await(); |
| 825 | + |
| 826 | + String queryId = waitForQueryId(stmt, 15); |
| 827 | + assertNotNull(queryId, "Query id was not assigned in time"); |
| 828 | + assertTrue(waitForQueryToStart(queryId, 15), "Insert did not start on the server in time"); |
| 829 | + |
| 830 | + // Cancel from the main thread - must not throw SESSION_IS_LOCKED. |
| 831 | + stmt.cancel(); |
| 832 | + |
| 833 | + worker.join(TimeUnit.SECONDS.toMillis(20)); |
| 834 | + assertFalse(worker.isAlive(), "Insert was not cancelled and is still running"); |
| 835 | + } finally { |
| 836 | + try (Statement cleanup = conn.createStatement()) { |
| 837 | + cleanup.execute("DROP TABLE IF EXISTS " + tableName); |
| 838 | + } |
| 839 | + } |
| 840 | + } |
| 841 | + } |
| 842 | + |
714 | 843 | @Test(groups = {"integration"}) |
715 | 844 | public void testTextFormatInResponse() throws Exception { |
716 | 845 | try (Connection conn = getJdbcConnection(); |
|
0 commit comments