Skip to content

Commit 547c032

Browse files
Neenu1995Neenu1995
committed
add cross region checks
1 parent 8289679 commit 547c032

1 file changed

Lines changed: 168 additions & 9 deletions

File tree

‎java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITBigQueryJDBCTest.java‎

Lines changed: 168 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,15 @@
3030
import com.google.cloud.bigquery.BigQueryOptions;
3131
import com.google.cloud.bigquery.DatasetId;
3232
import com.google.cloud.bigquery.Job;
33+
import com.google.cloud.bigquery.JobId;
3334
import com.google.cloud.bigquery.JobInfo;
3435
import com.google.cloud.bigquery.QueryJobConfiguration;
3536
import com.google.cloud.bigquery.exception.BigQueryJdbcException;
3637
import com.google.cloud.bigquery.exception.BigQueryJdbcSqlFeatureNotSupportedException;
3738
import com.google.cloud.bigquery.exception.BigQueryJdbcSqlSyntaxErrorException;
3839
import com.google.cloud.bigquery.jdbc.BigQueryConnection;
3940
import com.google.cloud.bigquery.jdbc.BigQueryDriver;
41+
import com.google.cloud.bigquery.jdbc.BigQueryStatement;
4042
import com.google.cloud.bigquery.jdbc.DataSource;
4143
import com.google.common.collect.ImmutableMap;
4244
import java.io.File;
@@ -429,14 +431,16 @@ public void testCancelLocation() throws Exception {
429431

430432
// Run the query in a separate thread so we can cancel it from the main thread
431433
ExecutorService executor = Executors.newSingleThreadExecutor();
432-
Future<SQLException> future = executor.submit(() -> {
433-
try {
434-
statement.executeQuery(query);
435-
return null;
436-
} catch (SQLException e) {
437-
return e;
438-
}
439-
});
434+
Future<SQLException> future =
435+
executor.submit(
436+
() -> {
437+
try {
438+
statement.executeQuery(query);
439+
return null;
440+
} catch (SQLException e) {
441+
return e;
442+
}
443+
});
440444

441445
// Wait a short moment to let the query submit and start running
442446
Thread.sleep(1500);
@@ -447,7 +451,162 @@ public void testCancelLocation() throws Exception {
447451
// Verify that the query threw a SQLException indicating it was cancelled
448452
SQLException exception = future.get(15, TimeUnit.SECONDS);
449453
assertNotNull(exception, "Expected SQLException to be thrown due to cancellation");
450-
assertTrue(exception.getMessage().contains("cancelled") || exception.getMessage().contains("Job was cancelled") || exception.getMessage().contains("Query was cancelled"),
454+
assertTrue(
455+
exception.getMessage().contains("cancelled")
456+
|| exception.getMessage().contains("Job was cancelled")
457+
|| exception.getMessage().contains("Query was cancelled"),
458+
"Unexpected exception message: " + exception.getMessage());
459+
460+
connection.close();
461+
executor.shutdown();
462+
}
463+
464+
@Test
465+
public void testEuConnectionQueryAndCancel() throws Exception {
466+
// 1. Create a JDBC connection explicitly set to the EU region
467+
String connectionUri =
468+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID="
469+
+ PROJECT_ID
470+
+ ";OAUTHTYPE=3;LOCATION=EU;JobCreationMode=1;useQueryCache=false";
471+
472+
Driver driver = BigQueryDriver.getRegisteredDriver();
473+
try (Connection connection = driver.connect(connectionUri, new Properties())) {
474+
Statement statement = connection.createStatement();
475+
476+
// 2. Run a query in the EU region
477+
String euQuery =
478+
"SELECT COUNT(*) FROM `bigquery-public-data.covid19_italy_eu.data_by_province` a "
479+
+ "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` b "
480+
+ "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` c /* "
481+
+ java.util.UUID.randomUUID()
482+
+ " */";
483+
484+
ExecutorService executor = Executors.newSingleThreadExecutor();
485+
Future<SQLException> future =
486+
executor.submit(
487+
() -> {
488+
try {
489+
statement.executeQuery(euQuery);
490+
return null;
491+
} catch (SQLException e) {
492+
return e;
493+
}
494+
});
495+
496+
// Retrieve the JobId from the statement with polling
497+
BigQueryStatement rawStmt = statement.unwrap(BigQueryStatement.class);
498+
java.lang.reflect.Field jobIdsField = BigQueryStatement.class.getDeclaredField("jobIds");
499+
jobIdsField.setAccessible(true);
500+
@SuppressWarnings("unchecked")
501+
java.util.List<JobId> jobIds = (java.util.List<JobId>) jobIdsField.get(rawStmt);
502+
long startTime = System.currentTimeMillis();
503+
while (jobIds.isEmpty() && (System.currentTimeMillis() - startTime < 15000)) {
504+
Thread.sleep(500);
505+
}
506+
assertFalse(jobIds.isEmpty(), "Expected at least one active Job ID in the statement");
507+
JobId runningJobId = jobIds.get(0);
508+
509+
// Verify that the JobId location is indeed "EU"
510+
assertEquals("EU", runningJobId.getLocation(), "Expected JobId to have location EU");
511+
512+
// 3. Verify that the job can be cancelled successfully
513+
statement.cancel();
514+
515+
SQLException cancelException = future.get(40, TimeUnit.SECONDS);
516+
assertNotNull(cancelException, "Expected SQLException to be thrown due to cancellation");
517+
assertTrue(
518+
cancelException.getMessage().contains("cancelled")
519+
|| cancelException.getMessage().contains("Job was cancelled")
520+
|| cancelException.getMessage().contains("Query was cancelled"),
521+
"Unexpected exception message: " + cancelException.getMessage());
522+
523+
executor.shutdown();
524+
525+
// 4. Use the same connection to run a query in the US region (should fail)
526+
String usQuery =
527+
"SELECT COUNT(*) FROM `bigquery-public-data.usa_names.usa_1910_current` LIMIT 10";
528+
SQLException usException =
529+
assertThrows(
530+
SQLException.class,
531+
() -> {
532+
statement.executeQuery(usQuery);
533+
},
534+
"Expected SQLException because US dataset cannot be queried via EU connection");
535+
536+
assertTrue(
537+
usException.getMessage().contains("Not found: Table")
538+
|| usException.getMessage().contains("Location")
539+
|| usException.getMessage().contains("europe")
540+
|| usException.getMessage().contains("EU")
541+
|| usException.getMessage().contains("permission")
542+
|| usException.getMessage().contains("not exist"),
543+
"Unexpected exception message for cross-region query: " + usException.getMessage());
544+
}
545+
}
546+
547+
@Test
548+
public void testJdbcAutoLocationCancel() throws Exception {
549+
// 1. Establish connection WITHOUT location parameter
550+
String connectionUri =
551+
"jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443;PROJECTID="
552+
+ PROJECT_ID
553+
+ ";OAUTHTYPE=3;JobCreationMode=1;useQueryCache=false";
554+
555+
Driver driver = BigQueryDriver.getRegisteredDriver();
556+
Connection connection = driver.connect(connectionUri, new Properties());
557+
Statement statement = connection.createStatement();
558+
559+
// Query a dataset in the EU. Since connection has no location, BigQuery routes the job to EU
560+
// automatically.
561+
String query =
562+
"SELECT COUNT(*) FROM `bigquery-public-data.covid19_italy_eu.data_by_province` a "
563+
+ "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` b "
564+
+ "CROSS JOIN `bigquery-public-data.covid19_italy_eu.data_by_province` c /* "
565+
+ java.util.UUID.randomUUID()
566+
+ " */";
567+
568+
ExecutorService executor = Executors.newSingleThreadExecutor();
569+
Future<SQLException> future =
570+
executor.submit(
571+
() -> {
572+
try {
573+
statement.executeQuery(query);
574+
return null;
575+
} catch (SQLException e) {
576+
return e;
577+
}
578+
});
579+
580+
// Retrieve the JobId from the statement with polling
581+
BigQueryStatement rawStmt = statement.unwrap(BigQueryStatement.class);
582+
java.lang.reflect.Field jobIdsField = BigQueryStatement.class.getDeclaredField("jobIds");
583+
jobIdsField.setAccessible(true);
584+
@SuppressWarnings("unchecked")
585+
java.util.List<JobId> jobIds = (java.util.List<JobId>) jobIdsField.get(rawStmt);
586+
long startTime = System.currentTimeMillis();
587+
while (jobIds.isEmpty() && (System.currentTimeMillis() - startTime < 15000)) {
588+
Thread.sleep(500);
589+
}
590+
assertFalse(jobIds.isEmpty(), "Expected at least one active Job ID in the statement");
591+
JobId runningJobId = jobIds.get(0);
592+
593+
// Verify that the JobId has location set to "EU" even though the connection did not specify any
594+
// location!
595+
assertEquals(
596+
"EU",
597+
runningJobId.getLocation(),
598+
"Expected JobId to have location EU populated by the backend");
599+
600+
// Cancel the query execution via the JDBC Statement
601+
statement.cancel();
602+
603+
// Verify that the query threw a SQLException indicating it was cancelled
604+
SQLException exception = future.get(40, TimeUnit.SECONDS);
605+
assertNotNull(exception, "Expected SQLException to be thrown due to cancellation");
606+
assertTrue(
607+
exception.getMessage().contains("cancelled")
608+
|| exception.getMessage().contains("Job was cancelled")
609+
|| exception.getMessage().contains("Query was cancelled"),
451610
"Unexpected exception message: " + exception.getMessage());
452611

453612
connection.close();

0 commit comments

Comments
 (0)