diff --git a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java index 9fe1c4f0f2fc..0d5cdac0650d 100644 --- a/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java +++ b/java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java @@ -56,6 +56,10 @@ import java.time.LocalTime; import java.util.Properties; import java.util.Random; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.function.BiFunction; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.Assertions; @@ -406,6 +410,50 @@ public void testLocation() throws SQLException { connection2.close(); } + @Test + public void testCancelLocation() throws Exception { + String connection_uri = + "jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID=" + + PROJECT_ID + + ";OAUTHTYPE=3;LOCATION=EU"; + + Driver driver = BigQueryDriver.getRegisteredDriver(); + Connection connection = driver.connect(connection_uri, new Properties()); + Statement statement = connection.createStatement(); + + // Query a dataset in the EU with a heavy query so that we can cancel it while it runs. + String query = + "SELECT COUNT(*) FROM `bigquery-public-data.covid19_italy_eu.data_by_province` a " + + "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` b " + + "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` c"; + + // Run the query in a separate thread so we can cancel it from the main thread + ExecutorService executor = Executors.newSingleThreadExecutor(); + Future future = executor.submit(() -> { + try { + statement.executeQuery(query); + return null; + } catch (SQLException e) { + return e; + } + }); + + // Wait a short moment to let the query submit and start running + Thread.sleep(1500); + + // Cancel the query execution + statement.cancel(); + + // Verify that the query threw a SQLException indicating it was cancelled + SQLException exception = future.get(15, TimeUnit.SECONDS); + assertNotNull(exception, "Expected SQLException to be thrown due to cancellation"); + assertTrue(exception.getMessage().contains("cancelled") || exception.getMessage().contains("Job was cancelled") || exception.getMessage().contains("Query was cancelled"), + "Unexpected exception message: " + exception.getMessage()); + + connection.close(); + executor.shutdown(); + } + @Test public void testIncorrectLocation() throws SQLException { String connection_uri =