Add support for Spark Overwrite save mode

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java

@@ -24,6 +24,7 @@ public enum SpannerErrorCode {
   DECIMAL_OUT_OF_RANGE(6),
   INVALID_ARGUMENT(7),
   SCHEMA_VALIDATION_ERROR(8),
+  DDL_EXCEPTION(9),
   // Should be last
   UNSUPPORTED(9998),
   UNKNOWN(9999);

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java

@@ -51,9 +51,9 @@ public class SpannerTable implements Table, SupportsRead, SupportsWrite {
   private final SpannerTableSchema dbSchema;
   private final StructType sparkSchema;
   private static final ImmutableSet<TableCapability> tableCapabilities =
-      ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE);
+      ImmutableSet.of(
+          TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE);
   private final Map<String, String> properties;
-
   private static final Logger log = LoggerFactory.getLogger(SpannerTable.class);
 
   public SpannerTable(Map<String, String> properties) {

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java

@@ -64,6 +64,7 @@
 import org.apache.spark.sql.types.Decimal;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 import org.apache.spark.unsafe.types.UTF8String;
 import org.threeten.bp.Duration;
 
@@ -190,6 +191,21 @@ public static BatchClientWithCloser batchClientFromProperties(Map<String, String
 
   public static BatchClientWithCloser batchClientFromProperties(
       Map<String, String> properties, SessionPoolOptions sessionPoolOptions) {
+    SpannerOptions options = buildSpannerOptions(properties, sessionPoolOptions);
+    Spanner spanner = options.getService();
+    DatabaseId databaseId =
+        DatabaseId.of(
+            options.getProjectId(), properties.get("instanceId"), properties.get("databaseId"));
+    return new BatchClientWithCloser(
+        spanner, spanner.getBatchClient(databaseId), spanner.getDatabaseClient(databaseId));
+  }
+
+  public static SpannerOptions buildSpannerOptions(CaseInsensitiveStringMap properties) {
+    return buildSpannerOptions(properties, defaultSessionPoolOptions);
+  }
+
+  public static SpannerOptions buildSpannerOptions(
+      Map<String, String> properties, SessionPoolOptions sessionPoolOptions) {
     SpannerOptions.Builder builder =
         SpannerOptions.newBuilder()
             .setSessionPoolOption(sessionPoolOptions)
@@ -222,13 +238,7 @@ public static BatchClientWithCloser batchClientFromProperties(
     }
     System.setProperty("com.google.cloud.spanner.watchdogTimeoutSeconds", "7200");
 
-    SpannerOptions options = builder.build();
-    Spanner spanner = options.getService();
-    DatabaseId databaseId =
-        DatabaseId.of(
-            options.getProjectId(), properties.get("instanceId"), properties.get("databaseId"));
-    return new BatchClientWithCloser(
-        spanner, spanner.getBatchClient(databaseId), spanner.getDatabaseClient(databaseId));
+    return builder.build();
   }
 
   public static void toSparkDecimal(GenericInternalRow dest, java.math.BigDecimal v, int at) {

spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java

@@ -14,11 +14,18 @@
 
 package com.google.cloud.spark.spanner;
 
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.Spanner;
+import com.google.cloud.spanner.SpannerException;
+import com.google.cloud.spanner.Statement;
 import org.apache.spark.sql.connector.write.BatchWrite;
 import org.apache.spark.sql.connector.write.LogicalWriteInfo;
+import org.apache.spark.sql.connector.write.SupportsTruncate;
 import org.apache.spark.sql.connector.write.WriteBuilder;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 
-public class SpannerWriteBuilder implements WriteBuilder {
+public class SpannerWriteBuilder implements WriteBuilder, SupportsTruncate {
   private final LogicalWriteInfo info;
 
   public SpannerWriteBuilder(LogicalWriteInfo info) {
@@ -29,4 +36,48 @@ public SpannerWriteBuilder(LogicalWriteInfo info) {
   public BatchWrite buildForBatch() {
     return new SpannerBatchWrite(info);
   }
+
+  @Override
+  public WriteBuilder truncate() {
+    CaseInsensitiveStringMap opts = new CaseInsensitiveStringMap(this.info.options());
+    String projectId = SpannerUtils.getRequiredOption(opts, "projectId");
+    String instanceId = SpannerUtils.getRequiredOption(opts, "instanceId");
+    String databaseId = SpannerUtils.getRequiredOption(opts, "databaseId");
+    String tableName = SpannerUtils.getRequiredOption(opts, "table");
+
+    try (Spanner spanner = SpannerUtils.buildSpannerOptions(opts).getService()) {
+      DatabaseClient dbClient =
+          spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId));
+      truncateTable(dbClient, tableName);
+    } catch (Exception e) {
+      throw new SpannerConnectorException(
+          SpannerErrorCode.DDL_EXCEPTION, "Error truncating table " + tableName, e);
+    }
+    return this;
+  }
+
+  public static long truncateTable(DatabaseClient dbClient, String tableName) {
+
+    // 1. Construct the DML Statement
+    // Spanner requires a WHERE clause for PDML, even if you are deleting everything.
+    String sql = "DELETE FROM " + tableName + " WHERE true";
+    Statement statement = Statement.of(sql);
+
+    System.out.println("Starting Partitioned DML execution: " + sql);
+
+    try {
+      // 2. Execute the Partitioned Update
+      // This is a blocking call. The Spanner client will divide the table into
+      // partitions and run concurrent background transactions to delete the data.
+      long deletedRowCount = dbClient.executePartitionedUpdate(statement);
+
+      System.out.println("Successfully deleted " + deletedRowCount + " rows.");
+      return deletedRowCount;
+
+    } catch (SpannerException e) {
+      // SpannerExceptions wrap underlying gRPC errors (e.g., DEADLINE_EXCEEDED, PERMISSION_DENIED)
+      System.err.println("Failed to execute Partitioned DML on table: " + tableName);
+      throw e;
+    }
+  }
 }