Support for Spark Catalogs and ErrorIfExists save mode. (#19)

CONTRIBUTING.md

@@ -43,14 +43,75 @@ a GCP Project with a valid service account.
 For instructions on how to generate a service account and corresponding
 credentials JSON see: [Creating a Service Account][1].
 
-Then run the following to build, package, run all unit tests and run all
-integration tests.
+Then run the following to build, package, and run all integration tests.
+There are two profiles for different test suites:
+1. `integration`: for integration tests (`*IntegrationTest.java`)
+2. `acceptance`: for acceptance tests (`*AcceptanceTest.java`)
 
 ```bash
+# To run all integration tests:
 export GOOGLE_APPLICATION_CREDENTIALS=/path/to/service/account.json
-mvn -Penable-integration-tests clean verify
+mvn clean verify -Pintegration
+
+# To run all acceptance tests:
+export GOOGLE_APPLICATION_CREDENTIALS=/path/to/service/account.json
+mvn clean verify -Pacceptance
+```
+
+#### Integration Test Setup
+
+The integration tests require a Cloud Spanner instance and database. There are two ways to set this up:
+
+1.  **Automatic Test-Managed Database (for CI/temporary use):**
+    By default, the test suite will automatically create a temporary database with a unique name for the test run and delete it afterwards. It will also attempt to create a Spanner instance if one is not found. This is convenient for isolated test runs.
+
+2.  **User-Managed Database (for local development):**
+    For local development, it's often better to use a stable, long-lived database. The provided `scripts/setup-test-db.sh` script is the recommended way to do this.
+
+    The script will:
+    *   Create a Cloud Spanner instance if it doesn't exist.
+    *   Create (or drop and recreate) GoogleSQL and PostgreSQL-dialect databases.
+    *   Populate the databases with the required schema and test data.
+
+    To use this script, first set the required environment variables, then run it:
+    ```bash
+    ./scripts/setup-test-db.sh
+    ```
+
+    After running the script, you must set an additional environment variable to tell the test suite to use this pre-existing database and not create a temporary one:
+    ```bash
+    export SPANNER_USE_EXISTING_DATABASE=true
+    ```
+
+#### Environment Variables
+
+The integration tests are configured through the following environment variables:
+
+*   `SPANNER_PROJECT_ID`: Your Google Cloud project ID.
+*   `SPANNER_INSTANCE_ID`: The ID for the Cloud Spanner instance to use or create.
+*   `SPANNER_DATABASE_ID`: The base name for the test databases. The setup script will create a GoogleSQL database with this name and a PostgreSQL database named `${SPANNER_DATABASE_ID}-pg`.
+*   `GOOGLE_APPLICATION_CREDENTIALS`: (Optional if you have already authenticated with `gcloud auth application-default login`) Path to your service account credentials JSON file.
+*   `SPANNER_USE_EXISTING_DATABASE`: If set to `true`, the tests will use the database specified by `SPANNER_DATABASE_ID` and `SPANNER_INSTANCE_ID` without trying to create or delete it. This is useful for running tests against a database created with `scripts/setup-test-db.sh`.
+*   `SPANNER_EMULATOR_HOST`: (Optional) If you are running a local Cloud Spanner emulator, set this to the emulator's address (e.g., `localhost:9010`). Note that the emulator does not support PostgreSQL dialect tests.
+
+#### Running Specific Tests
+
+You can run a single integration test class or a specific test method using Maven.
+
+To run all tests in `Spark31WriteIntegrationTest`, which is in the `spark-3.1-spanner-lib` module:
+```bash
+mvn verify -Pintegration -Dtest=com.google.cloud.spark.spanner.integration.Spark31WriteIntegrationTest
 ```
 
+To run a specific method, like `testWriteWithNulls`, within that class:
+```bash
+mvn verify -Pintegration -Dtest=com.google.cloud.spark.spanner.integration.Spark31WriteIntegrationTest#testWriteWithNulls
+```
+
+*Note: Some test classes like `Spark31WriteIntegrationTest` may be subclasses that inherit tests from a base class. You should specify the concrete subclass in the test command.*
+
+The same principles apply to other test classes, modules, and profiles like `-Pacceptance`.
+
 ## Code Samples
 
 All code samples must be in compliance with the [java sample formatting guide][3].

examples/SpannerSpark.java

@@ -14,26 +14,70 @@
 
 package com.google.cloud.spark.spanner.examples;
 
+import java.util.Arrays;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.SaveMode;
 import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.DataTypes;
+import org.apache.spark.sql.types.StructType;
 
 public class SpannerSpark {
     public static void main(String[] args) {
+        String projectId = System.getenv("SPANNER_PROJECT_ID");
+        String instanceId = System.getenv("SPANNER_INSTANCE_ID");
+        String databaseId = System.getenv("SPANNER_DATABASE_ID");
+
+        // Configure a Spark session with a Spanner catalog named "spanner".
+        // This allows reading and writing tables with SQL and DataFrame APIs
+        // without specifying connection options on every operation.
         SparkSession spark = SparkSession
             .builder()
-            .appName("cloud spanner for census 2020")
+            .appName("SpannerSpark Example")
+            .config("spark.sql.catalog.spanner",
+                "com.google.cloud.spark.spanner.SpannerCatalog")
+            .config("spark.sql.catalog.spanner.projectId", projectId)
+            .config("spark.sql.catalog.spanner.instanceId", instanceId)
+            .config("spark.sql.catalog.spanner.databaseId", databaseId)
             .getOrCreate();
 
+        // --- Read via the catalog ---
+        // Tables are referenced as <catalog>.<table>.
+        Dataset<Row> df = spark.sql("SELECT * FROM spanner.people");
+        df.show();
+        df.printSchema();
 
-        Dataset<Row> df = spark.read()
+        // --- Read via the DataSource API (format) ---
+        // This approach still works and does not require catalog configuration.
+        Dataset<Row> df2 = spark.read()
             .format("cloud-spanner")
+            .option("projectId", projectId)
+            .option("instanceId", instanceId)
+            .option("databaseId", databaseId)
             .option("table", "people")
-            .option("projectId", System.getenv("SPANNER_SPARK_PROJECT"))
-            .option("instanceId", System.getenv("SPANNER_SPARK_INSTANCE"))
-            .option("database", System.getenv("SPANNER_SPARK_DATABASE"))
             .load();
-        df.show();
-        df.printSchema();
+        df2.show();
+
+        // --- Write via the DataSource API ---
+        StructType schema = new StructType()
+            .add("id", DataTypes.LongType, false)
+            .add("name", DataTypes.StringType, true)
+            .add("email", DataTypes.StringType, true);
+
+        Dataset<Row> writeDf = spark.createDataFrame(
+            Arrays.asList(
+                RowFactory.create(1L, "John Doe", "john@example.com"),
+                RowFactory.create(2L, "Jane Doe", "jane@example.com")),
+            schema);
+
+        writeDf.write()
+            .format("cloud-spanner")
+            .option("projectId", projectId)
+            .option("instanceId", instanceId)
+            .option("databaseId", databaseId)
+            .option("table", "people")
+            .mode(SaveMode.Append)
+            .save();
     }
 }

examples/SpannerSpark.py

@@ -1,4 +1,4 @@
-#/usr/bin/env python
+#!/usr/bin/env python
 
 # Copyright 2023 Google LLC. All Rights Reserved.
 #
@@ -14,20 +14,68 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from pyspark.sql import SparkSession
 
 def main():
-    table = "TABLE_NAME"
-    spark = SparkSession.builder.appName("SparkSpannerDemo").getOrCreate()
-    df = spark.read.format('cloud-spanner') \
-                .option("projectId", "<PROJECT_ID>") \
-                .option("instanceId", "<INSTANCE_ID>") \
-                .option("databaseId", "<DATABASE_ID>") \
-                .option("enableDataBoost", "true") \
-                .option("table", "<TABLE_NAME>") \
-                .load()
+    project_id = os.environ["SPANNER_PROJECT_ID"]
+    instance_id = os.environ["SPANNER_INSTANCE_ID"]
+    database_id = os.environ["SPANNER_DATABASE_ID"]
+
+    # Configure a Spark session with a Spanner catalog named "spanner".
+    # This allows reading and writing tables with SQL and DataFrame APIs
+    # without specifying connection options on every operation.
+    spark = (SparkSession.builder
+        .appName("SparkSpannerDemo")
+        .config("spark.sql.catalog.spanner",
+                "com.google.cloud.spark.spanner.SpannerCatalog")
+        .config("spark.sql.catalog.spanner.projectId", project_id)
+        .config("spark.sql.catalog.spanner.instanceId", instance_id)
+        .config("spark.sql.catalog.spanner.databaseId", database_id)
+        .getOrCreate())
+
+    # --- Read via the catalog ---
+    # Tables are referenced as <catalog>.<table>.
+    df = spark.sql("SELECT * FROM spanner.people")
     df.printSchema()
     df.show()
 
+    # --- Read via the DataSource API (format) ---
+    # This approach still works and does not require catalog configuration.
+    df2 = spark.read.format('cloud-spanner') \
+        .option("projectId", project_id) \
+        .option("instanceId", instance_id) \
+        .option("databaseId", database_id) \
+        .option("table", "people") \
+        .load()
+    df2.show()
+
+    # --- Write via the DataSource API ---
+    columns = ['id', 'name', 'email']
+    data = [(1, 'John Doe', 'john@example.com'),
+            (2, 'Jane Doe', 'jane@example.com')]
+    write_df = spark.createDataFrame(data, columns)
+
+    write_df.write.format('cloud-spanner') \
+        .option("projectId", project_id) \
+        .option("instanceId", instance_id) \
+        .option("databaseId", database_id) \
+        .option("table", "people") \
+        .mode("append") \
+        .save()
+
+    # --- Write via the catalog (SQL) ---
+    # INSERT INTO uses the catalog, so no connection options are needed.
+    spark.sql("""
+        INSERT INTO spanner.people
+        VALUES (3, 'Bob Smith', 'bob@example.com')
+    """)
+
+    # --- Write via the catalog (DataFrame) ---
+    # writeTo references the catalog table directly.
+    write_df2 = spark.createDataFrame(
+        [(4, 'Alice Wong', 'alice@example.com')], columns)
+    write_df2.writeTo("spanner.people").append()
+
 if __name__ == '__main__':
     main()