Simplify table creation for integration tests

connector/src/main/java/com/google/cloud/hive/bigquery/connector/input/arrow/ArrowRecordReader.java

@@ -17,7 +17,6 @@
 
 import com.google.cloud.hive.bigquery.connector.BigQuerySerDe;
 import com.google.cloud.hive.bigquery.connector.input.BigQueryInputSplit;
-
 import java.io.IOException;
 import java.util.List;
 import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;

connector/src/main/java/com/google/cloud/hive/bigquery/connector/output/direct/DirectRecordWriter.java

@@ -20,7 +20,6 @@
 import com.google.cloud.hive.bigquery.connector.BigQuerySerDe;
 import com.google.cloud.hive.bigquery.connector.JobDetails;
 import com.google.cloud.hive.bigquery.connector.utils.hive.HiveUtils;
-
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import org.apache.hadoop.fs.FSDataOutputStream;

connector/src/test/java/com/google/cloud/hive/bigquery/connector/BigLakeIntegrationTests.java

@@ -33,7 +33,7 @@ public void testReadBigLakeTable(
     runBqQuery(BIGQUERY_BIGLAKE_TABLE_CREATE_QUERY);
     // Create Hive table
     initHive(engine, readDataFormat);
-    runHiveScript(HIVE_BIGLAKE_TABLE_CREATE_QUERY);
+    createExternalTable(BIGLAKE_TABLE_NAME, HIVE_BIGLAKE_TABLE_DDL);
     // Read data
     List<Object[]> rows = runHiveStatement(String.format("SELECT * FROM %s", BIGLAKE_TABLE_NAME));
     assertArrayEquals(

connector/src/test/java/com/google/cloud/hive/bigquery/connector/ConfigValidationIntegrationTests.java

@@ -15,7 +15,7 @@
  */
 package com.google.cloud.hive.bigquery.connector;
 
-import static com.google.cloud.hive.bigquery.connector.TestUtils.HIVE_TEST_TABLE_CREATE_QUERY;
+import static com.google.cloud.hive.bigquery.connector.TestUtils.HIVE_TEST_TABLE_DDL;
 import static com.google.cloud.hive.bigquery.connector.TestUtils.TEST_TABLE_NAME;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -56,7 +56,7 @@ public void testMissingGcsTempPath() {
     hive.setHiveConfValue(
         HiveBigQueryConfig.WRITE_METHOD_KEY, HiveBigQueryConfig.WRITE_METHOD_INDIRECT);
     initHive("tez", HiveBigQueryConfig.AVRO, "");
-    runHiveScript(HIVE_TEST_TABLE_CREATE_QUERY);
+    createExternalTable(TEST_TABLE_NAME, HIVE_TEST_TABLE_DDL);
     Throwable exception =
         assertThrows(
             RuntimeException.class,
@@ -78,7 +78,7 @@ public void testMissingBucketPermissions() {
     hive.setHiveConfValue(
         HiveBigQueryConfig.WRITE_METHOD_KEY, HiveBigQueryConfig.WRITE_METHOD_INDIRECT);
     initHive("tez", HiveBigQueryConfig.AVRO, "gs://random-bucket-abcdef-12345");
-    runHiveScript(HIVE_TEST_TABLE_CREATE_QUERY);
+    createExternalTable(TEST_TABLE_NAME, HIVE_TEST_TABLE_DDL);
     Throwable exception =
         assertThrows(
             RuntimeException.class,

connector/src/test/java/com/google/cloud/hive/bigquery/connector/IntegrationTestsBase.java

@@ -109,6 +109,44 @@ public String renderQueryTemplate(String queryTemplate) {
     return StrSubstitutor.replace(queryTemplate, params, "${", "}");
   }
 
+  public void createHiveTable(
+      String tableName, String hiveDDL, boolean isExternal, String hiveProps) {
+    runHiveScript(
+        String.join(
+            "\n",
+            "CREATE " + (isExternal ? "EXTERNAL" : "") + " TABLE " + tableName + " (",
+            hiveDDL,
+            ")",
+            "STORED BY" + " 'com.google.cloud.hive.bigquery.connector.BigQueryStorageHandler'",
+            "TBLPROPERTIES (",
+            "  'bq.project'='${project}',",
+            "  'bq.dataset'='${dataset}',",
+            "  'bq.table'='" + tableName + "'",
+            (hiveProps != null ? "," + hiveProps : ""),
+            ");"));
+  }
+
+  public void createExternalTable(String tableName, String hiveDDL) {
+    createHiveTable(tableName, hiveDDL, true, null);
+  }
+
+  public void createExternalTable(String tableName, String hiveDDL, String bqDDL) {
+    createExternalTable(tableName, hiveDDL);
+    createBqTable(tableName, bqDDL);
+  }
+
+  public void createManagedTable(String tableName, String hiveDDL) {
+    createHiveTable(tableName, hiveDDL, false, null);
+  }
+
+  public void createManagedTableWithProps(String tableName, String hiveDDL, String hiveProps) {
+    createHiveTable(tableName, hiveDDL, false, hiveProps);
+  }
+
+  public void createBqTable(String tableName, String bqDDL) {
+    runBqQuery("CREATE OR REPLACE TABLE ${dataset}." + tableName + " (" + bqDDL + ")");
+  }
+
   public TableResult runBqQuery(String queryTemplate) {
     return getBigqueryClient().query(renderQueryTemplate(queryTemplate));
   }

connector/src/test/java/com/google/cloud/hive/bigquery/connector/ManagedAndExternalHiveTableIntegrationTests.java

@@ -32,9 +32,9 @@ public void testCreateManagedTable() {
     dropBqTableIfExists(dataset, MANAGED_TEST_TABLE_NAME);
     assertFalse(bQTableExists(dataset, MANAGED_TEST_TABLE_NAME));
     // Create the managed table using Hive
-    runHiveScript(HIVE_MANAGED_TEST_TABLE_CREATE_QUERY);
+    createManagedTable(MANAGED_TEST_TABLE_NAME, HIVE_MANAGED_TEST_TABLE_DDL);
     // Create another BQ table with the same schema
-    runBqQuery(BIGQUERY_ALL_TYPES_TABLE_CREATE_QUERY);
+    createBqTable(ALL_TYPES_TABLE_NAME, BIGQUERY_ALL_TYPES_TABLE_DDL);
     // Make sure that the managed table was created in BQ
     // and that the two schemas are the same
     TableInfo managedTableInfo = getTableInfo(dataset, MANAGED_TEST_TABLE_NAME);
@@ -51,12 +51,12 @@ public void testCreateManagedTable() {
   public void testCreateManagedTableAlreadyExists() {
     initHive();
     // Create the table in BigQuery
-    runBqQuery(BIGQUERY_MANAGED_TEST_TABLE_CREATE_QUERY);
+    createBqTable(MANAGED_TEST_TABLE_NAME, BIGQUERY_MANAGED_TEST_TABLE_DDL);
     // Try to create the managed table using Hive
     Throwable exception =
         assertThrows(
             IllegalArgumentException.class,
-            () -> runHiveScript(HIVE_MANAGED_TEST_TABLE_CREATE_QUERY));
+            () -> createManagedTable(MANAGED_TEST_TABLE_NAME, HIVE_MANAGED_TEST_TABLE_DDL));
     assertTrue(exception.getMessage().contains("BigQuery table already exists"));
   }
 
@@ -69,7 +69,7 @@ public void testDropManagedTable() {
     dropBqTableIfExists(dataset, MANAGED_TEST_TABLE_NAME);
     assertFalse(bQTableExists(dataset, MANAGED_TEST_TABLE_NAME));
     // Create the managed table using Hive
-    runHiveScript(HIVE_MANAGED_TEST_TABLE_CREATE_QUERY);
+    createManagedTable(MANAGED_TEST_TABLE_NAME, HIVE_MANAGED_TEST_TABLE_DDL);
     // Check that the table was created in BigQuery
     assertTrue(bQTableExists(dataset, MANAGED_TEST_TABLE_NAME));
     // Drop the managed table using hive
@@ -83,11 +83,8 @@ public void testDropManagedTable() {
   @Test
   public void testDropExternalTable() {
     initHive();
-    // Create the table in BigQuery
-    runBqQuery(BIGQUERY_TEST_TABLE_CREATE_QUERY);
-    // Create the corresponding external table in Hive
-    runHiveScript(HIVE_TEST_TABLE_CREATE_QUERY);
-    // Drop the external table
+    createExternalTable(TEST_TABLE_NAME, HIVE_TEST_TABLE_DDL, BIGQUERY_TEST_TABLE_DDL);
+    // Drop the external table in Hive
     runHiveScript("DROP TABLE " + TEST_TABLE_NAME);
     // Check that the table still exists in BigQuery
     assertTrue(bQTableExists(dataset, TEST_TABLE_NAME));

connector/src/test/java/com/google/cloud/hive/bigquery/connector/PartitionIntegrationTests.java

@@ -30,11 +30,15 @@ public class PartitionIntegrationTests extends IntegrationTestsBase {
 
   @Test
   public void testFieldTimePartition() {
+    initHive();
     // Make sure the BQ table doesn't exist
     dropBqTableIfExists(dataset, FIELD_TIME_PARTITIONED_TABLE_NAME);
     // Create the table using Hive
-    initHive();
-    runHiveScript(HIVE_FIELD_TIME_PARTITIONED_TABLE_CREATE_QUERY);
+    createManagedTableWithProps(
+        FIELD_TIME_PARTITIONED_TABLE_NAME,
+        HIVE_FIELD_TIME_PARTITIONED_TABLE_DDL,
+        HIVE_FIELD_TIME_PARTITIONED_TABLE_PROPS);
+    // Verify that the BQ table has the right partition & clustering options
     StandardTableDefinition tableDef =
         getTableInfo(dataset, FIELD_TIME_PARTITIONED_TABLE_NAME).getDefinition();
     TimePartitioning timePartitioning = tableDef.getTimePartitioning();
@@ -47,11 +51,14 @@ public void testFieldTimePartition() {
 
   @Test
   public void testCreateIngestionTimePartition() {
+    initHive();
     // Make sure the BQ table doesn't exist
     dropBqTableIfExists(dataset, INGESTION_TIME_PARTITIONED_TABLE_NAME);
     // Create the table using Hive
-    initHive();
-    runHiveScript(HIVE_INGESTION_TIME_PARTITIONED_TABLE_CREATE_QUERY);
+    createManagedTableWithProps(
+        INGESTION_TIME_PARTITIONED_TABLE_NAME,
+        HIVE_INGESTION_TIME_PARTITIONED_DDL,
+        HIVE_INGESTION_TIME_PARTITIONED_PROPS);
     // Retrieve the table metadata from BigQuery
     StandardTableDefinition tableDef =
         getTableInfo(dataset, INGESTION_TIME_PARTITIONED_TABLE_NAME).getDefinition();
@@ -79,7 +86,10 @@ public void testQueryIngestionTimePartition() {
     // Make sure the BQ table doesn't exist
     dropBqTableIfExists(dataset, INGESTION_TIME_PARTITIONED_TABLE_NAME);
     // Create the table using Hive
-    runHiveScript(HIVE_INGESTION_TIME_PARTITIONED_TABLE_CREATE_QUERY);
+    createManagedTableWithProps(
+        INGESTION_TIME_PARTITIONED_TABLE_NAME,
+        HIVE_INGESTION_TIME_PARTITIONED_DDL,
+        HIVE_INGESTION_TIME_PARTITIONED_PROPS);
     runHiveScript(
         String.format(
             "SELECT * from %s WHERE `_PARTITIONTIME` > TIMESTAMP'2018-09-05 00:10:04.19'",