Write acceptance tests

spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocAcceptanceTestBase.java

@@ -63,7 +63,6 @@ public class DataprocAcceptanceTestBase {
       Preconditions.checkNotNull(
           System.getenv("SPANNER_INSTANCE_ID"),
           "Please set the 'SPANNER_INSTANCE_ID' environment variable");
-  private static final String TABLE = "ATable";
   private static Spanner spanner =
       SpannerOptions.newBuilder().setProjectId(PROJECT_ID).build().getService();
   private static final Logger logger = LoggerFactory.getLogger(DataprocAcceptanceTestBase.class);
@@ -89,6 +88,22 @@ public void testRead() throws Exception {
     assertThat(output.trim()).isEqualTo("41");
   }
 
+  @Test
+  public void testWrite() throws Exception {
+    logger.info("testWrite started");
+    String testName = "test-write";
+    Job result =
+        createAndRunPythonJob(
+            testName,
+            "write_test_table.py",
+            null,
+            Arrays.asList(context.getResultsDirUri(testName), PROJECT_ID, INSTANCE_ID, DATABASE_ID),
+            120);
+    assertThat(result.getStatus().getState()).isEqualTo(JobStatus.State.DONE);
+    String output = AcceptanceTestUtils.getCsv(context.getResultsDirUri(testName));
+    assertThat(output.trim()).startsWith("PASS");
+  }
+
   private Job createAndRunPythonJob(
       String testName, String pythonFile, String pythonZipUri, List<String> args, long duration)
       throws Exception {

spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocServerlessAcceptanceTestBase.java

@@ -15,6 +15,7 @@
 package com.google.cloud.spark.spanner.acceptance;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertTrue;
 
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.longrunning.OperationSnapshot;
@@ -47,6 +48,7 @@
 import java.util.Locale;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -90,14 +92,15 @@ public class DataprocServerlessAcceptanceTestBase {
   String testId = String.format("%s-%s", testName, System.currentTimeMillis());
   String testBaseGcsDir = AcceptanceTestUtils.createTestBaseGcsDir(testId);
   String connectorJarUri = testBaseGcsDir + "/connector.jar";
-  AcceptanceTestContext context =
-      new AcceptanceTestContext(
-          testId, generateClusterName(testId), testBaseGcsDir, connectorJarUri);
+  static AcceptanceTestContext context;
 
   private final String s8sImageVersion;
   private final String connectorJarDirectory;
   private final String connectorJarPrefix;
 
+  private static Exception initializationException;
+  private static boolean initialized = false;
+
   public DataprocServerlessAcceptanceTestBase(
       String connectorJarDirectory, String connectorJarPrefix, String s8sImageVersion) {
     this.connectorJarDirectory = connectorJarDirectory;
@@ -107,9 +110,24 @@ public DataprocServerlessAcceptanceTestBase(
 
   @Before
   public void createBatchControllerClient() throws Exception {
-    AcceptanceTestUtils.uploadConnectorJar(
-        connectorJarDirectory, connectorJarPrefix, context.connectorJarUri);
-    createSpannerDataset();
+    if (initializationException != null) {
+      throw initializationException;
+    }
+
+    // Lazy-initialize exactly once
+    if (!initialized) {
+      try {
+        context =
+            new AcceptanceTestContext(
+                testId, generateClusterName(testId), testBaseGcsDir, connectorJarUri);
+        AcceptanceTestUtils.uploadConnectorJar(
+            connectorJarDirectory, connectorJarPrefix, context.connectorJarUri);
+        createSpannerDataset();
+      } catch (Exception e) {
+        initializationException = e;
+        throw e;
+      }
+    }
 
     batchController =
         BatchControllerClient.create(
@@ -119,8 +137,20 @@ public void createBatchControllerClient() throws Exception {
   @After
   public void tearDown() throws Exception {
     batchController.close();
-    AcceptanceTestUtils.deleteGcsDir(context.testBaseGcsDir);
-    deleteSpannerDatasetAndTables();
+  }
+
+  @AfterClass
+  public static void tearDownOnce() {
+    if (initialized) {
+      try {
+        AcceptanceTestUtils.deleteGcsDir(context.testBaseGcsDir);
+      } catch (Exception e) {
+        System.err.println("Cleanup failed: " + e.getMessage());
+      }
+      deleteSpannerDatasetAndTables();
+      context = null;
+      initialized = false;
+    }
   }
 
   @Test
@@ -139,6 +169,22 @@ public void testBatch() throws Exception {
     assertThat(output.trim()).isEqualTo("41");
   }
 
+  @Test
+  public void testWrite() throws Exception {
+    OperationSnapshot operationSnapshot =
+        createAndRunPythonBatch(
+            context,
+            testName,
+            "write_test_table.py",
+            null,
+            Arrays.asList(
+                context.getResultsDirUri(testName), PROJECT_ID, INSTANCE_ID, DATABASE_ID));
+    assertThat(operationSnapshot.isDone()).isTrue();
+    assertThat(operationSnapshot.getErrorMessage()).isEmpty();
+    String output = AcceptanceTestUtils.getCsv(context.getResultsDirUri(testName));
+    assertTrue(output.trim().startsWith("PASS"));
+  }
+
   protected static void createSpannerDataset() throws Exception {
     // 1. Create the Spanner instance.
     InstanceAdminClient instanceAdminClient = spanner.getInstanceAdminClient();

spark-3.1-spanner-lib/src/test/resources/acceptance/write_test_table.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python
+# Copyright 2023 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from datetime import datetime, date
+from decimal import Decimal
+from pyspark.sql import SparkSession, Row
+from pyspark.sql.types import StructType, StructField, StringType, LongType, BinaryType, TimestampType, DecimalType, BooleanType, DoubleType, DateType
+
+def main():
+
+    # Initialize Spark Session
+    spark = SparkSession.builder.appName('Write Acceptance Test on Spark').getOrCreate()
+
+    # 1. Define the Schema (Column Name, Type, Nullable)
+    schema = StructType([
+        StructField("A", LongType(), False),
+        StructField("B", StringType(), True),
+        StructField("C", BinaryType(), True),
+        StructField("D", TimestampType(), True),
+        StructField("E", DecimalType(38, 9), True),
+        StructField("F", BooleanType(), True),
+        StructField("G", DoubleType(), True),
+        StructField("H", DateType(), True)
+    ])
+
+    # 2. Prepare Data as a list of tuples
+    data = [
+        (1,  "2",  None, datetime.fromisoformat("2023-08-22T12:22:00"), Decimal("1000.282111401"), True, 123.456, date(2023, 12, 25)),
+        (10, "20", None, datetime.fromisoformat("2023-08-22T12:23:00"), Decimal("10000.282111603"), False, 987.654, date(2023, 12, 24)),
+        (30, "30", None, datetime.fromisoformat("2023-08-22T12:24:00"), Decimal("30000.282111805"), True, -2121.1212, date(2023, 12, 23))
+    ]
+
+    # 3. Create the DataFrame
+    dfw = spark.createDataFrame(data, schema)
+    dfw.show()
+
+    table = 'AWriteTable'
+
+    # Configure Spanner properties
+    spanner_base_options = {
+        "instanceId": sys.argv[3],
+        "databaseId": sys.argv[4],
+        "projectId": sys.argv[2],
+        "table": table
+    }
+
+    spanner_write_options = {
+        **spanner_base_options,
+        "mutationType": "insert_or_update", # Use this to avoid ALREADY_EXISTS errors
+        "enablePartialRowUpdates": "true"   # Required since not all columns are being populated
+    }
+
+    spanner_read_options = {
+        **spanner_base_options
+    }
+
+    dfw.write.format('cloud-spanner') \
+        .options(**spanner_write_options) \
+        .mode("append") \
+        .save()
+
+    # Read the table to verify the write operation
+    df = spark.read.format('cloud-spanner') \
+      .options(**spanner_read_options) \
+      .load(table)
+
+    print('The resulting schema is')
+    df.printSchema()
+    df.show()
+
+    df_result = verify_data_to_df(dfw, df, spark)
+    df_result.show()
+
+    # coalesce 1 to ensure results are written in single partition and avoid empty file creation.
+    df_result.coalesce(1).write.csv(sys.argv[1])
+
+
+def verify_data_to_df(df_expected, df_actual, spark):
+    issues = []
+
+    # 1. Validation Logic
+    if df_expected.schema != df_actual.schema:
+        issues.append("Schema mismatch")
+
+    # Cache DFs for performance since we'll perform multiple actions.
+    df_expected.cache()
+    df_actual.cache()
+
+    missing_rows_count = df_expected.subtract(df_actual).count()
+    if missing_rows_count > 0:
+        issues.append(f"Missing rows in actual: {missing_rows_count}")
+
+    extra_rows_count = df_actual.subtract(df_expected).count()
+    if extra_rows_count > 0:
+        issues.append(f"Extra rows in actual: {extra_rows_count}")
+
+    df_expected.unpersist()
+    df_actual.unpersist()
+
+    # 2. Determine Final Status
+    status_msg = "PASS" if not issues else "FAIL: " + " | ".join(issues)
+
+    # 3. Create a DataFrame from the result string
+    return spark.createDataFrame([Row(summary=status_msg)])
+
+if __name__ == '__main__':
+  main()

spark-3.1-spanner-lib/src/test/resources/db/populate_ddl.sql

@@ -48,6 +48,17 @@ CREATE TABLE ATable(
   G JSON
 ) PRIMARY KEY(A);
 
+CREATE TABLE AWriteTable (
+  A INT64 NOT NULL,
+  B STRING(100),
+  C BYTES(MAX),
+  D TIMESTAMP,
+  E NUMERIC,
+  F BOOL,
+  G FLOAT64,
+  H DATE
+) PRIMARY KEY(A);
+
 CREATE TABLE compositeTable (
   id STRING(36) NOT NULL,
   A ARRAY<INT64>,