Rename option location to bigquery_location

Author: BoqianShi

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java

@@ -427,7 +427,7 @@ public static SparkBigQueryConfig from(
     config.parentProjectId =
         getAnyOption(globalOptions, options, "parentProject").or(defaultBilledProject());
     config.catalogProjectId = getOption(options, "projectId");
-    config.catalogLocation = getOption(options, "location");
+    config.catalogLocation = getOption(options, "bigquery_location");
     config.useParentProjectForMetadataOperations =
         getAnyBooleanOption(globalOptions, options, "useParentProjectForMetadataOperations", false);
     config.accessTokenProviderFQCN = getAnyOption(globalOptions, options, "gcpAccessTokenProvider");

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java

@@ -1296,7 +1296,7 @@ public void testCatalogProjectId() {
   @Test
   public void testCatalogLocation() {
     Map<String, String> optionsWithLocation = new HashMap<>();
-    optionsWithLocation.put("location", "US");
+    optionsWithLocation.put("bigquery_location", "US");
     optionsWithLocation.put("table", "dataset.table");
 
     SparkBigQueryConfig configWithLocation =
@@ -1332,7 +1332,7 @@ public void testCatalogLocation() {
   public void testCatalogProjectIdAndLocation() {
     Map<String, String> optionsWithProjectAndLocation = new HashMap<>();
     optionsWithProjectAndLocation.put("projectId", "my-catalog-project");
-    optionsWithProjectAndLocation.put("location", "EU");
+    optionsWithProjectAndLocation.put("bigquery_location", "EU");
     optionsWithProjectAndLocation.put("table", "dataset.table");
 
     SparkBigQueryConfig config =

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/CatalogIntegrationTestBase.java

@@ -29,7 +29,9 @@
 import org.apache.spark.sql.RowFactory;
 import org.apache.spark.sql.SparkSession;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Ignore;
 import org.junit.Test;
@@ -41,7 +43,32 @@ public class CatalogIntegrationTestBase {
 
   BigQuery bigquery = IntegrationTestUtils.getBigquery();
 
+  protected static SparkSession spark;
   private String testTable;
+  // 2. Initialize the SparkSession ONCE before all tests
+  @BeforeClass
+  public static void setupSparkSession() {
+    spark =
+        SparkSession.builder()
+            .appName("catalog test")
+            .master("local[*]")
+            .config("spark.sql.legacy.createHiveTableByDefault", "false")
+            .config("spark.sql.sources.default", "bigquery")
+            .config("spark.datasource.bigquery.writeMethod", "direct")
+            .config("spark.sql.defaultCatalog", "bigquery")
+            .config("spark.sql.catalog.bigquery", "com.google.cloud.spark.bigquery.BigQueryCatalog")
+            .getOrCreate();
+  }
+
+  // 4. Stop the SparkSession ONCE after all tests are done
+  // This fixes the local IllegalStateException (race condition)
+  @AfterClass
+  public static void teardownSparkSession() {
+    if (spark != null) {
+      spark.stop();
+      spark = null;
+    }
+  }
 
   @Before
   public void renameTestTable() {
@@ -71,12 +98,10 @@ public void testCreateTableInCustomNamespace() throws Exception {
 
   private void internalTestCreateTable(String dataset) throws InterruptedException {
     assertThat(bigquery.getDataset(DatasetId.of(dataset))).isNotNull();
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql("CREATE TABLE " + fullTableName(dataset) + "(id int, data string);");
-      Table table = bigquery.getTable(TableId.of(dataset, testTable));
-      assertThat(table).isNotNull();
-      assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(0L);
-    }
+    spark.sql("CREATE TABLE " + fullTableName(dataset) + "(id int, data string);");
+    Table table = bigquery.getTable(TableId.of(dataset, testTable));
+    assertThat(table).isNotNull();
+    assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(0L);
   }
 
   @Test
@@ -91,13 +116,11 @@ public void testCreateTableAndInsertInCustomNamespace() throws Exception {
 
   private void internalTestCreateTableAndInsert(String dataset) throws InterruptedException {
     assertThat(bigquery.getDataset(DatasetId.of(dataset))).isNotNull();
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql("CREATE TABLE " + fullTableName(dataset) + "(id int, data string);");
-      spark.sql(String.format("INSERT INTO `%s`.`%s` VALUES (1, 'foo');", dataset, testTable));
-      Table table = bigquery.getTable(TableId.of(dataset, testTable));
-      assertThat(table).isNotNull();
-      assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(1L);
-    }
+    spark.sql("CREATE TABLE " + fullTableName(dataset) + "(id int, data string);");
+    spark.sql(String.format("INSERT INTO `%s`.`%s` VALUES (1, 'foo');", dataset, testTable));
+    Table table = bigquery.getTable(TableId.of(dataset, testTable));
+    assertThat(table).isNotNull();
+    assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(1L);
   }
 
   @Test
@@ -112,12 +135,10 @@ public void testCreateTableAsSelectInCustomNamespace() throws Exception {
 
   private void internalTestCreateTableAsSelect(String dataset) throws InterruptedException {
     assertThat(bigquery.getDataset(DatasetId.of(dataset))).isNotNull();
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql("CREATE TABLE " + fullTableName(dataset) + " AS SELECT 1 AS id, 'foo' AS data;");
-      Table table = bigquery.getTable(TableId.of(dataset, testTable));
-      assertThat(table).isNotNull();
-      assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(1L);
-    }
+    spark.sql("CREATE TABLE " + fullTableName(dataset) + " AS SELECT 1 AS id, 'foo' AS data;");
+    Table table = bigquery.getTable(TableId.of(dataset, testTable));
+    assertThat(table).isNotNull();
+    assertThat(selectCountStarFrom(dataset, testTable)).isEqualTo(1L);
   }
 
   @Test
@@ -135,23 +156,21 @@ public void testCreateTableWithExplicitTargetInCustomNamespace() throws Exceptio
   private void internalTestCreateTableWithExplicitTarget(String dataset)
       throws InterruptedException {
     assertThat(bigquery.getDataset(DatasetId.of(dataset))).isNotNull();
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql(
-          "CREATE TABLE "
-              + fullTableName(dataset)
-              + " OPTIONS (table='bigquery-public-data.samples.shakespeare')");
-      List<Row> result =
-          spark
-              .sql(
-                  "SELECT word, SUM(word_count) FROM "
-                      + fullTableName(dataset)
-                      + " WHERE word='spark' GROUP BY word;")
-              .collectAsList();
-      assertThat(result).hasSize(1);
-      Row resultRow = result.get(0);
-      assertThat(resultRow.getString(0)).isEqualTo("spark");
-      assertThat(resultRow.getLong(1)).isEqualTo(10L);
-    }
+    spark.sql(
+        "CREATE TABLE "
+            + fullTableName(dataset)
+            + " OPTIONS (table='bigquery-public-data.samples.shakespeare')");
+    List<Row> result =
+        spark
+            .sql(
+                "SELECT word, SUM(word_count) FROM "
+                    + fullTableName(dataset)
+                    + " WHERE word='spark' GROUP BY word;")
+            .collectAsList();
+    assertThat(result).hasSize(1);
+    Row resultRow = result.get(0);
+    assertThat(resultRow.getString(0)).isEqualTo("spark");
+    assertThat(resultRow.getLong(1)).isEqualTo(10L);
   }
 
   private String fullTableName(String dataset) {
@@ -176,14 +195,11 @@ private long selectCountStarFrom(String dataset, String table) throws Interrupte
 
   @Test
   public void testReadFromDifferentBigQueryProject() throws Exception {
-    try (SparkSession spark = createSparkSession()) {
-      List<Row> df =
-          spark
-              .sql(
-                  "SELECT * from `bigquery-public-data`.`samples`.`shakespeare` WHERE word='spark'")
-              .collectAsList();
-      assertThat(df).hasSize(9);
-    }
+    List<Row> df =
+        spark
+            .sql("SELECT * from `bigquery-public-data`.`samples`.`shakespeare` WHERE word='spark'")
+            .collectAsList();
+    assertThat(df).hasSize(9);
   }
 
   @Test
@@ -192,45 +208,36 @@ public void testListNamespaces() throws Exception {
         String.format("show_databases_test_%s_%s", System.currentTimeMillis(), System.nanoTime());
     DatasetId datasetId = DatasetId.of(database);
     bigquery.create(Dataset.newBuilder(datasetId).build());
-    try (SparkSession spark = createSparkSession()) {
-      List<Row> databases = spark.sql("SHOW DATABASES").collectAsList();
-      assertThat(databases).contains(RowFactory.create(database));
-    } finally {
-      bigquery.delete(datasetId);
-    }
+    List<Row> databases = spark.sql("SHOW DATABASES").collectAsList();
+    assertThat(databases).contains(RowFactory.create(database));
+    bigquery.delete(datasetId);
   }
 
   @Test
   public void testCreateNamespace() throws Exception {
     String database =
         String.format("create_database_test_%s_%s", System.currentTimeMillis(), System.nanoTime());
     DatasetId datasetId = DatasetId.of(database);
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql("CREATE DATABASE " + database + ";");
-      Dataset dataset = bigquery.getDataset(datasetId);
-      assertThat(dataset).isNotNull();
-    } finally {
-      bigquery.delete(datasetId);
-    }
+    spark.sql("CREATE DATABASE " + database + ";");
+    Dataset dataset = bigquery.getDataset(datasetId);
+    assertThat(dataset).isNotNull();
+    bigquery.delete(datasetId);
   }
 
   @Test
   public void testCreateNamespaceWithLocation() throws Exception {
     String database =
         String.format("create_database_test_%s_%s", System.currentTimeMillis(), System.nanoTime());
     DatasetId datasetId = DatasetId.of(database);
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql(
-          "CREATE DATABASE "
-              + database
-              + " COMMENT 'foo' WITH DBPROPERTIES (bigquery_location = 'us-east1');");
-      Dataset dataset = bigquery.getDataset(datasetId);
-      assertThat(dataset).isNotNull();
-      assertThat(dataset.getLocation()).isEqualTo("us-east1");
-      assertThat(dataset.getDescription()).isEqualTo("foo");
-    } finally {
-      bigquery.delete(datasetId);
-    }
+    spark.sql(
+        "CREATE DATABASE "
+            + database
+            + " COMMENT 'foo' WITH DBPROPERTIES (bigquery_location = 'us-east1');");
+    Dataset dataset = bigquery.getDataset(datasetId);
+    assertThat(dataset).isNotNull();
+    assertThat(dataset.getLocation()).isEqualTo("us-east1");
+    assertThat(dataset.getDescription()).isEqualTo("foo");
+    bigquery.delete(datasetId);
   }
 
   @Test
@@ -239,85 +246,73 @@ public void testDropDatabase() {
         String.format("drop_database_test_%s_%s", System.currentTimeMillis(), System.nanoTime());
     DatasetId datasetId = DatasetId.of(database);
     bigquery.create(Dataset.newBuilder(datasetId).build());
-    try (SparkSession spark = createSparkSession()) {
-      spark.sql("DROP DATABASE " + database + ";");
-      Dataset dataset = bigquery.getDataset(datasetId);
-      assertThat(dataset).isNull();
-    }
+    spark.sql("DROP DATABASE " + database + ";");
+    Dataset dataset = bigquery.getDataset(datasetId);
+    assertThat(dataset).isNull();
   }
 
   @Test
   public void testCatalogInitializationWithProject() {
-    try (SparkSession spark = createSparkSession()) {
-      spark
-          .conf()
-          .set(
-              "spark.sql.catalog.public_catalog",
-              "com.google.cloud.spark.bigquery.BigQueryCatalog");
-      spark.conf().set("spark.sql.catalog.public_catalog.projectId", "bigquery-public-data");
-      List<Row> rows = spark.sql("SHOW DATABASES IN public_catalog").collectAsList();
-      List<String> databaseNames =
-          rows.stream().map(row -> row.getString(0)).collect(Collectors.toList());
-      assertThat(databaseNames).contains("samples");
-      System.out.println(databaseNames);
-      spark.sql("SHOW TABLES IN public_catalog.samples").show();
-      List<Row> data =
-          spark.sql("SELECT * FROM public_catalog.samples.shakespeare LIMIT 10").collectAsList();
-      assertThat(data).hasSize(10);
-    }
+    spark
+        .conf()
+        .set("spark.sql.catalog.public_catalog", "com.google.cloud.spark.bigquery.BigQueryCatalog");
+    spark.conf().set("spark.sql.catalog.public_catalog.project", "bigquery-public-data");
+
+    List<Row> rows = spark.sql("SHOW DATABASES IN public_catalog").collectAsList();
+    List<String> databaseNames =
+        rows.stream().map(row -> row.getString(0)).collect(Collectors.toList());
+    assertThat(databaseNames).contains("samples");
+
+    List<Row> data =
+        spark.sql("SELECT * FROM public_catalog.samples.shakespeare LIMIT 10").collectAsList();
+    assertThat(data).hasSize(10);
   }
 
   @Test
-  public void testCreateCatalogWithLocation() {
+  public void testCreateCatalogWithLocation() throws Exception {
     String database = String.format("create_db_with_location_%s", System.nanoTime());
     DatasetId datasetId = DatasetId.of(database);
-    try (SparkSession spark = createSparkSession()) {
-      spark
-          .conf()
-          .set("spark.sql.catalog.test_catalog", "com.google.cloud.spark.bigquery.BigQueryCatalog");
-      spark.conf().set("spark.sql.catalog.test_catalog.location", "EU");
-
-      spark.sql("CREATE DATABASE test_catalog." + database);
-
-      Dataset dataset = bigquery.getDataset(datasetId);
-      assertThat(dataset).isNotNull();
-      assertThat(dataset.getLocation()).isEqualTo("EU");
-    } finally {
-      bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
-    }
+    spark
+        .conf()
+        .set(
+            "spark.sql.catalog.test_location_catalog",
+            "com.google.cloud.spark.bigquery.BigQueryCatalog");
+    spark.conf().set("spark.sql.catalog.test_location_catalog.bigquery_location", "EU");
+    spark.sql("CREATE DATABASE test_location_catalog." + database);
+    Dataset dataset = bigquery.getDataset(datasetId);
+    assertThat(dataset).isNotNull();
+    assertThat(dataset.getLocation()).isEqualTo("EU");
+    bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
   }
 
   @Test
   public void testCreateTableAsSelectWithProjectAndLocation() {
     String database = String.format("ctas_db_with_location_%s", System.nanoTime());
     String newTable = "ctas_table_from_public";
     DatasetId datasetId = DatasetId.of(database);
-    try (SparkSession spark = createSparkSession()) {
-      spark
-          .conf()
-          .set(
-              "spark.sql.catalog.public_catalog",
-              "com.google.cloud.spark.bigquery.BigQueryCatalog");
-      spark.conf().set("spark.sql.catalog.public_catalog.projectId", "bigquery-public-data");
-      spark
-          .conf()
-          .set("spark.sql.catalog.test_catalog", "com.google.cloud.spark.bigquery.BigQueryCatalog");
-      spark.conf().set("spark.sql.catalog.test_catalog.location", "EU");
-      spark.sql("CREATE DATABASE test_catalog." + database);
-      spark.sql(
-          "CREATE TABLE test_catalog."
-              + database
-              + "."
-              + newTable
-              + " AS SELECT * FROM public_catalog.samples.shakespeare LIMIT 10");
-      Dataset dataset = bigquery.getDataset(datasetId);
-      assertThat(dataset).isNotNull();
-      assertThat(dataset.getLocation()).isEqualTo("EU");
-      Table table = bigquery.getTable(TableId.of(datasetId.getDataset(), newTable));
-      assertThat(table).isNotNull();
-    } finally {
-      bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
-    }
+    spark
+        .conf()
+        .set("spark.sql.catalog.public_catalog", "com.google.cloud.spark.bigquery.BigQueryCatalog");
+    spark.conf().set("spark.sql.catalog.public_catalog.projectId", "bigquery-public-data");
+    spark
+        .conf()
+        .set(
+            "spark.sql.catalog.test_catalog_as_select",
+            "com.google.cloud.spark.bigquery.BigQueryCatalog");
+    spark.conf().set("spark.sql.catalog.test_catalog_as_select.bigquery_location", "EU");
+    spark.sql("CREATE DATABASE test_catalog_as_select." + database);
+    spark.sql(
+        "CREATE TABLE test_catalog_as_select."
+            + database
+            + "."
+            + newTable
+            + " AS SELECT * FROM public_catalog.samples.shakespeare LIMIT 10");
+    Dataset dataset = bigquery.getDataset(datasetId);
+    assertThat(dataset).isNotNull();
+    assertThat(dataset.getLocation()).isEqualTo("EU");
+    Table table = bigquery.getTable(TableId.of(datasetId.getDataset(), newTable));
+    assertThat(table).isNotNull();
+    bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
   }
 
   private static SparkSession createSparkSession() {