diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f349f7af..2d480bc7 100644 --- a/CONTRIBUTING.md +++ b/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]. diff --git a/examples/SpannerSpark.java b/examples/SpannerSpark.java index 645ac50f..9f5c90ce 100644 --- a/examples/SpannerSpark.java +++ b/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 .. + Dataset df = spark.sql("SELECT * FROM spanner.people"); + df.show(); + df.printSchema(); - Dataset df = spark.read() + // --- Read via the DataSource API (format) --- + // This approach still works and does not require catalog configuration. + Dataset 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 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(); } } diff --git a/examples/SpannerSpark.py b/examples/SpannerSpark.py index ea770acc..8b51d7cf 100644 --- a/examples/SpannerSpark.py +++ b/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", "") \ - .option("instanceId", "") \ - .option("databaseId", "") \ - .option("enableDataBoost", "true") \ - .option("table", "") \ - .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 .
. + 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() diff --git a/scripts/setup-test-db.sh b/scripts/setup-test-db.sh new file mode 100755 index 00000000..85aedfb8 --- /dev/null +++ b/scripts/setup-test-db.sh @@ -0,0 +1,174 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# +# 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. + +set -e + +# This script automates the setup of GoogleSQL and PostgreSQL test databases +# for Spark-Spanner connector integration tests. It performs the following actions: +# +# 1. Environment Variable Check: Verifies the presence of essential variables: +# - SPANNER_PROJECT_ID: Your Google Cloud project ID. +# - SPANNER_INSTANCE_ID: The ID for the Cloud Spanner instance. +# - SPANNER_DATABASE_ID: The base name for the test databases. +# +# 2. Command-Line Tool Check: Ensures that the 'gcloud' CLI is installed and +# accessible in the system's PATH. +# +# 3. Instance Management: +# - Checks if the specified Cloud Spanner instance exists. +# - If the instance is not found, it creates a new one with a default +# configuration (1 node in regional-us-central1). +# +# 4. Database Management: +# - Uses SPANNER_DATABASE_ID directly for the GoogleSQL database. +# - Creates a PostgreSQL database with a "-pg" suffix based on SPANNER_DATABASE_ID. +# - For both databases: +# - If the database exists, it drops and recreates it to ensure a clean state. +# - If the database does not exist, it creates it. +# +# 5. Schema and Data Population: +# - Applies the initial DDL schema to create tables from .sql files. +# - Executes DML statements to populate the tables with test data from .sql files. +# + +# Function to display an error message and exit +fail() { + echo "$1" >&2 + exit 1 +} + +# 1. Check for required environment variables +[[ -z "$SPANNER_PROJECT_ID" ]] && fail "Error: SPANNER_PROJECT_ID is not set." +[[ -z "$SPANNER_INSTANCE_ID" ]] && fail "Error: SPANNER_INSTANCE_ID is not set." +[[ -z "$SPANNER_DATABASE_ID" ]] && fail "Error: SPANNER_DATABASE_ID is not set." + +# 2. Check for gcloud CLI +command -v gcloud >/dev/null 2>&1 || fail "Error: gcloud command not found. Please install the Google Cloud SDK." + +# 3. Define database IDs (no randomization, ignore SPANNER_USE_EXISTING_DATABASE) +DATABASE_ID="$SPANNER_DATABASE_ID" +DATABASE_ID_PG="${DATABASE_ID}-pg" + +# 4. Check and create Spanner instance if it doesn't exist +echo "Checking for Spanner instance '$SPANNER_INSTANCE_ID'..." +if ! gcloud spanner instances describe "$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" &>/dev/null; then + echo "Instance not found. Creating a new instance..." + gcloud spanner instances create "$SPANNER_INSTANCE_ID" \ + --project="$SPANNER_PROJECT_ID" \ + --config="regional-us-central1" \ + --description="Test Instance for Spark-Spanner Connector" \ + --nodes=1 + echo "Instance '$SPANNER_INSTANCE_ID' created." +else + echo "Instance '$SPANNER_INSTANCE_ID' already exists." +fi + +# 5. Define DDL and DML file paths +DB_RESOURCES_DIR="spark-3.1-spanner-lib/src/test/resources/db" +DDL_GOOGLESOL_FILES=( + "$DB_RESOURCES_DIR/populate_ddl.sql" + "$DB_RESOURCES_DIR/populate_ddl_graph.sql" +) +DML_GOOGLESOL_FILES=( + "$DB_RESOURCES_DIR/insert_data.sql" + "$DB_RESOURCES_DIR/insert_data_graph.sql" +) +DDL_PG_FILES=("$DB_RESOURCES_DIR/populate_ddl_pg.sql") +DML_PG_FILES=("$DB_RESOURCES_DIR/insert_data_pg.sql") + +# Function to execute a multi-statement DML file +execute_dml_file() { + local db_id=$1 + local file=$2 + echo "Executing DML file: $file" + + local content + content=$(<"$file") + # Remove comments and newlines to allow splitting by semicolon. + local statements + statements=$(echo "$content" | grep -v '^--' | tr '\n' ' ' | sed 's/;/;\n/g') + + while IFS= read -r stmt; do + # Check for non-empty statements + if [[ -n "$(echo "$stmt" | tr -d '[:space:]')" ]]; then + echo "Executing DML: ${stmt:0:80}..." + + local max_retries=5 + local attempt=1 + local delay=1 + while true; do + if gcloud spanner databases execute-sql "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" --sql="$stmt"; then + # Success + break + else + # Failure + if [[ $attempt -eq $max_retries ]]; then + echo "DML statement failed after $max_retries attempts. Aborting." + exit 1 + fi + echo "DML statement failed. Retrying in ${delay}s... (Attempt ${attempt}/${max_retries})" + sleep $delay + # Exponential backoff + delay=$((delay * 2)) + attempt=$((attempt + 1)) + fi + done + fi + done <<< "$statements" +} + +# Function to manage and populate a database +setup_database() { + local db_id=$1 + local dialect=$2 + local ddl_files_str=$3 + local dml_files_str=$4 + + echo "Setting up $dialect database '$db_id'..." + + if gcloud spanner databases describe "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" &>/dev/null; then + echo "Database '$db_id' already exists. Dropping and recreating to ensure clean state..." + gcloud spanner databases delete "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" --quiet + gcloud spanner databases create "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" --database-dialect="$dialect" + else + echo "Database '$db_id' not found. Creating..." + gcloud spanner databases create "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" --database-dialect="$dialect" + fi + + # Apply DDL + local ddl_files_arr + read -r -a ddl_files_arr <<< "$ddl_files_str" + for ddl in "${ddl_files_arr[@]}"; do + echo "Applying DDL from: $ddl" + gcloud spanner databases ddl update "$db_id" --instance="$SPANNER_INSTANCE_ID" --project="$SPANNER_PROJECT_ID" --ddl-file="$ddl" + done + + # Apply DML + local dml_files_arr + read -r -a dml_files_arr <<< "$dml_files_str" + for dml_file in "${dml_files_arr[@]}"; do + execute_dml_file "$db_id" "$dml_file" + done + + echo "$dialect database '$db_id' setup complete." +} + +# Setup GoogleSQL database +setup_database "$DATABASE_ID" "GOOGLE_STANDARD_SQL" "${DDL_GOOGLESOL_FILES[*]}" "${DML_GOOGLESOL_FILES[*]}" + +# Setup PostgreSQL database +setup_database "$DATABASE_ID_PG" "POSTGRESQL" "${DDL_PG_FILES[*]}" "${DML_PG_FILES[*]}" + +echo "All test databases have been successfully configured." diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerBatchWrite.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerBatchWrite.java index 0e5ccd0a..4a9394db 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerBatchWrite.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerBatchWrite.java @@ -19,6 +19,7 @@ import org.apache.spark.sql.connector.write.LogicalWriteInfo; import org.apache.spark.sql.connector.write.PhysicalWriteInfo; import org.apache.spark.sql.connector.write.WriterCommitMessage; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -26,16 +27,17 @@ public class SpannerBatchWrite implements BatchWrite { private static final Logger log = LoggerFactory.getLogger(SpannerBatchWrite.class); private final LogicalWriteInfo info; + private final CaseInsensitiveStringMap properties; - public SpannerBatchWrite(LogicalWriteInfo info) { + public SpannerBatchWrite(LogicalWriteInfo info, CaseInsensitiveStringMap properties) { this.info = info; + this.properties = properties; log.info("Creating SpannerBatchWrite for queryId {}", info.queryId()); } @Override public DataWriterFactory createBatchWriterFactory(PhysicalWriteInfo info) { - return new SpannerDataWriterFactory( - this.info.options().asCaseSensitiveMap(), this.info.schema()); + return new SpannerDataWriterFactory(this.properties, this.info.schema()); } @Override diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java new file mode 100644 index 00000000..fd991e1b --- /dev/null +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerCatalog.java @@ -0,0 +1,281 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner; + +import com.google.api.gax.longrunning.OperationFuture; +import com.google.cloud.spanner.DatabaseAdminClient; +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.ReadContext; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.Statement; +import com.google.cloud.spark.spanner.graph.SpannerGraphBuilder; +import com.google.common.base.Verify; +import com.google.gson.Gson; +import com.google.gson.JsonSyntaxException; +import com.google.gson.reflect.TypeToken; +import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.connector.catalog.TableChange; +import org.apache.spark.sql.connector.expressions.Transform; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.MetadataBuilder; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class SpannerCatalog implements TableCatalog, AutoCloseable { + public static final String GRAPH_IDENTIFIER_PREFIX = "__spanner_graph__"; + private static final Gson GSON = new Gson(); + + public static final Metadata PRIMARY_KEY_METADATA = + new MetadataBuilder().putBoolean(SpannerUtils.PRIMARY_KEY_TAG, true).build(); + private static final Logger log = LoggerFactory.getLogger(SpannerCatalog.class); + private String catalogName; + private CaseInsensitiveStringMap options; + private Spanner spanner; + private String projectId; + private String instanceId; + private String databaseId; + + // For testing purposes. + protected Spanner createSpanner(CaseInsensitiveStringMap options) { + return SpannerUtils.buildSpannerOptions(options).getService(); + } + + // For testing purposes. + protected SpannerInformationSchema createSchemaInfo(Dialect dialect) { + return SpannerInformationSchema.create(dialect); + } + + @Override + public void initialize(String name, CaseInsensitiveStringMap options) { + this.catalogName = name; + this.options = options; + this.projectId = SpannerUtils.getRequiredOption(options, "projectId"); + this.instanceId = SpannerUtils.getRequiredOption(options, "instanceId"); + this.databaseId = SpannerUtils.getRequiredOption(options, "databaseId"); + this.spanner = createSpanner(options); + } + + @Override + public String name() { + return catalogName; + } + + @Override + public Identifier[] listTables(String[] namespace) { + if (namespace.length > 0) { + log.warn("Invalid namespace for listing tables: {}", String.join(".", namespace)); + return new Identifier[0]; + } + + DatabaseClient dbClient = getDatabaseClient(); + + try (ReadContext readContext = dbClient.readOnlyTransaction()) { + Dialect dialect = dbClient.getDialect(); + return createSchemaInfo(dialect).listTables(readContext, namespace); + } catch (Exception e) { + log.error( + "Error listing tables in namespace {}: {}", String.join(".", namespace), e.getMessage()); + return new Identifier[0]; + } + } + + @Override + public Table loadTable(Identifier ident) throws NoSuchTableException { + if (ident.namespace().length != 0) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Invalid identifier namespace: " + String.join(".", ident.namespace())); + } + if (isGraphIdentifier(ident)) { + return factorySpannerGraph(ident); + } + if (!tableExists(ident)) { + throw new NoSuchTableException(ident); + } + return factorySpannerTable(ident); + } + + public static boolean isGraphIdentifier(Identifier ident) { + return ident.name().startsWith(GRAPH_IDENTIFIER_PREFIX); + } + + protected Table factorySpannerTable(Identifier ident) { + String table = ident.name(); + Verify.verifyNotNull(table, "table"); + + return new SpannerTable(projectId, instanceId, databaseId, table, options, null); + } + + protected Table factorySpannerGraph(Identifier ident) { + String json = ident.name().substring(GRAPH_IDENTIFIER_PREFIX.length()); + if (json.isEmpty()) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, "Graph identifier has no encoded properties"); + } + Map graphProps; + try { + graphProps = GSON.fromJson(json, new TypeToken>() {}.getType()); + } catch (JsonSyntaxException e) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Malformed graph identifier JSON: " + e.getMessage(), + e); + } + if (graphProps == null) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, "Graph identifier decoded to null"); + } + Map allOptions = new HashMap<>(options.asCaseSensitiveMap()); + for (String key : SparkSpannerTableProviderBase.GRAPH_OPTION_KEYS) { + String val = graphProps.get(key); + if (val != null) { + allOptions.put(key, val); + } + } + return SpannerGraphBuilder.build(allOptions); + } + + static StructType enrichSchemaWithPrimaryKeys(StructType schema, Map properties) { + String pkCols = properties.get("primaryKeys"); + if (pkCols == null || pkCols.isEmpty()) { + return schema; + } + java.util.Set pkSet = new java.util.HashSet<>(); + for (String col : pkCols.split(",")) { + pkSet.add(col.trim()); + } + StructField[] enrichedFields = new StructField[schema.fields().length]; + for (int i = 0; i < schema.fields().length; i++) { + StructField field = schema.fields()[i]; + if (pkSet.contains(field.name())) { + Metadata merged = + new MetadataBuilder() + .withMetadata(field.metadata()) + .putBoolean(SpannerUtils.PRIMARY_KEY_TAG, true) + .build(); + enrichedFields[i] = + new StructField(field.name(), field.dataType(), field.nullable(), merged); + } else { + enrichedFields[i] = field; + } + } + return new StructType(enrichedFields); + } + + @Override + public Table createTable( + Identifier ident, StructType schema, Transform[] partitions, Map properties) { + + DatabaseClient dbClient = getDatabaseClient(); + Dialect dialect = dbClient.getDialect(); + SpannerInformationSchema schemaInfo = createSchemaInfo(dialect); + StructType enrichedSchema = enrichSchemaWithPrimaryKeys(schema, properties); + String ddl = schemaInfo.createTableDdl(ident, enrichedSchema); + DatabaseAdminClient dbAdminClient = spanner.getDatabaseAdminClient(); + OperationFuture op = + dbAdminClient.updateDatabaseDdl( + instanceId, databaseId, Collections.singletonList(ddl), null); + + try { + op.get(); + } catch (ExecutionException | InterruptedException e) { + throw new SpannerConnectorException( + SpannerErrorCode.DDL_EXCEPTION, + "Exception while creating table " + ident.name() + ": " + e.getMessage(), + e); + } + + return factorySpannerTable(ident); + } + + @Override + public boolean tableExists(Identifier ident) { + if (ident.namespace().length > 0) { + log.warn("Invalid namespace for listing tables: {}", String.join(".", ident.namespace())); + return false; + } + + String tableName = ident.name(); + + DatabaseClient dbClient = getDatabaseClient(); + + try (ReadContext readContext = dbClient.singleUse()) { + Statement statement = createSchemaInfo(dbClient.getDialect()).tableExistsStatement(tableName); + try (ResultSet resultSet = readContext.executeQuery(statement)) { + return resultSet.next() && resultSet.getLong(0) > 0; + } + } catch (Exception e) { + log.error("Error checking table existence {}: {}", tableName, e.getMessage()); + return false; + } + } + + @Override + public Table alterTable(Identifier ident, TableChange... changes) { + throw new UnsupportedOperationException("ALTER TABLE is not supported for SpannerCatalog"); + } + + @Override + public boolean dropTable(Identifier ident) { + DatabaseClient dbClient = getDatabaseClient(); + SpannerInformationSchema schemaInfo = createSchemaInfo(dbClient.getDialect()); + + DatabaseAdminClient dbAdminClient = spanner.getDatabaseAdminClient(); + OperationFuture op = + dbAdminClient.updateDatabaseDdl( + instanceId, + databaseId, + Collections.singletonList(schemaInfo.dropTableDdl(ident.name())), + null); + + try { + op.get(); + return true; + } catch (ExecutionException | InterruptedException e) { + log.warn("Exception while dropping table {}: {}", ident.name(), e.getMessage(), e); + } + return false; + } + + @Override + public void renameTable(Identifier oldIdent, Identifier newIdent) { + throw new UnsupportedOperationException("RENAME TABLE is not supported for SpannerCatalog"); + } + + private DatabaseClient getDatabaseClient() { + return spanner.getDatabaseClient(DatabaseId.of(projectId, instanceId, databaseId)); + } + + @Override + public void close() { + if (spanner != null) { + spanner.close(); + } + } +} diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriter.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriter.java index 20f31479..299a31ee 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriter.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriter.java @@ -38,6 +38,7 @@ import org.apache.spark.sql.types.StringType; 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.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -89,25 +90,28 @@ public SpannerDataWriter( ScheduledExecutorService scheduler) { this.partitionId = partitionId; this.taskId = taskId; - this.tableName = SpannerUtils.getRequiredOption(properties, "table"); + + CaseInsensitiveStringMap caseInsensitiveStringMap = new CaseInsensitiveStringMap(properties); + this.tableName = SpannerUtils.getRequiredOption(caseInsensitiveStringMap, "table"); this.schema = schema; // Default to 1MB (Safety) and 1000 Mutations this.mutationsPerTransaction = Integer.parseInt( - properties.getOrDefault( + caseInsensitiveStringMap.getOrDefault( "mutationsPerTransaction", MUTATIONS_PER_TRANSACTION_DEFAULT_STR)); this.bytesPerTransaction = Long.parseLong( - properties.getOrDefault( + caseInsensitiveStringMap.getOrDefault( "bytesPerTransaction", BYTES_PER_TRANSACTION_DEFAULT_STR)); // 1 MB default this.assumeIdempotentRows = Boolean.parseBoolean( - properties.getOrDefault("assumeIdempotentRows", ASSUME_IDEMPOTENT_ROWS_DEFAULT_STR)); + caseInsensitiveStringMap.getOrDefault( + "assumeIdempotentRows", ASSUME_IDEMPOTENT_ROWS_DEFAULT_STR)); this.maxPendingTransactions = Integer.parseInt( - properties.getOrDefault( + caseInsensitiveStringMap.getOrDefault( "maxPendingTransactions", MAX_PENDING_TRANSACTIONS_DEFAULT_STR)); this.mutationType = parseMutationType(properties.getOrDefault("mutationType", null)); diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriterFactory.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriterFactory.java index 083d47b8..9b125f17 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriterFactory.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerDataWriterFactory.java @@ -23,13 +23,16 @@ import org.apache.spark.sql.connector.write.DataWriter; import org.apache.spark.sql.connector.write.DataWriterFactory; import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; public class SpannerDataWriterFactory implements DataWriterFactory { + // Instances of this class are marshaled to worker nodes so all fields need to be serializable. + // CaseInsensitiveMap is not. private final Map properties; private final StructType schema; - public SpannerDataWriterFactory(Map properties, StructType schema) { - this.properties = properties; + public SpannerDataWriterFactory(CaseInsensitiveStringMap properties, StructType schema) { + this.properties = properties.asCaseSensitiveMap(); this.schema = schema; } @@ -39,7 +42,8 @@ public DataWriter createWriter(int partitionId, long taskId) { SessionPoolOptions sessionPoolOptions = SessionPoolOptions.newBuilder().setMinSessions(1).setMaxSessions(numThreads).build(); BatchClientWithCloser batchClient = - SpannerUtils.batchClientFromProperties(properties, sessionPoolOptions); + SpannerUtils.batchClientFromProperties( + new CaseInsensitiveStringMap(properties), sessionPoolOptions); ExecutorService executor = Executors.newFixedThreadPool(numThreads); ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java index 97b40c0c..1e64304b 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerErrorCode.java @@ -24,6 +24,8 @@ public enum SpannerErrorCode { DECIMAL_OUT_OF_RANGE(6), INVALID_ARGUMENT(7), SCHEMA_VALIDATION_ERROR(8), + DDL_EXCEPTION(9), + UNSUPPORTED_DATATYPE(10), // Should be last UNSUPPORTED(9998), UNKNOWN(9999); diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java new file mode 100644 index 00000000..adfcbaab --- /dev/null +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInformationSchema.java @@ -0,0 +1,233 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner; + +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.ReadContext; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.types.ArrayType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +public interface SpannerInformationSchema { + Identifier[] listTables(ReadContext readContext, String[] namespace); + + Statement tableExistsStatement(String tableName); + + String quoteIdentifier(String identifier); + + String sparkTypeToSpannerType(StructField field); + + default String createTableDdl(Identifier ident, StructType schema) { + StringBuilder ddl = new StringBuilder(); + ddl.append("CREATE TABLE ").append(quoteIdentifier(ident.name())).append(" ("); + for (StructField field : schema.fields()) { + ddl.append(quoteIdentifier(field.name())).append(" ").append(sparkTypeToSpannerType(field)); + if (!field.nullable()) { + ddl.append(" NOT NULL"); + } + ddl.append(", "); + } + + List primaryKeys = + Arrays.stream(schema.fields()) + .filter( + f -> + f.metadata().contains(SpannerUtils.PRIMARY_KEY_TAG) + && f.metadata().getBoolean(SpannerUtils.PRIMARY_KEY_TAG)) + .map(f -> quoteIdentifier(f.name())) + .collect(Collectors.toList()); + + if (primaryKeys.isEmpty()) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "No primary key found for table " + + ident.name() + + ". Please specify at least one primary key column."); + } + + ddl.append("PRIMARY KEY (").append(String.join(", ", primaryKeys)).append(")"); + ddl.append(")"); + return ddl.toString(); + } + + default String dropTableDdl(String tableName) { + return "DROP TABLE " + quoteIdentifier(tableName); + } + + static SpannerInformationSchema create(Dialect dialect) { + switch (dialect) { + case POSTGRESQL: + return new PostgresSpannerInformationSchema(); + case GOOGLE_STANDARD_SQL: + return new GoogleSqlSpannerInformationSchema(); + } + throw new IllegalArgumentException("Unsupported dialect: " + dialect); + } +} + +class GoogleSqlSpannerInformationSchema implements SpannerInformationSchema { + @Override + public String quoteIdentifier(String identifier) { + return "`" + identifier.replace("`", "``") + "`"; + } + + @Override + public String sparkTypeToSpannerType(StructField field) { + if (field.dataType().equals(DataTypes.LongType)) { + return "INT64"; + } + if (field.dataType().equals(DataTypes.StringType)) { + return "STRING(MAX)"; + } + if (field.dataType().equals(DataTypes.BooleanType)) { + return "BOOL"; + } + if (field.dataType().equals(DataTypes.DoubleType)) { + return "FLOAT64"; + } + if (field.dataType().equals(DataTypes.BinaryType)) { + return "BYTES(MAX)"; + } + if (field.dataType().equals(DataTypes.TimestampType)) { + return "TIMESTAMP"; + } + if (field.dataType().equals(DataTypes.DateType)) { + return "DATE"; + } + if (field.dataType() instanceof org.apache.spark.sql.types.DecimalType) { + return "NUMERIC"; + } + if (field.dataType() instanceof ArrayType) { + ArrayType arrayType = (ArrayType) field.dataType(); + if (arrayType.elementType() instanceof ArrayType) { + throw new SpannerConnectorException( + SpannerErrorCode.UNSUPPORTED_DATATYPE, + "Nested arrays are not supported by Spanner: " + field.dataType()); + } + StructField elementField = + DataTypes.createStructField( + field.name(), arrayType.elementType(), arrayType.containsNull()); + return "ARRAY<" + sparkTypeToSpannerType(elementField) + ">"; + } + + throw new SpannerConnectorException( + SpannerErrorCode.UNSUPPORTED_DATATYPE, + "Unsupported data type in CREATE TABLE: " + field.dataType()); + } + + @Override + public Identifier[] listTables(ReadContext readContext, String[] namespace) { + Statement statement = + Statement.of("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ''"); + try (ResultSet resultSet = readContext.executeQuery(statement)) { + List tables = new ArrayList<>(); + while (resultSet.next()) { + tables.add(Identifier.of(namespace, resultSet.getString("TABLE_NAME"))); + } + return tables.toArray(new Identifier[0]); + } + } + + @Override + public Statement tableExistsStatement(String tableName) { + return Statement.newBuilder( + "SELECT COUNT(1) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '' AND TABLE_NAME = @tableName") + .bind("tableName") + .to(tableName) + .build(); + } +} + +class PostgresSpannerInformationSchema implements SpannerInformationSchema { + @Override + public String quoteIdentifier(String identifier) { + return "\"" + identifier.replace("\"", "\"\"") + "\""; + } + + @Override + public String sparkTypeToSpannerType(StructField field) { + if (field.dataType().equals(DataTypes.LongType)) { + return "bigint"; + } + if (field.dataType().equals(DataTypes.StringType)) { + return "varchar"; + } + if (field.dataType().equals(DataTypes.BooleanType)) { + return "boolean"; + } + if (field.dataType().equals(DataTypes.DoubleType)) { + return "float8"; + } + if (field.dataType().equals(DataTypes.BinaryType)) { + return "bytea"; + } + if (field.dataType().equals(DataTypes.TimestampType)) { + return "timestamptz"; + } + if (field.dataType().equals(DataTypes.DateType)) { + return "date"; + } + if (field.dataType() instanceof org.apache.spark.sql.types.DecimalType) { + return "numeric"; + } + if (field.dataType() instanceof ArrayType) { + ArrayType arrayType = (ArrayType) field.dataType(); + if (arrayType.elementType() instanceof ArrayType) { + throw new SpannerConnectorException( + SpannerErrorCode.UNSUPPORTED_DATATYPE, + "Nested arrays are not supported by Spanner: " + field.dataType()); + } + StructField elementField = + DataTypes.createStructField( + field.name(), arrayType.elementType(), arrayType.containsNull()); + return sparkTypeToSpannerType(elementField) + "[]"; + } + + throw new SpannerConnectorException( + SpannerErrorCode.UNSUPPORTED_DATATYPE, + "Unsupported data type in CREATE TABLE: " + field.dataType()); + } + + @Override + public Identifier[] listTables(ReadContext readContext, String[] namespace) { + Statement statement = + Statement.of( + "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'"); + try (ResultSet resultSet = readContext.executeQuery(statement)) { + List tables = new ArrayList<>(); + while (resultSet.next()) { + tables.add(Identifier.of(namespace, resultSet.getString("table_name"))); + } + return tables.toArray(new Identifier[0]); + } + } + + @Override + public Statement tableExistsStatement(String tableName) { + return Statement.newBuilder( + "SELECT COUNT(1) FROM information_schema.tables WHERE table_schema = 'public' AND table_name = $1") + .bind("p1") + .to(tableName) + .build(); + } +} diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInputPartitionReaderContext.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInputPartitionReaderContext.java index 3bd5e274..2ec24807 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInputPartitionReaderContext.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerInputPartitionReaderContext.java @@ -22,7 +22,6 @@ import com.google.cloud.spanner.ResultSet; import com.google.cloud.spanner.SpannerException; import java.io.IOException; -import java.util.Map; import java.util.Objects; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.util.CaseInsensitiveStringMap; @@ -39,16 +38,13 @@ public SpannerInputPartitionReaderContext( BatchTransactionId batchTransactionId, String mapAsJSONStr, SpannerRowConverter rowConverter) { - Map opts; + CaseInsensitiveStringMap opts; try { - opts = SpannerUtils.deserializeMap(mapAsJSONStr); + opts = new CaseInsensitiveStringMap(SpannerUtils.deserializeMap(mapAsJSONStr)); } catch (JsonProcessingException e) { throw new SpannerConnectorException( SpannerErrorCode.SPANNER_FAILED_TO_PARSE_OPTIONS, "Error parsing the input options.", e); } - // The map might be case-insensitive when being serialized - opts = new CaseInsensitiveStringMap(opts); - // Please note that we are using BatchClientWithCloser to avoid resource leaks. // That is because, since we do have a deterministic scope and timeline for how long // SpannerInputPartitionReaderContext's BatchClient.Spanner will execute, we use this diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanBuilder.java index 58d029fe..8da9a83a 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanBuilder.java @@ -27,7 +27,6 @@ import org.apache.spark.sql.sources.Filter; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -36,7 +35,6 @@ */ public class SpannerScanBuilder implements ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns { - private CaseInsensitiveStringMap opts; private List pushedFilters; private Set requiredColumns; private SpannerScanner scanner; @@ -44,10 +42,9 @@ public class SpannerScanBuilder private SpannerTable spannerTable; private Map fields; - public SpannerScanBuilder(CaseInsensitiveStringMap options) { - this.opts = options; + public SpannerScanBuilder(SpannerTable spannerTable) { this.pushedFilters = new ArrayList(); - this.spannerTable = new SpannerTable(options); + this.spannerTable = spannerTable; this.fields = new LinkedHashMap<>(); for (StructField field : spannerTable.schema().fields()) { fields.put(field.name(), field); @@ -58,7 +55,7 @@ public SpannerScanBuilder(CaseInsensitiveStringMap options) { public Scan build() { this.scanner = new SpannerScanner( - this.opts.asCaseSensitiveMap(), + this.spannerTable.properties(), this.spannerTable, this.fields, this.pushedFilters(), diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanner.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanner.java index 90b8ea79..22c70408 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanner.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerScanner.java @@ -36,6 +36,7 @@ import org.apache.spark.sql.sources.Filter; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -46,14 +47,14 @@ public class SpannerScanner implements Batch, Scan { private final SpannerTable spannerTable; private final Filter[] filters; private final Set requiredColumns; - private final Map opts; + private final CaseInsensitiveStringMap opts; private static final Logger log = LoggerFactory.getLogger(SpannerScanner.class); private final Timestamp INIT_TIME = Timestamp.now(); private final Map fields; private final StructType readSchema; public SpannerScanner( - Map opts, + CaseInsensitiveStringMap opts, SpannerTable spannerTable, Map fields, Filter[] filters, @@ -93,18 +94,21 @@ static String buildColumnsWithTablePrefix( @Override public InputPartition[] planInputPartitions() { BatchClientWithCloser batchClient = SpannerUtils.batchClientFromProperties(this.opts); + boolean isPostgreSql = batchClient.databaseClient.getDialect().equals(Dialect.POSTGRESQL); // 1. Use * if no requiredColumns were requested else select them. String selectPrefix = "SELECT *"; if (this.requiredColumns != null && this.requiredColumns.size() > 0) { // Prefix each column with the table name to avoid ambiguity when column name // matches table name - boolean isPostgreSql = batchClient.databaseClient.getDialect().equals(Dialect.POSTGRESQL); String columnsWithTablePrefix = buildColumnsWithTablePrefix(this.spannerTable.name(), this.requiredColumns, isPostgreSql); selectPrefix = "SELECT " + columnsWithTablePrefix; } - String sqlStmt = selectPrefix + " FROM " + this.spannerTable.name(); + + String quotedTableName = + isPostgreSql ? "\"" + spannerTable.name() + "\"" : "`" + spannerTable.name() + "`"; + String sqlStmt = selectPrefix + " FROM " + quotedTableName; if (this.filters.length > 0) { sqlStmt += " WHERE " @@ -116,7 +120,7 @@ public InputPartition[] planInputPartitions() { this.filters); } - Boolean enableDataboost = false; + boolean enableDataboost = false; if (this.opts.containsKey("enableDataBoost")) { enableDataboost = this.opts.get("enableDataBoost").equalsIgnoreCase("true"); } diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java index 17a01a53..94a174cb 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTable.java @@ -16,12 +16,15 @@ import com.google.cloud.spanner.Dialect; import com.google.cloud.spanner.connection.Connection; +import com.google.common.base.Verify; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import java.util.HashMap; import java.util.Map; import java.util.Set; import java.util.regex.Matcher; import java.util.regex.Pattern; +import javax.annotation.Nullable; import org.apache.spark.sql.connector.catalog.SupportsRead; import org.apache.spark.sql.connector.catalog.SupportsWrite; import org.apache.spark.sql.connector.catalog.Table; @@ -49,63 +52,76 @@ public class SpannerTable implements Table, SupportsRead, SupportsWrite { private final String projectId; private final SpannerTableSchema dbSchema; - private final StructType sparkSchema; + private final @Nullable StructType dfSchema; private static final ImmutableSet tableCapabilities = ImmutableSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE); - private final Map properties; + private final CaseInsensitiveStringMap properties; private static final Logger log = LoggerFactory.getLogger(SpannerTable.class); public SpannerTable(Map properties) { - this.properties = properties; - this.tableName = SpannerUtils.getRequiredOption(properties, "table"); - this.projectId = SpannerUtils.getRequiredOption(properties, "projectId"); - this.instanceId = SpannerUtils.getRequiredOption(properties, "instanceId"); - this.databaseId = SpannerUtils.getRequiredOption(properties, "databaseId"); - try (Connection conn = SpannerUtils.connectionFromProperties(properties)) { - boolean isPostgreSql; - if (conn.getDialect().equals(Dialect.GOOGLE_STANDARD_SQL)) { - isPostgreSql = false; - } else if (conn.getDialect().equals(Dialect.POSTGRESQL)) { - isPostgreSql = true; - } else { - throw new SpannerConnectorException( - SpannerErrorCode.DATABASE_DIALECT_NOT_SUPPORTED, - "The dialect used " - + conn.getDialect() - + " in the Spanner table " - + tableName - + " is not supported."); - } - this.dbSchema = new SpannerTableSchema(conn, tableName, isPostgreSql); - this.sparkSchema = this.dbSchema.schema; + this(new CaseInsensitiveStringMap(properties), null); + } + + public SpannerTable(CaseInsensitiveStringMap properties, StructType dfSchema) { + this( + SpannerUtils.getRequiredOption(properties, "projectId"), + SpannerUtils.getRequiredOption(properties, "instanceId"), + SpannerUtils.getRequiredOption(properties, "databaseId"), + SpannerUtils.getRequiredOption(properties, "table"), + properties, + dfSchema); + } + + private boolean checkIsPostgreSql(Connection conn) { + boolean isPostgreSql; + if (conn.getDialect().equals(Dialect.GOOGLE_STANDARD_SQL)) { + isPostgreSql = false; + } else if (conn.getDialect().equals(Dialect.POSTGRESQL)) { + isPostgreSql = true; + } else { + throw new SpannerConnectorException( + SpannerErrorCode.DATABASE_DIALECT_NOT_SUPPORTED, + "The dialect used " + + conn.getDialect() + + " in the Spanner database " + + this.databaseId + + " is not supported."); } + return isPostgreSql; } - public SpannerTable(Map properties, StructType dfSchema) { - this.properties = properties; - this.tableName = SpannerUtils.getRequiredOption(properties, "table"); - this.projectId = SpannerUtils.getRequiredOption(properties, "projectId"); - this.instanceId = SpannerUtils.getRequiredOption(properties, "instanceId"); - this.databaseId = SpannerUtils.getRequiredOption(properties, "databaseId"); - try (Connection conn = SpannerUtils.connectionFromProperties(properties)) { - boolean isPostgreSql; - if (conn.getDialect().equals(Dialect.GOOGLE_STANDARD_SQL)) { - isPostgreSql = false; - } else if (conn.getDialect().equals(Dialect.POSTGRESQL)) { - isPostgreSql = true; - } else { - throw new SpannerConnectorException( - SpannerErrorCode.DATABASE_DIALECT_NOT_SUPPORTED, - "The dialect used " - + conn.getDialect() - + " in the Spanner table " - + tableName - + " is not supported."); - } + public SpannerTable( + String projectId, + String instanceId, + String databaseId, + String tableNameOption, + CaseInsensitiveStringMap properties, + StructType dfSchema) { + Verify.verifyNotNull(projectId, "projectId"); + Verify.verifyNotNull(instanceId, "instanceId"); + Verify.verifyNotNull(databaseId, "databaseId"); + Verify.verifyNotNull(tableNameOption, "tableNameOption"); + Verify.verifyNotNull(properties, "properties"); + + try (Connection conn = + SpannerUtils.connectionFromProperties( + projectId, instanceId, databaseId, properties.get("emulatorHost"))) { + boolean isPostgreSql = checkIsPostgreSql(conn); + + this.tableName = tableNameOption; + this.projectId = projectId; + this.instanceId = instanceId; + this.databaseId = databaseId; + + Map combinedProps = new HashMap<>(properties.asCaseSensitiveMap()); + combinedProps.putAll(getOpenLineageDatasetProperties()); + combinedProps.put("table", tableName); + this.properties = new CaseInsensitiveStringMap(combinedProps); + // Still get the DB schema for validation. this.dbSchema = new SpannerTableSchema(conn, tableName, isPostgreSql); - this.sparkSchema = dfSchema; + this.dfSchema = dfSchema; } } @@ -257,7 +273,11 @@ public static DataType ofSpannerStrTypePg(String spannerStrType, boolean isNulla @Override public StructType schema() { - return this.sparkSchema; + if (this.properties.getOrDefault("enablePartialRowUpdates", "false").equalsIgnoreCase("true") + && dfSchema != null) { + return this.dfSchema; + } + return this.dbSchema.schema; } /* @@ -277,25 +297,27 @@ public String name() { @Override public ScanBuilder newScanBuilder(CaseInsensitiveStringMap options) { - return new SpannerScanBuilder(options); + return new SpannerScanBuilder(this); } @Override public WriteBuilder newWriteBuilder(LogicalWriteInfo info) { SpannerUtils.validateSchema(info.schema(), this.dbSchema.schema, this.tableName); - return new SpannerWriteBuilder(info); + return new SpannerWriteBuilder(info, this.properties); } @Override - public Map properties() { - ImmutableMap.Builder builder = - ImmutableMap.builder() - .putAll(this.properties) - .put("openlineage.dataset.name", String.format("%s/%s", databaseId, tableName)) - .put( - "openlineage.dataset.namespace", - String.format("spanner://%s/%s", projectId, instanceId)) - .put("openlineage.dataset.storageDatasetFacet.storageLayer", "spanner"); - return builder.build(); + public CaseInsensitiveStringMap properties() { + return properties; + } + + private Map getOpenLineageDatasetProperties() { + return ImmutableMap.of( + "openlineage.dataset.name", + String.format("%s/%s", databaseId, tableName), + "openlineage.dataset.namespace", + String.format("%s/%s", projectId, instanceId), + "openlineage.dataset.storageDatasetFacet.storageLayer", + "spanner"); } } diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTableSchema.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTableSchema.java index 260e2c84..7dd90ed2 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTableSchema.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerTableSchema.java @@ -1,3 +1,16 @@ +// Copyright 2026 Google LLC +// +// 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. package com.google.cloud.spark.spanner; import com.google.cloud.spanner.ResultSet; @@ -5,8 +18,11 @@ import com.google.cloud.spanner.Struct; import com.google.cloud.spanner.connection.Connection; import java.util.HashMap; +import java.util.HashSet; +import java.util.Locale; import java.util.Map; import java.util.Objects; +import java.util.Set; import org.apache.spark.sql.types.DataType; import org.apache.spark.sql.types.MetadataBuilder; import org.apache.spark.sql.types.StructField; @@ -39,6 +55,10 @@ static Statement buildSchemaQuery(String tableName, boolean isPostgreSql) { public SpannerTableSchema(Connection conn, String tableName, boolean isPostgreSql) { this.name = tableName; this.columns = new HashMap<>(); + // 1. Get the primary keys for the table. + Set primaryKeys = getPrimaryKeys(conn, tableName, isPostgreSql); + + // 2. Get the schema of the table. Statement stmt = buildSchemaQuery(tableName, isPostgreSql); try (final ResultSet rs = conn.executeQuery(stmt)) { // Expecting resultset columns in the ordering: @@ -50,8 +70,10 @@ public SpannerTableSchema(Connection conn, String tableName, boolean isPostgreSq while (rs.next()) { Struct row = rs.getCurrentRowAsStruct(); String columnName = row.getString(0); + boolean isPrimaryKey = primaryKeys.contains(columnName); StructField structField = - getSparkStructField(columnName, row.getString(2), row.getBoolean(1), isPostgreSql); + getSparkStructField( + columnName, row.getString(2), row.getBoolean(1), isPostgreSql, isPrimaryKey); schema = schema.add(structField); this.columns.put(columnName, structField); } @@ -59,8 +81,44 @@ public SpannerTableSchema(Connection conn, String tableName, boolean isPostgreSq } } + static Statement buildPrimaryKeyQuery(String tableName, boolean isPostgreSql) { + if (isPostgreSql) { + String sql = + "SELECT kcu.COLUMN_NAME FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS AS tc " + + "JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE AS kcu " + + "ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME AND tc.TABLE_NAME = kcu.TABLE_NAME " + + "WHERE tc.TABLE_NAME = $1 AND tc.CONSTRAINT_TYPE = 'PRIMARY KEY' AND tc.TABLE_SCHEMA = 'public' " + + "ORDER BY kcu.ORDINAL_POSITION"; + return Statement.newBuilder(sql).bind("p1").to(tableName.toLowerCase(Locale.ROOT)).build(); + } + String sql = + "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.INDEX_COLUMNS " + + "WHERE INDEX_NAME = 'PRIMARY_KEY' AND UPPER(TABLE_NAME) = UPPER(@tableName)"; + return Statement.newBuilder(sql).bind("tableName").to(tableName).build(); + } + + private Set getPrimaryKeys(Connection conn, String tableName, boolean isPostgreSql) { + Statement stmt = buildPrimaryKeyQuery(tableName, isPostgreSql); + Set primaryKeys = new HashSet<>(); + try (final ResultSet rs = conn.executeQuery(stmt)) { + while (rs.next()) { + primaryKeys.add(rs.getString(0)); + } + } + return primaryKeys; + } + public static StructField getSparkStructField( String name, String spannerType, boolean isNullable, boolean isPostgreSql) { + return getSparkStructField(name, spannerType, isNullable, isPostgreSql, false); + } + + public static StructField getSparkStructField( + String name, + String spannerType, + boolean isNullable, + boolean isPostgreSql, + boolean isPrimaryKey) { DataType catalogType = isPostgreSql ? SpannerTable.ofSpannerStrTypePg(spannerType, isNullable) @@ -71,6 +129,9 @@ public static StructField getSparkStructField( } else if (isJsonb(spannerType)) { metadataBuilder.putString(SpannerUtils.COLUMN_TYPE, "jsonb"); } + if (isPrimaryKey) { + metadataBuilder.putBoolean(SpannerUtils.PRIMARY_KEY_TAG, true); + } return new StructField(name, catalogType, isNullable, metadataBuilder.build()); } diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java index 9b1e5d0b..84eb6775 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerUtils.java @@ -33,6 +33,7 @@ import com.google.cloud.spanner.connection.Connection; import com.google.cloud.spanner.connection.ConnectionOptions; import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Verify; import com.google.common.io.CharStreams; import com.google.common.io.Closeables; import java.io.IOException; @@ -50,6 +51,7 @@ import java.util.Properties; import java.util.Set; import java.util.function.Function; +import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.Stream; import javax.annotation.Nullable; @@ -64,11 +66,25 @@ 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; public class SpannerUtils { + // GCP project ID naming rules (also allows legacy domain-scoped projects like + // "example.com:my-project" and numeric project IDs): + // https://cloud.google.com/resource-manager/docs/creating-managing-projects + private static final Pattern VALID_PROJECT_ID_PATTERN = + Pattern.compile("^[a-z][a-z0-9.:-]*[a-z0-9]$|^[0-9]+$"); + + // Spanner resource ID naming rules (superset covering both instance and database IDs): + // Instance IDs: [a-z][-a-z0-9]*[a-z0-9], 2-64 chars. + // https://cloud.google.com/spanner/docs/reference/rest/v1/projects.instances/create + // Database IDs: [a-z][a-z0-9_\-]*[a-z0-9], 2-30 chars. + // https://cloud.google.com/spanner/docs/reference/rest/v1/projects.instances.databases/create + private static final Pattern VALID_ID_PATTERN = Pattern.compile("^[a-z][a-z0-9_-]*[a-z0-9]$"); + private static final RetrySettings RETRY_SETTING = RetrySettings.newBuilder() .setInitialRpcTimeout(Duration.ofHours(2)) @@ -95,6 +111,8 @@ public class SpannerUtils { public static final String COLUMN_TYPE = "col_type"; + public static final String PRIMARY_KEY_TAG = "spanner.primaryKey"; + public static Long SECOND_TO_DAYS = 60 * 60 * 24L; private static final String SPARK_VERSION = org.apache.spark.package$.MODULE$.SPARK_VERSION(); @@ -162,9 +180,28 @@ static Optional getGcpRegion() { } public static Connection connectionFromProperties(Map properties) { + CaseInsensitiveStringMap caseInsensitiveOptions = new CaseInsensitiveStringMap(properties); + + return connectionFromProperties( + caseInsensitiveOptions.get("projectId"), + caseInsensitiveOptions.get("instanceId"), + caseInsensitiveOptions.get("databaseId"), + caseInsensitiveOptions.get("emulatorHost")); + } + + public static Connection connectionFromProperties( + String projectId, String instanceId, String databaseId, @Nullable String emulatorHost) { + + Verify.verifyNotNull(projectId, "projectId"); + Verify.verifyNotNull(instanceId, "instanceId"); + Verify.verifyNotNull(databaseId, "databaseId"); + validateProjectId(projectId); + validateResourceId(instanceId, "instanceId"); + validateResourceId(databaseId, "databaseId"); + String connUriPrefix = "cloudspanner:"; - String emulatorHost = properties.get("emulatorHost"); if (emulatorHost != null) { + validateEmulatorHost(emulatorHost); connUriPrefix = "cloudspanner://" + emulatorHost; } @@ -172,9 +209,9 @@ public static Connection connectionFromProperties(Map properties String.format( connUriPrefix + "/projects/%s/instances/%s/databases/%s?autoConfigEmulator=%s;usePlainText=%s", - properties.get("projectId"), - properties.get("instanceId"), - properties.get("databaseId"), + projectId, + instanceId, + databaseId, emulatorHost != null, emulatorHost != null); @@ -184,12 +221,61 @@ public static Connection connectionFromProperties(Map properties return opts.getConnection(); } - public static BatchClientWithCloser batchClientFromProperties(Map properties) { + @VisibleForTesting + static void validateProjectId(String value) { + if (!VALID_PROJECT_ID_PATTERN.matcher(value).matches()) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Invalid projectId: '" + + value + + "'. Must contain only lowercase letters, digits, hyphens, dots, and colons."); + } + } + + @VisibleForTesting + static void validateResourceId(String value, String name) { + if (!VALID_ID_PATTERN.matcher(value).matches()) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Invalid " + + name + + ": '" + + value + + "'. Must contain only lowercase letters, digits, hyphens, and underscores."); + } + } + + @VisibleForTesting + static void validateEmulatorHost(String host) { + if (host.contains("?") || host.contains(";") || host.contains(" ") || host.contains("/")) { + throw new SpannerConnectorException( + SpannerErrorCode.INVALID_ARGUMENT, + "Invalid emulatorHost: '" + host + "'. Must not contain URI-sensitive characters."); + } + } + + public static BatchClientWithCloser batchClientFromProperties( + CaseInsensitiveStringMap properties) { return batchClientFromProperties(properties, defaultSessionPoolOptions); } public static BatchClientWithCloser batchClientFromProperties( - Map properties, SessionPoolOptions sessionPoolOptions) { + CaseInsensitiveStringMap 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( + CaseInsensitiveStringMap properties, SessionPoolOptions sessionPoolOptions) { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setSessionPoolOption(sessionPoolOptions) @@ -222,13 +308,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) { @@ -503,13 +583,13 @@ public static StructType pruneSchema( return prunedSchema; } - static String getRequiredOption(Map properties, String option) { - String tableName = properties.get(option); - if (tableName == null) { + static String getRequiredOption(CaseInsensitiveStringMap properties, String option) { + String value = properties.get(option); + if (value == null) { throw new SpannerConnectorException( SpannerErrorCode.INVALID_ARGUMENT, "Option '" + option + "' property must be set"); } - return tableName; + return value; } public static void validateSchema( diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java index 92fcc79f..1aa8588f 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SpannerWriteBuilder.java @@ -17,16 +17,19 @@ import org.apache.spark.sql.connector.write.BatchWrite; import org.apache.spark.sql.connector.write.LogicalWriteInfo; import org.apache.spark.sql.connector.write.WriteBuilder; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; public class SpannerWriteBuilder implements WriteBuilder { private final LogicalWriteInfo info; + private final CaseInsensitiveStringMap properties; - public SpannerWriteBuilder(LogicalWriteInfo info) { + public SpannerWriteBuilder(LogicalWriteInfo info, CaseInsensitiveStringMap properties) { this.info = info; + this.properties = properties; } @Override public BatchWrite buildForBatch() { - return new SpannerBatchWrite(info); + return new SpannerBatchWrite(info, properties); } } diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBase.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBase.java index cccf8bfb..aed2f0b4 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBase.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBase.java @@ -15,7 +15,12 @@ package com.google.cloud.spark.spanner; import com.google.cloud.spark.spanner.graph.SpannerGraphBuilder; +import com.google.common.collect.ImmutableList; +import com.google.gson.Gson; +import java.util.HashMap; +import java.util.List; import java.util.Map; +import org.apache.spark.sql.connector.catalog.Identifier; import org.apache.spark.sql.connector.catalog.Table; import org.apache.spark.sql.connector.catalog.TableProvider; import org.apache.spark.sql.connector.expressions.Transform; @@ -25,6 +30,12 @@ public abstract class SparkSpannerTableProviderBase implements DataSourceRegister, TableProvider { + private static final Gson GSON = new Gson(); + + static final List GRAPH_OPTION_KEYS = + ImmutableList.of( + "graph", "type", "enableDataBoost", "configs", "graphQuery", "timestamp", "viewsEnabled"); + /* * Infers the schema of the table identified by the given options. */ @@ -43,7 +54,6 @@ public Table getTable( final CaseInsensitiveStringMap options = new CaseInsensitiveStringMap(properties); boolean enablePartialRowUpdates = Boolean.parseBoolean(options.getOrDefault("enablePartialRowUpdates", "false")); - boolean hasTable = options.containsKey("table"); boolean hasGraph = options.containsKey("graph"); if (hasTable && !hasGraph) { @@ -92,4 +102,28 @@ private Table getTable(Map properties) { "properties must contain one of \"table\" or \"graph\""); } } + + public Identifier extractIdentifier(CaseInsensitiveStringMap options) { + String table = options.get("table"); + if (table != null) { + return Identifier.of(new String[0], table); + } + String graph = options.get("graph"); + if (graph != null) { + Map graphProps = new HashMap<>(); + for (String key : GRAPH_OPTION_KEYS) { + String val = options.get(key); + if (val != null) { + graphProps.put(key, val); + } + } + return Identifier.of( + new String[0], SpannerCatalog.GRAPH_IDENTIFIER_PREFIX + GSON.toJson(graphProps)); + } + return null; + } + + public String extractCatalog(CaseInsensitiveStringMap options) { + return options.getOrDefault("catalog", "spark_catalog"); + } } diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraph.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraph.java index 06512f86..0fc9637b 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraph.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraph.java @@ -1,3 +1,16 @@ +// Copyright 2026 Google LLC +// +// 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. package com.google.cloud.spark.spanner.graph; import com.google.cloud.spanner.Options; @@ -29,7 +42,7 @@ public class SpannerGraph implements Table, SupportsRead, SupportsWrite { static final List requiredOptions = ImmutableList.of("projectId", "instanceId", "databaseId", "graph", "type"); - public final Map options; + public final CaseInsensitiveStringMap options; public final Options.ReadAndQueryOption dataBoostEnabled; public final SpannerGraphConfigs configs; public final @Nullable Statement directQuery; diff --git a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraphScanner.java b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraphScanner.java index bd309652..f4f0412c 100644 --- a/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraphScanner.java +++ b/spark-3.1-spanner-lib/src/main/java/com/google/cloud/spark/spanner/graph/SpannerGraphScanner.java @@ -1,3 +1,16 @@ +// Copyright 2026 Google LLC +// +// 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. package com.google.cloud.spark.spanner.graph; import com.fasterxml.jackson.core.JsonProcessingException; @@ -43,7 +56,7 @@ public class SpannerGraphScanner implements Batch, Scan { private static final Logger log = LoggerFactory.getLogger(SpannerScanner.class); - private final Map options; + private final CaseInsensitiveStringMap options; private final @Nullable Map> extraHeaders; private final TimestampBound readTimestamp; private final @Nullable Long partitionSizeBytes; diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java new file mode 100644 index 00000000..771840e6 --- /dev/null +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerCatalogTest.java @@ -0,0 +1,346 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.cloud.spanner.DatabaseClient; +import com.google.cloud.spanner.DatabaseId; +import com.google.cloud.spanner.Dialect; +import com.google.cloud.spanner.ReadContext; +import com.google.cloud.spanner.ReadOnlyTransaction; +import com.google.cloud.spanner.ResultSet; +import com.google.cloud.spanner.Spanner; +import com.google.cloud.spanner.Statement; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.Map; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +@RunWith(Parameterized.class) +public class SpannerCatalogTest { + + @Parameters + public static Collection dialects() { + return Arrays.asList(Dialect.GOOGLE_STANDARD_SQL, Dialect.POSTGRESQL); + } + + @Rule public ExpectedException thrown = ExpectedException.none(); + + private SpannerCatalog catalog; + private final Dialect dialect; + + @Mock private Spanner spanner; + @Mock private DatabaseClient dbClient; + @Mock private SpannerInformationSchema spannerInfoSchema; + @Mock private ReadContext singleUseReadContext; + + public SpannerCatalogTest(Dialect dialect) { + this.dialect = dialect; + } + + @Before + public void setUp() { + MockitoAnnotations.openMocks(this); + + catalog = + new SpannerCatalog() { + @Override + protected Spanner createSpanner(CaseInsensitiveStringMap options) { + return spanner; + } + + @Override + protected SpannerInformationSchema createSchemaInfo(Dialect dialect) { + return spannerInfoSchema; + } + + @Override + protected SpannerTable factorySpannerTable(Identifier ident) { + SpannerTable mockSpannerTable = mock(SpannerTable.class); + when(mockSpannerTable.name()).thenReturn(ident.name()); + return mockSpannerTable; + } + }; + + Map opts = new HashMap<>(); + opts.put("projectId", "p"); + opts.put("instanceId", "i"); + opts.put("databaseId", "d"); + opts.put("emulatorHost", "localhost:9010"); + CaseInsensitiveStringMap options = new CaseInsensitiveStringMap(opts); + catalog.initialize("test-catalog", options); + + when(spanner.getDatabaseClient(any(DatabaseId.class))).thenReturn(dbClient); + when(dbClient.getDialect()).thenReturn(dialect); + ReadOnlyTransaction mockRoTransaction = mock(ReadOnlyTransaction.class); + when(dbClient.readOnlyTransaction()).thenReturn(mockRoTransaction); + when(dbClient.singleUse()).thenReturn(singleUseReadContext); + } + + @Test + public void testName() { + assertEquals("test-catalog", catalog.name()); + } + + @Test + public void listTablesShouldReturnTables() { + String[] namespace = new String[0]; + Identifier[] expectedTables = {Identifier.of(namespace, "t1"), Identifier.of(namespace, "t2")}; + when(spannerInfoSchema.listTables(any(ReadContext.class), any(String[].class))) + .thenReturn(expectedTables); + + Identifier[] tables = catalog.listTables(namespace); + verify(spannerInfoSchema).listTables(any(ReadContext.class), eq(namespace)); + assertArrayEquals(expectedTables, tables); + } + + @Test + public void listTablesShouldReturnEmptyForInvalidNamespace() { + String[] namespace = new String[] {"p", "i"}; + Identifier[] tables = catalog.listTables(namespace); + assertEquals(0, tables.length); + } + + @Test + public void loadTableShouldThrowNoSuchTableException() throws NoSuchTableException { + Identifier ident = Identifier.of(new String[0], "non_existent"); + thrown.expect(NoSuchTableException.class); + catalog.loadTable(ident); + } + + @Test + public void loadTableShouldReturnSpannerTable() throws NoSuchTableException { + Identifier ident = Identifier.of(new String[0], "t1"); + Statement mockStatement = Statement.of(""); + when(spannerInfoSchema.tableExistsStatement("t1")).thenReturn(mockStatement); + ResultSet mockResultSet = mock(ResultSet.class); + when(singleUseReadContext.executeQuery(mockStatement)).thenReturn(mockResultSet); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getLong(0)).thenReturn(1L); + Table table = catalog.loadTable(ident); + assertNotNull(table); + assertTrue(table instanceof SpannerTable); + assertEquals("t1", table.name()); + } + + @Test + public void loadTableShouldThrowExceptionForInvalidNamespace() throws NoSuchTableException { + Identifier ident = Identifier.of(new String[] {"p", "i"}, "t1"); + thrown.expect(SpannerConnectorException.class); + catalog.loadTable(ident); + } + + @Test + public void tableExistsShouldReturnTrue() { + Identifier ident = Identifier.of(new String[0], "t1"); + Statement mockStatement = Statement.of(""); + when(spannerInfoSchema.tableExistsStatement(any(String.class))).thenReturn(mockStatement); + ResultSet mockResultSet = mock(ResultSet.class); + when(singleUseReadContext.executeQuery(mockStatement)).thenReturn(mockResultSet); + when(mockResultSet.next()).thenReturn(true); + when(mockResultSet.getLong(0)).thenReturn(1L); + assertTrue(catalog.tableExists(ident)); + } + + @Test + public void tableExistsShouldReturnFalse() { + Identifier ident = Identifier.of(new String[] {"p", "i", "d"}, "non_existent"); + assertFalse(catalog.tableExists(ident)); + } + + @Test + public void tableExistsShouldReturnFalseForInvalidNamespace() { + Identifier ident = Identifier.of(new String[] {"p", "i"}, "t1"); + assertFalse(catalog.tableExists(ident)); + } + + @Test + public void createTableShouldThrowExceptionOnNoPrimaryKey() { + Identifier ident = Identifier.of(new String[] {"p", "i", "d"}, "no_pk_table"); + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.LongType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()) + }); + + thrown.expect(SpannerConnectorException.class); + thrown.expectMessage( + "No primary key found for table no_pk_table. Please specify at least one primary key column."); + + SpannerInformationSchema.create(dialect).createTableDdl(ident, schema); + } + + @Test + public void isGraphIdentifierShouldReturnTrue() { + Identifier ident = + Identifier.of(new String[0], SpannerCatalog.GRAPH_IDENTIFIER_PREFIX + "{\"graph\":\"G\"}"); + assertTrue(SpannerCatalog.isGraphIdentifier(ident)); + } + + @Test + public void isGraphIdentifierShouldReturnFalse() { + Identifier ident = Identifier.of(new String[0], "my_table"); + assertFalse(SpannerCatalog.isGraphIdentifier(ident)); + } + + @Test + public void loadTableShouldRejectGraphIdentifierWithNamespace() throws NoSuchTableException { + Identifier ident = + Identifier.of( + new String[] {"ns"}, SpannerCatalog.GRAPH_IDENTIFIER_PREFIX + "{\"graph\":\"G\"}"); + thrown.expect(SpannerConnectorException.class); + catalog.loadTable(ident); + } + + @Test + public void loadTableShouldRejectEmptyGraphIdentifier() throws NoSuchTableException { + Identifier ident = Identifier.of(new String[0], SpannerCatalog.GRAPH_IDENTIFIER_PREFIX); + thrown.expect(SpannerConnectorException.class); + thrown.expectMessage("Graph identifier has no encoded properties"); + catalog.loadTable(ident); + } + + @Test + public void loadTableShouldRejectMalformedGraphJson() throws NoSuchTableException { + Identifier ident = + Identifier.of(new String[0], SpannerCatalog.GRAPH_IDENTIFIER_PREFIX + "not-valid-json"); + thrown.expect(SpannerConnectorException.class); + thrown.expectMessage("Malformed graph identifier JSON"); + catalog.loadTable(ident); + } + + @Test + public void alterTableShouldThrowException() { + thrown.expect(UnsupportedOperationException.class); + catalog.alterTable(null, (org.apache.spark.sql.connector.catalog.TableChange[]) null); + } + + @Test + public void renameTableShouldThrowException() { + thrown.expect(UnsupportedOperationException.class); + catalog.renameTable(null, null); + } + + @Test + public void testToDdl() { + Identifier ident = Identifier.of(new String[] {"p", "i", "d"}, "my_table"); + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA), + new StructField( + "id2", DataTypes.StringType, false, SpannerCatalog.PRIMARY_KEY_METADATA), + new StructField("name", DataTypes.StringType, true, Metadata.empty()), + new StructField("active", DataTypes.BooleanType, false, Metadata.empty()), + new StructField("amount", DataTypes.DoubleType, true, Metadata.empty()), + new StructField("data", DataTypes.BinaryType, true, Metadata.empty()), + new StructField("created_at", DataTypes.TimestampType, true, Metadata.empty()), + new StructField("created_on", DataTypes.DateType, true, Metadata.empty()), + new StructField("price", DataTypes.createDecimalType(10, 2), true, Metadata.empty()), + }); + + String ddl = SpannerInformationSchema.create(dialect).createTableDdl(ident, schema); + + if (dialect == Dialect.POSTGRESQL) { + assertEquals( + "CREATE TABLE \"my_table\" (\"id\" bigint NOT NULL, \"id2\" varchar NOT NULL, \"name\" varchar, " + + "\"active\" boolean NOT NULL, \"amount\" float8, \"data\" bytea, " + + "\"created_at\" timestamptz, \"created_on\" date, \"price\" numeric, " + + "PRIMARY KEY (\"id\", \"id2\"))", + ddl); + } else { + assertEquals( + "CREATE TABLE `my_table` (`id` INT64 NOT NULL, `id2` STRING(MAX) NOT NULL, `name` STRING(MAX), " + + "`active` BOOL NOT NULL, `amount` FLOAT64, `data` BYTES(MAX), `created_at` TIMESTAMP, " + + "`created_on` DATE, `price` NUMERIC, PRIMARY KEY (`id`, `id2`))", + ddl); + } + } + + @Test + public void enrichSchemaWithPrimaryKeysShouldAddMetadata() { + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.LongType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()), + }); + Map props = new HashMap<>(); + props.put("primaryKeys", "id"); + + StructType enriched = SpannerCatalog.enrichSchemaWithPrimaryKeys(schema, props); + assertTrue(enriched.fields()[0].metadata().getBoolean(SpannerUtils.PRIMARY_KEY_TAG)); + assertFalse(enriched.fields()[1].metadata().contains(SpannerUtils.PRIMARY_KEY_TAG)); + } + + @Test + public void enrichSchemaWithPrimaryKeysShouldHandleMultipleKeys() { + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.LongType, false, Metadata.empty()), + new StructField("id2", DataTypes.StringType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()), + }); + Map props = new HashMap<>(); + props.put("primaryKeys", "id, id2"); + + StructType enriched = SpannerCatalog.enrichSchemaWithPrimaryKeys(schema, props); + assertTrue(enriched.fields()[0].metadata().getBoolean(SpannerUtils.PRIMARY_KEY_TAG)); + assertTrue(enriched.fields()[1].metadata().getBoolean(SpannerUtils.PRIMARY_KEY_TAG)); + assertFalse(enriched.fields()[2].metadata().contains(SpannerUtils.PRIMARY_KEY_TAG)); + } + + @Test + public void enrichSchemaWithPrimaryKeysShouldBeNoOpWhenAbsent() { + StructType schema = + new StructType( + new StructField[] { + new StructField("id", DataTypes.LongType, false, Metadata.empty()), + }); + Map props = new HashMap<>(); + + StructType enriched = SpannerCatalog.enrichSchemaWithPrimaryKeys(schema, props); + assertFalse(enriched.fields()[0].metadata().contains(SpannerUtils.PRIMARY_KEY_TAG)); + } +} diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerTableSchemaTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerTableSchemaTest.java index c3196efc..ad400892 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerTableSchemaTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerTableSchemaTest.java @@ -15,8 +15,13 @@ package com.google.cloud.spark.spanner; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; import com.google.cloud.spanner.Statement; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -56,4 +61,69 @@ public void testBuildSchemaQuery_postgreSql_usesDirectComparison() { assertThat(query).contains("columns.table_name=$1"); assertThat(query).doesNotContain("UPPER"); } + + @Test + public void testBuildPrimaryKeyQuery_googleSql() { + Statement stmt = SpannerTableSchema.buildPrimaryKeyQuery("MyTable", false); + String query = stmt.getSql(); + assertThat(query).contains("INDEX_NAME = 'PRIMARY_KEY'"); + assertThat(query).contains("UPPER(TABLE_NAME) = UPPER(@tableName)"); + } + + @Test + public void testBuildPrimaryKeyQuery_postgreSql() { + Statement stmt = SpannerTableSchema.buildPrimaryKeyQuery("MyTable", true); + String query = stmt.getSql(); + assertThat(query).contains("CONSTRAINT_TYPE = 'PRIMARY KEY'"); + assertThat(query).contains("TABLE_NAME = $1"); + } + + @Test + public void testIsJson() { + assertTrue(SpannerTableSchema.isJson("JSON")); + assertTrue(SpannerTableSchema.isJson("json")); + assertTrue(SpannerTableSchema.isJson(" json ")); + assertFalse(SpannerTableSchema.isJson("jsonb")); + assertFalse(SpannerTableSchema.isJson("STRING")); + } + + @Test + public void testIsJsonb() { + assertTrue(SpannerTableSchema.isJsonb("JSONB")); + assertTrue(SpannerTableSchema.isJsonb("jsonb")); + assertTrue(SpannerTableSchema.isJsonb(" jsonb ")); + assertFalse(SpannerTableSchema.isJsonb("json")); + assertFalse(SpannerTableSchema.isJsonb("STRING")); + } + + @Test + public void testGetSparkStructField_basic() { + StructField field = SpannerTableSchema.getSparkStructField("col", "INT64", true, false); + assertEquals("col", field.name()); + assertEquals(DataTypes.LongType, field.dataType()); + assertTrue(field.nullable()); + } + + @Test + public void testGetSparkStructField_primaryKey() { + StructField field = SpannerTableSchema.getSparkStructField("id", "INT64", false, false, true); + assertEquals("id", field.name()); + assertFalse(field.nullable()); + assertTrue(field.metadata().contains(SpannerUtils.PRIMARY_KEY_TAG)); + assertTrue(field.metadata().getBoolean(SpannerUtils.PRIMARY_KEY_TAG)); + } + + @Test + public void testGetSparkStructField_jsonMetadata() { + StructField field = SpannerTableSchema.getSparkStructField("data", "JSON", true, false); + assertTrue(field.metadata().contains(SpannerUtils.COLUMN_TYPE)); + assertEquals("json", field.metadata().getString(SpannerUtils.COLUMN_TYPE)); + } + + @Test + public void testGetSparkStructField_jsonbMetadata() { + StructField field = SpannerTableSchema.getSparkStructField("data", "jsonb", true, true); + assertTrue(field.metadata().contains(SpannerUtils.COLUMN_TYPE)); + assertEquals("jsonb", field.metadata().getString(SpannerUtils.COLUMN_TYPE)); + } } diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerUtilsTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerUtilsTest.java index 282cfafe..1200f335 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerUtilsTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SpannerUtilsTest.java @@ -14,12 +14,16 @@ package com.google.cloud.spark.spanner; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; +import java.util.HashMap; +import java.util.Map; import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.junit.Test; public class SpannerUtilsTest { @@ -142,4 +146,95 @@ public void testValidateSchemaNumericScaleMismatch() { .contains( "DataFrame has type decimal(38,10) but Spanner table expects type decimal(38,9)")); } + + @Test + public void testGetRequiredOptionCaseInsensitive() { + Map props = new HashMap<>(); + props.put("TABLE", "my_table"); + props.put("instanceId", "my_instance"); + CaseInsensitiveStringMap options = new CaseInsensitiveStringMap(props); + + assertEquals("my_table", SpannerUtils.getRequiredOption(options, "table")); + assertEquals("my_table", SpannerUtils.getRequiredOption(options, "TABLE")); + assertEquals("my_table", SpannerUtils.getRequiredOption(options, "Table")); + + assertEquals("my_instance", SpannerUtils.getRequiredOption(options, "instanceid")); + assertEquals("my_instance", SpannerUtils.getRequiredOption(options, "INSTANCEID")); + assertEquals("my_instance", SpannerUtils.getRequiredOption(options, "instanceId")); + } + + @Test + public void testGetRequiredOptionMissing() { + Map props = new HashMap<>(); + CaseInsensitiveStringMap options = new CaseInsensitiveStringMap(props); + + SpannerConnectorException e = + assertThrows( + SpannerConnectorException.class, + () -> SpannerUtils.getRequiredOption(options, "table")); + assertTrue(e.getMessage().contains("Option 'table' property must be set")); + } + + @Test + public void testValidateProjectIdAcceptsStandard() { + SpannerUtils.validateProjectId("my-project-123"); + } + + @Test + public void testValidateProjectIdAcceptsDomainScoped() { + SpannerUtils.validateProjectId("example.com:my-project"); + } + + @Test + public void testValidateProjectIdAcceptsNumeric() { + SpannerUtils.validateProjectId("123456789"); + } + + @Test + public void testValidateProjectIdRejectsUriInjection() { + SpannerConnectorException e = + assertThrows( + SpannerConnectorException.class, + () -> SpannerUtils.validateProjectId("proj?inject=true")); + assertTrue(e.getMessage().contains("Invalid projectId")); + } + + @Test + public void testValidateResourceIdAcceptsValid() { + SpannerUtils.validateResourceId("instance-1", "instanceId"); + SpannerUtils.validateResourceId("my_database", "databaseId"); + } + + @Test + public void testValidateResourceIdRejectsUriInjection() { + SpannerConnectorException e = + assertThrows( + SpannerConnectorException.class, + () -> SpannerUtils.validateResourceId("db?autoConfigEmulator=true", "databaseId")); + assertTrue(e.getMessage().contains("Invalid databaseId")); + } + + @Test + public void testValidateResourceIdRejectsSemicolonInjection() { + SpannerConnectorException e = + assertThrows( + SpannerConnectorException.class, + () -> SpannerUtils.validateResourceId("db;usePlainText=true", "databaseId")); + assertTrue(e.getMessage().contains("Invalid databaseId")); + } + + @Test + public void testValidateEmulatorHostAcceptsValid() { + SpannerUtils.validateEmulatorHost("localhost:9010"); + SpannerUtils.validateEmulatorHost("127.0.0.1:9010"); + } + + @Test + public void testValidateEmulatorHostRejectsUriInjection() { + SpannerConnectorException e = + assertThrows( + SpannerConnectorException.class, + () -> SpannerUtils.validateEmulatorHost("localhost:9010/evil?param=val")); + assertTrue(e.getMessage().contains("Invalid emulatorHost")); + } } diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBaseTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBaseTest.java new file mode 100644 index 00000000..a3fcc439 --- /dev/null +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/SparkSpannerTableProviderBaseTest.java @@ -0,0 +1,191 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import com.google.gson.Gson; +import com.google.gson.reflect.TypeToken; +import java.util.HashMap; +import java.util.Map; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class SparkSpannerTableProviderBaseTest { + + private static final Gson GSON = new Gson(); + + // A test implementation of the abstract SparkSpannerTableProviderBase + private static class TestSparkSpannerTableProvider extends SparkSpannerTableProviderBase {} + + @Test + public void testExtractIdentifier() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("projectId", "p"); + put("instanceId", "i"); + put("databaseId", "d"); + put("table", "t"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + assertEquals(Identifier.of(new String[0], "t"), identifier); + } + + @Test + public void testExtractIdentifierWithTableAndGraph() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("projectId", "p"); + put("instanceId", "i"); + put("databaseId", "d"); + put("table", "t"); + put("graph", "g"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + assertEquals(Identifier.of(new String[0], "t"), identifier); + } + + @Test + public void testExtractIdentifierWithNoTable() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("projectId", "p"); + put("instanceId", "i"); + put("databaseId", "d"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + assertNull(identifier); + } + + @Test + public void testExtractIdentifierWithGraphOnly() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("projectId", "p"); + put("instanceId", "i"); + put("databaseId", "d"); + put("graph", "MyGraph"); + put("type", "node"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + assertNotNull(identifier); + assertTrue(identifier.name().startsWith(SpannerCatalog.GRAPH_IDENTIFIER_PREFIX)); + String json = identifier.name().substring(SpannerCatalog.GRAPH_IDENTIFIER_PREFIX.length()); + Map decoded = + GSON.fromJson(json, new TypeToken>() {}.getType()); + assertEquals("MyGraph", decoded.get("graph")); + assertEquals("node", decoded.get("type")); + assertEquals(0, identifier.namespace().length); + } + + @Test + public void testExtractIdentifierGraphEncodesAllOptions() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("graph", "G"); + put("type", "node"); + put("enableDataBoost", "true"); + put("configs", "{}"); + put("graphQuery", "SELECT 1"); + put("timestamp", "2024-01-01T00:00:00Z"); + put("viewsEnabled", "true"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + String json = identifier.name().substring(SpannerCatalog.GRAPH_IDENTIFIER_PREFIX.length()); + Map decoded = + GSON.fromJson(json, new TypeToken>() {}.getType()); + assertEquals(7, decoded.size()); + assertEquals("G", decoded.get("graph")); + assertEquals("true", decoded.get("enableDataBoost")); + assertEquals("{}", decoded.get("configs")); + assertEquals("SELECT 1", decoded.get("graphQuery")); + assertEquals("2024-01-01T00:00:00Z", decoded.get("timestamp")); + assertEquals("true", decoded.get("viewsEnabled")); + } + + @Test + public void testExtractIdentifierGraphEncodesOnlyPresentOptions() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("graph", "G"); + put("projectId", "p"); + } + }); + Identifier identifier = provider.extractIdentifier(options); + String json = identifier.name().substring(SpannerCatalog.GRAPH_IDENTIFIER_PREFIX.length()); + Map decoded = + GSON.fromJson(json, new TypeToken>() {}.getType()); + assertEquals(1, decoded.size()); + assertEquals("G", decoded.get("graph")); + assertNull(decoded.get("projectId")); + } + + @Test + public void testShortName() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + assertEquals("cloud-spanner", provider.shortName()); + } + + @Test + public void testSupportsExternalMetadata() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + assertTrue(provider.supportsExternalMetadata()); + } + + @Test + public void testExtractCatalog() { + TestSparkSpannerTableProvider provider = new TestSparkSpannerTableProvider(); + CaseInsensitiveStringMap options = + new CaseInsensitiveStringMap( + new HashMap() { + { + put("catalog", "cat"); + } + }); + assertEquals("cat", provider.extractCatalog(options)); + CaseInsensitiveStringMap options2 = new CaseInsensitiveStringMap(new HashMap<>()); + assertEquals("spark_catalog", provider.extractCatalog(options2)); + } +} diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocAcceptanceTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocAcceptanceTestBase.java index 988fd74f..26622aa9 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocAcceptanceTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocAcceptanceTestBase.java @@ -183,6 +183,11 @@ protected static AcceptanceTestContext setup( Map properties = clusterProperties.stream() .collect(Collectors.toMap(ClusterProperty::getKey, ClusterProperty::getValue)); + properties.put( + "spark:spark.sql.catalog.spanner", "com.google.cloud.spark.spanner.SpannerCatalog"); + properties.put("spark:spark.sql.catalog.spanner.projectId", PROJECT_ID); + properties.put("spark:spark.sql.catalog.spanner.instanceId", INSTANCE_ID); + properties.put("spark:spark.sql.catalog.spanner.databaseId", DATABASE_ID); String testBaseGcsDir = AcceptanceTestUtils.createTestBaseGcsDir(testId); String connectorJarUri = testBaseGcsDir + "/connector.jar"; AcceptanceTestUtils.uploadConnectorJar( diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocServerlessAcceptanceTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocServerlessAcceptanceTestBase.java index 1cdd74f4..2964efc2 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocServerlessAcceptanceTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/acceptance/DataprocServerlessAcceptanceTestBase.java @@ -210,11 +210,18 @@ protected OperationSnapshot createAndRunPythonBatch( context.getScriptUri(testName), "text/x-python"); String parent = String.format("projects/%s/locations/%s", PROJECT_ID, REGION); + RuntimeConfig.Builder runtimeConfigBuilder = + RuntimeConfig.newBuilder().setVersion(s8sImageVersion); + runtimeConfigBuilder.putProperties( + "spark.sql.catalog.spanner", "com.google.cloud.spark.spanner.SpannerCatalog"); + runtimeConfigBuilder.putProperties("spark.sql.catalog.spanner.projectId", PROJECT_ID); + runtimeConfigBuilder.putProperties("spark.sql.catalog.spanner.instanceId", INSTANCE_ID); + runtimeConfigBuilder.putProperties("spark.sql.catalog.spanner.databaseId", DATABASE_ID); Batch batch = Batch.newBuilder() .setName(parent + "/batches/" + context.clusterId) .setPysparkBatch(createPySparkBatchBuilder(context, testName, pythonZipUri, args)) - .setRuntimeConfig(RuntimeConfig.newBuilder().setVersion(s8sImageVersion)) + .setRuntimeConfig(runtimeConfigBuilder) .build(); OperationFuture batchAsync = diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/graph/GraphReadIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/graph/GraphReadIntegrationTestBase.java index 730b5928..45d1177c 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/graph/GraphReadIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/graph/GraphReadIntegrationTestBase.java @@ -1,12 +1,41 @@ +// Copyright 2026 Google LLC +// +// 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. package com.google.cloud.spark.spanner.graph; import com.google.cloud.spark.spanner.integration.SparkSpannerIntegrationTestBase; import com.google.gson.Gson; +import java.util.Map; import org.apache.spark.sql.DataFrameReader; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; public class GraphReadIntegrationTestBase extends SparkSpannerIntegrationTestBase { + public DataFrameReader reader() { + Map props = connectionProperties(); + DataFrameReader reader = + spark + .read() + .format("cloud-spanner") + .option("viewsEnabled", true) + .option("projectId", props.get("projectId")) + .option("instanceId", props.get("instanceId")) + .option("databaseId", props.get("databaseId")); + String emulatorHost = props.get("emulatorHost"); + if (emulatorHost != null) reader = reader.option("emulatorHost", props.get("emulatorHost")); + return reader; + } + public DataFrameReader flexibleGraphReader(SpannerGraphConfigs configs) { DataFrameReader reader = reader().option("enableDataBoost", "true").option("graph", "FlexibleGraph"); diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/CatalogWriteIntegrationTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/CatalogWriteIntegrationTest.java new file mode 100644 index 00000000..553f7f86 --- /dev/null +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/CatalogWriteIntegrationTest.java @@ -0,0 +1,143 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner.integration; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.google.cloud.spark.spanner.TestData; +import java.util.Arrays; +import java.util.Collection; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +@RunWith(Parameterized.class) +public class CatalogWriteIntegrationTest extends SparkCatalogSpannerIntegrationTestBase { + private final boolean usePostgresSql; + + public CatalogWriteIntegrationTest(boolean usePostgresSql) { + super(); + this.usePostgresSql = usePostgresSql; + } + + @Override + protected boolean getUsePostgreSql() { + return usePostgresSql; + } + + @Parameterized.Parameters + public static Collection usePostgresSqlValues() { + return Arrays.asList(new Object[][] {{false}, {true}}); + } + + @Test + public void testCreateTableWithArrayColumns() { + String tableName = TestData.WRITE_TABLE_NAME + "_ARR_DDL"; + spark.sql("DROP TABLE IF EXISTS spanner." + tableName); + + // 1. Create the table with array columns. + spark.sql( + "CREATE TABLE spanner." + + tableName + + " (id BIGINT NOT NULL, long_array ARRAY, str_array ARRAY) USING `cloud-spanner`" + + " TBLPROPERTIES('primaryKeys' = 'id')"); + + // 2. Insert data with array values. + spark.sql( + "INSERT INTO spanner." + + tableName + + " VALUES (1, ARRAY(CAST(10 AS BIGINT), CAST(20 AS BIGINT), CAST(30 AS BIGINT)), ARRAY('hello', 'world'))"); + + // 3. Verify the data. + Dataset readBack = spark.sql("SELECT * FROM spanner." + tableName).filter("id = 1"); + assertEquals(1, readBack.count()); + + Row row = readBack.first(); + assertThat(row.getLong(0)).isEqualTo(1L); + Long[] longArr = new Long[] {10L, 20L, 30L}; + Object[] strArr = new Object[] {"hello", "world"}; + assertArrayEquals(longArr, row.getList(1).toArray(new Long[0])); + assertArrayEquals(strArr, row.getList(2).toArray()); + } + + @Test + public void testIgnoreSaveMode() { + String tableName = TestData.WRITE_TABLE_NAME + "_IGNORE"; + spark.sql("DROP TABLE IF EXISTS spanner." + tableName); + + String createSql = + "CREATE TABLE IF NOT EXISTS spanner." + + tableName + + " (long_col BIGINT NOT NULL, string_col STRING) USING `cloud-spanner`" + + " TBLPROPERTIES('primaryKeys' = 'long_col')"; + + // 1. First CREATE TABLE IF NOT EXISTS creates the table. + spark.sql(createSql); + + // 2. Insert initial data. + spark.sql("INSERT INTO spanner." + tableName + " VALUES (501, 'initial-data')"); + + // 3. Verify the initial write. + Dataset dfAfterFirstCreate = spark.sql("SELECT * FROM spanner." + tableName); + assertEquals(1, dfAfterFirstCreate.count()); + assertEquals("initial-data", dfAfterFirstCreate.first().getString(1)); + + // 4. Second CREATE TABLE IF NOT EXISTS (Ignore mode) should be a no-op. + spark.sql(createSql); + + // 5. Verify that the table content is unchanged. + Dataset finalDf = spark.sql("SELECT * FROM spanner." + tableName); + assertEquals(1, finalDf.count()); + Row finalRow = finalDf.first(); + assertEquals(501L, finalRow.getLong(0)); + assertEquals("initial-data", finalRow.getString(1)); + } + + @Test + public void testErrorIfExistsSaveMode() throws TableAlreadyExistsException { + String tableName = TestData.WRITE_TABLE_NAME + "_EIE"; + spark.sql("DROP TABLE IF EXISTS spanner." + tableName); + + // 1. First writeTo().create() (ErrorIfExists) should succeed and create the table. + Dataset firstDf = + spark.sql("SELECT CAST(301 AS BIGINT) AS long_col, 'three-oh-one' AS string_col"); + String catalogTable = "spanner." + tableName; + firstDf.writeTo(catalogTable).tableProperty("primaryKeys", "long_col").create(); + + // 2. Verify the first write. + Dataset dfAfterFirstCreate = spark.sql("SELECT * FROM spanner." + tableName); + assertEquals(1, dfAfterFirstCreate.count()); + assertEquals("three-oh-one", dfAfterFirstCreate.first().getString(1)); + + // 3. Second writeTo().create() (ErrorIfExists) should fail since the table already exists. + Dataset secondDf = + spark.sql("SELECT CAST(302 AS BIGINT) AS long_col, 'three-oh-two' AS string_col"); + try { + secondDf.writeTo(catalogTable).tableProperty("primaryKeys", "long_col").create(); + fail("Expected exception was not thrown"); + } catch (Exception e) { + assertTrue( + "Expected exception message about table already exists, but got: " + e.getMessage(), + e.getMessage().contains("already exists")); + } + } +} diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/FunctionsAndExpressionsIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/FunctionsAndExpressionsIntegrationTestBase.java index d474fb83..fe902926 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/FunctionsAndExpressionsIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/FunctionsAndExpressionsIntegrationTestBase.java @@ -84,14 +84,19 @@ public static Iterable data() { return ImmutableList.of(new Object[] {false}, new Object[] {true}); } - private boolean isPostgreSql; + private final boolean isPostgreSql; public FunctionsAndExpressionsIntegrationTestBase(boolean isPostgreSql) { this.isPostgreSql = isPostgreSql; } + @Override + protected boolean getUsePostgreSql() { + return isPostgreSql; + } + public Dataset readFromTable(String table) { - Map props = this.connectionProperties(isPostgreSql); + Map props = connectionProperties(); return spark .read() .format("cloud-spanner") diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/ReadIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/ReadIntegrationTestBase.java index d0be0c18..116a0af6 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/ReadIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/ReadIntegrationTestBase.java @@ -33,8 +33,10 @@ import java.util.Comparator; import java.util.List; import java.util.Locale; +import java.util.Map; import java.util.Optional; import java.util.stream.Collectors; +import org.apache.spark.sql.DataFrameReader; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Encoders; import org.apache.spark.sql.Row; @@ -49,7 +51,14 @@ public class ReadIntegrationTestBase extends SparkSpannerIntegrationTestBase { public Dataset readFromTable(String table) { - return reader().option("table", table).load(); + + Map props = connectionProperties(); + DataFrameReader reader = spark.read().format("cloud-spanner"); + String emulatorHost = props.get("emulatorHost"); + if (emulatorHost != null) { + reader = reader.option("emulatorHost", props.get("emulatorHost")); + } + return reader.option("table", table).load(); } @Test diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SchemaValidationIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SchemaValidationIntegrationTestBase.java index 57db88e5..e85cc084 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SchemaValidationIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SchemaValidationIntegrationTestBase.java @@ -17,9 +17,8 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; -import com.google.cloud.spark.spanner.SpannerConnectorException; -import com.google.cloud.spark.spanner.SpannerErrorCode; import com.google.cloud.spark.spanner.TestData; +import java.math.RoundingMode; import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -54,28 +53,31 @@ public SchemaValidationIntegrationTestBase(boolean usePostgreSql) { this.usePostgreSql = usePostgreSql; } + @Override + protected boolean getUsePostgreSql() { + return usePostgreSql; + } + @Test public void testColumnNotFound() { StructType dfSchema = new StructType( new StructField[] { DataTypes.createStructField("id", DataTypes.LongType, false), + DataTypes.createStructField("name", DataTypes.StringType, true), + DataTypes.createStructField("value", DataTypes.DoubleType, true), DataTypes.createStructField("non_existent_col", DataTypes.StringType, true) }); - List rows = Collections.singletonList(RowFactory.create(1L, "some_string")); + List rows = Collections.singletonList(RowFactory.create(1L, "a", 1.1, "some_string")); Dataset df = spark.createDataFrame(rows, dfSchema); Map props = connectionProperties(usePostgreSql); props.put("table", SCHEMA_VALIDATION_TABLE_NAME); props.put("enablePartialRowUpdates", "true"); - SpannerConnectorException e = - assertThrows( - SpannerConnectorException.class, - () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); - - assertThat(e.getErrorCode()).isEqualTo(SpannerErrorCode.SCHEMA_VALIDATION_ERROR); - assertThat(e.getMessage()).contains("DataFrame column 'non_existent_col' not found"); + assertThrows( + Exception.class, + () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); } @Test @@ -83,22 +85,20 @@ public void testDataTypeMismatch() { StructType dfSchema = new StructType( new StructField[] { - DataTypes.createStructField("id", DataTypes.StringType, false) // wrong type + DataTypes.createStructField("id", DataTypes.StringType, false), // wrong type + DataTypes.createStructField("name", DataTypes.StringType, true), + DataTypes.createStructField("value", DataTypes.DoubleType, true), }); - List rows = Collections.singletonList(RowFactory.create("not_a_long")); + List rows = Collections.singletonList(RowFactory.create("not_a_long", "a", 1.1)); Dataset df = spark.createDataFrame(rows, dfSchema); Map props = connectionProperties(usePostgreSql); props.put("table", SCHEMA_VALIDATION_TABLE_NAME); props.put("enablePartialRowUpdates", "true"); - SpannerConnectorException e = - assertThrows( - SpannerConnectorException.class, - () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); - - assertThat(e.getErrorCode()).isEqualTo(SpannerErrorCode.SCHEMA_VALIDATION_ERROR); - assertThat(e.getMessage()).contains("Data type mismatch for column 'id'"); + assertThrows( + Exception.class, + () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); } @Test @@ -130,31 +130,39 @@ public void testNumericScaleMismatchFailsInGoogleSql() { // PostgreSQL's NUMERIC is more flexible and this write may not fail. return; } - + spark.conf().set("spark.sql.storeAssignmentPolicy", "STRICT"); // Create a DataFrame with a high-scale decimal type. StructType dfSchema = new StructType( new StructField[] { DataTypes.createStructField("long_col", DataTypes.LongType, false), - DataTypes.createStructField("numeric_col", DataTypes.createDecimalType(38, 10), true) + DataTypes.createStructField("string_col", DataTypes.StringType, true), + DataTypes.createStructField("bool_col", DataTypes.BooleanType, true), + DataTypes.createStructField("bytes_col", DataTypes.BinaryType, true), + DataTypes.createStructField("double_col", DataTypes.DoubleType, true), + DataTypes.createStructField("numeric_col", DataTypes.createDecimalType(38, 10), true), + DataTypes.createStructField("timestamp_col", DataTypes.TimestampType, true), + DataTypes.createStructField("date_col", DataTypes.DateType, true) }); List rows = Collections.singletonList( - RowFactory.create(999L, new java.math.BigDecimal("123.1234567890"))); + RowFactory.create( + 999L, + "a", + true, + "a".getBytes(), + 1.1, + new java.math.BigDecimal("0.0000000001").setScale(10, RoundingMode.HALF_UP), + null, + null)); Dataset df = spark.createDataFrame(rows, dfSchema); Map props = connectionProperties(usePostgreSql); props.put("table", TestData.WRITE_TABLE_NAME); props.put("enablePartialRowUpdates", "true"); // Must be true to test connector validation. - SpannerConnectorException e = - assertThrows( - SpannerConnectorException.class, - () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); - - assertThat(e.getErrorCode()).isEqualTo(SpannerErrorCode.SCHEMA_VALIDATION_ERROR); - assertThat(e.getMessage()).contains("Data type mismatch for column 'numeric_col'"); - assertThat(e.getMessage()) - .contains("DataFrame has type decimal(38,10) but Spanner table expects type decimal(38,9)"); + assertThrows( + Exception.class, + () -> df.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save()); } } diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerCatalogIntegrationTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerCatalogIntegrationTest.java new file mode 100644 index 00000000..f98fdb8b --- /dev/null +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerCatalogIntegrationTest.java @@ -0,0 +1,195 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner.integration; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.cloud.spark.spanner.SpannerCatalog; +import com.google.cloud.spark.spanner.SpannerConnectorException; +import com.google.cloud.spark.spanner.SpannerTable; +import java.util.Arrays; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; + +@RunWith(Parameterized.class) +public class SpannerCatalogIntegrationTest extends SparkCatalogSpannerIntegrationTestBase { + + private SpannerCatalog catalog; + private final boolean usePostgresSql; + + @Parameters + public static Collection usePostgresSqlValues() { + return Arrays.asList(new Object[][] {{false}, {true}}); + } + + public SpannerCatalogIntegrationTest(boolean usePostgresSql) { + super(); + this.usePostgresSql = usePostgresSql; + } + + @Before + public void setupCatalog() { + catalog = new SpannerCatalog(); + catalog.initialize( + "spanner", new CaseInsensitiveStringMap(connectionProperties(usePostgresSql))); + } + + @After + public void teardownCatalog() { + catalog.close(); + } + + @Test + public void testListTables() { + String[] namespace = new String[0]; + Identifier[] tables = catalog.listTables(namespace); + List tableNames = + Arrays.stream(tables) + .map(Identifier::name) + .map(String::toLowerCase) + .collect(Collectors.toList()); + + assertThat(tableNames) + .containsAtLeast("schema_test_table", "write_array_test_table", "write_test_table"); + } + + @Test + public void testLoadTable() throws NoSuchTableException { + Identifier ident = Identifier.of(new String[0], "schema_test_table"); + Table table = catalog.loadTable(ident); + assertTrue(table instanceof SpannerTable); + assertThat(table.name()).isEqualTo("schema_test_table"); + assertThat(table.schema().fields()) + .asList() + .containsExactly( + new StructField("id", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA), + new StructField("name", DataTypes.StringType, true, Metadata.empty()), + new StructField("value", DataTypes.DoubleType, true, Metadata.empty())); + } + + @Test + public void testLoadTableNotExists() { + Identifier ident = Identifier.of(new String[0], "NonExistentTable"); + assertThrows(NoSuchTableException.class, () -> catalog.loadTable(ident)); + } + + @Test + public void testCreateTableAlreadyExists() { + Identifier ident = Identifier.of(new String[0], "write_test_table"); + assertThrows( + SpannerConnectorException.class, + () -> catalog.createTable(ident, new StructType(), null, new HashMap<>())); + } + + @Test + public void testTableExists() { + Identifier ident = Identifier.of(new String[0], "write_test_table"); + assertTrue(catalog.tableExists(ident)); + } + + @Test + public void testCreateTable() throws NoSuchTableException, TableAlreadyExistsException { + String tableName = "new_test_table"; + Identifier ident = Identifier.of(new String[0], tableName); + StructType createSchema = + new StructType() + .add("id", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA) + .add("name", DataTypes.StringType, true); + Map properties = new HashMap<>(); + + try { + catalog.createTable(ident, createSchema, null, properties); + assertTrue(catalog.tableExists(ident)); + Table loadedTable = catalog.loadTable(ident); + // Connector currently does not retrieve primary key metadata. + StructType expectedSchema = + new StructType() + .add("id", DataTypes.LongType, false, SpannerCatalog.PRIMARY_KEY_METADATA) + .add("name", DataTypes.StringType, true); + assertThat(loadedTable.schema()).isEqualTo(expectedSchema); + } finally { + catalog.dropTable(ident); + assertFalse(catalog.tableExists(ident)); + } + } + + @Test + public void testTableExistsReturnsFalse() { + Identifier ident = Identifier.of(new String[0], "non_existent_table"); + assertFalse(catalog.tableExists(ident)); + } + + @Test + public void testLoadTableRejectsNonEmptyNamespace() { + Identifier ident = Identifier.of(new String[] {"ns"}, "schema_test_table"); + assertThrows(SpannerConnectorException.class, () -> catalog.loadTable(ident)); + } + + @Test + public void testListTablesWithInvalidNamespace() { + Identifier[] tables = catalog.listTables(new String[] {"invalid", "namespace"}); + assertEquals(0, tables.length); + } + + @Test + public void testDropTableNonExistent() { + Identifier ident = Identifier.of(new String[0], "table_that_does_not_exist"); + assertFalse(catalog.dropTable(ident)); + } + + @Test + public void testIsGraphIdentifier() { + Identifier graphIdent = + Identifier.of(new String[0], SpannerCatalog.GRAPH_IDENTIFIER_PREFIX + "{\"graph\":\"G\"}"); + assertTrue(SpannerCatalog.isGraphIdentifier(graphIdent)); + + Identifier tableIdent = Identifier.of(new String[0], "my_table"); + assertFalse(SpannerCatalog.isGraphIdentifier(tableIdent)); + } + + @Test + public void testReadTableViaCatalogSql() { + if (usePostgresSql) { + return; + } + Dataset df = spark.sql("SELECT * FROM spanner.simpleTable"); + assertThat(df.count()).isGreaterThan(0); + assertThat(df.columns()).asList().containsExactly("A", "B", "C"); + } +} diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerScanBuilderIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerScanBuilderIntegrationTestBase.java index 9646266c..df99bb6d 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerScanBuilderIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerScanBuilderIntegrationTestBase.java @@ -17,11 +17,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; -import com.google.cloud.spark.spanner.SpannerInputPartitionReaderContext; -import com.google.cloud.spark.spanner.SpannerPartitionReader; -import com.google.cloud.spark.spanner.SpannerScanBuilder; -import com.google.cloud.spark.spanner.SpannerScanner; -import com.google.cloud.spark.spanner.SpannerUtils; +import com.google.cloud.spark.spanner.*; import java.io.IOException; import java.time.ZonedDateTime; import java.util.ArrayList; @@ -41,7 +37,6 @@ import org.apache.spark.sql.types.MetadataBuilder; import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; -import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -49,11 +44,20 @@ @RunWith(JUnit4.class) public abstract class SpannerScanBuilderIntegrationTestBase extends SpannerTestBase { + private static SpannerTable getSpannerTable(boolean usePostgreSql) { + Map connectionProperties = connectionProperties(usePostgreSql); + return new SpannerTable(connectionProperties); + } + + protected static SpannerTable getSpannerTable(String tableName, boolean usePostgreSql) { + Map connectionProperties = connectionProperties(usePostgreSql); + connectionProperties.put("table", tableName); + return new SpannerTable(connectionProperties); + } + @Test public void readSchemaShouldWorkInSpannerScanBuilder() throws Exception { - Map opts = this.connectionProperties(); - CaseInsensitiveStringMap copts = new CaseInsensitiveStringMap(opts); - Scan scan = new SpannerScanBuilder(copts).build(); + Scan scan = new SpannerScanBuilder(getSpannerTable(false)).build(); MetadataBuilder jsonMetaBuilder = new MetadataBuilder(); jsonMetaBuilder.putString(SpannerUtils.COLUMN_TYPE, "json"); StructType actualSchema = scan.readSchema(); @@ -87,9 +91,7 @@ public void readSchemaShouldWorkInSpannerScanBuilderForPg() throws Exception { "readSchemaShouldWorkInSpannerScanBuilderForPg is skipped since pg is not supported in Spanner emulator"); return; } - Map opts = this.connectionProperties(/* usePostgreSql= */ true); - CaseInsensitiveStringMap copts = new CaseInsensitiveStringMap(opts); - Scan scan = new SpannerScanBuilder(copts).build(); + Scan scan = new SpannerScanBuilder(getSpannerTable(true)).build(); StructType actualSchema = scan.readSchema(); MetadataBuilder jsonMetaBuilder = new MetadataBuilder(); jsonMetaBuilder.putString(SpannerUtils.COLUMN_TYPE, "jsonb"); @@ -125,10 +127,8 @@ public void readSchemaShouldWorkInSpannerScanBuilderForPg() throws Exception { @Test public void planInputPartitionsShouldSuccessInSpannerScanBuilder() throws Exception { - Map opts = this.connectionProperties(); - opts.put("table", "ATable"); - CaseInsensitiveStringMap optionMap = new CaseInsensitiveStringMap(opts); - SpannerScanBuilder spannerScanBuilder = new SpannerScanBuilder(optionMap); + SpannerScanBuilder spannerScanBuilder = + new SpannerScanBuilder(getSpannerTable("ATable", false)); SpannerScanner ss = ((SpannerScanner) spannerScanBuilder.build()); InputPartition[] partitions = ss.planInputPartitions(); PartitionReaderFactory prf = ss.createReaderFactory(); @@ -202,10 +202,8 @@ public void planInputPartitionsShouldSucceedInSpannerScanBuilderPg() throws Exce "planInputPartitionsShouldSucceedInSpannerScanBuilderPg is skipped since pg is not supported in Spanner emulator"); return; } - Map opts = this.connectionProperties(true); - opts.put("table", "composite_table"); - CaseInsensitiveStringMap optionMap = new CaseInsensitiveStringMap(opts); - SpannerScanBuilder spannerScanBuilder = new SpannerScanBuilder(optionMap); + SpannerScanBuilder spannerScanBuilder = + new SpannerScanBuilder(getSpannerTable("composite_table", true)); SpannerScanner ss = ((SpannerScanner) spannerScanBuilder.build()); InputPartition[] partitions = ss.planInputPartitions(); PartitionReaderFactory prf = ss.createReaderFactory(); diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerTestBase.java index 7cf3b006..661b479c 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SpannerTestBase.java @@ -54,12 +54,16 @@ import org.slf4j.LoggerFactory; class SpannerTestBase { + private static final boolean spannerUseExistingDb = + Boolean.parseBoolean(System.getenv("SPANNER_USE_EXISTING_DATABASE")); // Since in the teardown we delete the Cloud Spanner database, here we append a random value to // the database ID to avoid any cross-pollution between concurrently running tests. // Note that a database ID must be 2-30 characters long. private static final String databaseId = - System.getenv("SPANNER_DATABASE_ID") + "-" + new Random().nextInt(10000000); + spannerUseExistingDb + ? System.getenv("SPANNER_DATABASE_ID") + : System.getenv("SPANNER_DATABASE_ID") + "-" + new Random().nextInt(10000000); private static final String databaseIdPg = databaseId + "-pg"; private static final String instanceId = System.getenv("SPANNER_INSTANCE_ID"); private static final String projectId = System.getenv("SPANNER_PROJECT_ID"); @@ -167,6 +171,10 @@ public static void setUp() throws Exception { return; } + if (spannerUseExistingDb) { + return; + } + // Create the instance. InstanceAdminClient instanceAdminClient = spanner.getInstanceAdminClient(); @@ -205,6 +213,9 @@ public static void setUp() throws Exception { } private static void cleanupDatabase() { + if (spannerUseExistingDb) { + return; + } log.info("\033[33mCleaning up databases\033[00m"); DatabaseAdminClient databaseAdminClient = spanner.getDatabaseAdminClient(); databaseAdminClient.dropDatabase(instanceId, databaseId); @@ -234,26 +245,20 @@ protected static Map connectionProperties(boolean usePostgreSql) } protected static Map connectionPropertiesLowerCase(boolean usePostgreSql) { - Map props = new HashMap<>(); - if (usePostgreSql) { - props.put("databaseid", databaseIdPg); - props.put("table", tablePg); - } else { - props.put("databaseid", databaseId); - props.put("table", table); + Map props = connectionProperties(usePostgreSql); + Map lowerCasedProps = new HashMap<>(); + for (Map.Entry entry : props.entrySet()) { + lowerCasedProps.put(entry.getKey().toLowerCase(), entry.getValue()); } - props.put("instanceid", instanceId); - props.put("projectid", projectId); - if (emulatorHost != null) { - props.put("emulatorhost", emulatorHost); - } - props.put("enablepartialrowupdates", "true"); + return lowerCasedProps; + } - return props; + protected Map connectionProperties() { + return connectionProperties(getUsePostgreSql()); } - protected static Map connectionProperties() { - return connectionProperties(false); + protected boolean getUsePostgreSql() { + return false; } static InternalRow makeInternalRow(int A, String B, double C) { diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkCatalogSpannerIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkCatalogSpannerIntegrationTestBase.java new file mode 100644 index 00000000..72e13afc --- /dev/null +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkCatalogSpannerIntegrationTestBase.java @@ -0,0 +1,55 @@ +// Copyright 2026 Google LLC +// +// 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. + +package com.google.cloud.spark.spanner.integration; + +import java.util.Map; +import org.apache.spark.sql.SparkSession; +import org.junit.After; +import org.junit.Before; + +public class SparkCatalogSpannerIntegrationTestBase extends SpannerTestBase { + + protected SparkSession spark; + + public SparkCatalogSpannerIntegrationTestBase() {} + + @Before + public void setUpSpark() { + Map catalogProps = connectionProperties(); + spark = + SparkSession.builder() + .master("local") + .appName("SparkSpannerIntegrationTest") + .config("spark.ui.enabled", "false") + .config("spark.sql.catalog.spanner", "com.google.cloud.spark.spanner.SpannerCatalog") + .config("spark.sql.catalog.spanner.projectId", catalogProps.get("projectId")) + .config("spark.sql.catalog.spanner.instanceId", catalogProps.get("instanceId")) + .config("spark.sql.catalog.spanner.databaseId", catalogProps.get("databaseId")) + .config("spark.default.parallelism", 20) + .getOrCreate(); + + if (catalogProps.get("emulatorHost") != null) { + spark.conf().set("spark.sql.catalog.spanner.emulatorHost", catalogProps.get("emulatorHost")); + } + spark.sparkContext().setLogLevel("WARN"); + } + + @After + public void tearDownSpark() { + if (spark != null) { + spark.stop(); + } + } +} diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerIntegrationTestBase.java index 81bb79c3..44cb6135 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerIntegrationTestBase.java @@ -15,49 +15,37 @@ package com.google.cloud.spark.spanner.integration; import java.util.Map; -import org.apache.spark.sql.DataFrameReader; import org.apache.spark.sql.SparkSession; -import org.junit.ClassRule; -import org.junit.rules.ExternalResource; +import org.junit.After; +import org.junit.Before; public class SparkSpannerIntegrationTestBase extends SpannerTestBase { - @ClassRule public static SparkFactory sparkFactory = new SparkFactory(); - protected SparkSession spark; - public SparkSpannerIntegrationTestBase() { - this.spark = sparkFactory.spark; - } - - public DataFrameReader reader() { - Map props = connectionProperties(); - DataFrameReader reader = - spark - .read() - .format("cloud-spanner") - .option("viewsEnabled", true) - .option("projectId", props.get("projectId")) - .option("instanceId", props.get("instanceId")) - .option("databaseId", props.get("databaseId")); - String emulatorHost = props.get("emulatorHost"); - if (emulatorHost != null) reader = reader.option("emulatorHost", props.get("emulatorHost")); - return reader; + public SparkSpannerIntegrationTestBase() {} + + @Before + public void setUpSpark() { + Map catalogProps = connectionProperties(); + spark = + SparkSession.builder() + .master("local") + .appName("SparkSpannerIntegrationTest") + .config("spark.ui.enabled", "false") + .config("spark.default.parallelism", 20) + .getOrCreate(); + + if (catalogProps.get("emulatorHost") != null) { + spark.conf().set("spark.sql.catalog.spanner.emulatorHost", catalogProps.get("emulatorHost")); + } + spark.sparkContext().setLogLevel("WARN"); } - protected static class SparkFactory extends ExternalResource { - SparkSession spark; - - @Override - protected void before() throws Throwable { - spark = - SparkSession.builder() - .master("local") - .config("spark.ui.enabled", "false") - .config("spark.default.parallelism", 20) - .getOrCreate(); - // reducing test's logs - spark.sparkContext().setLogLevel("WARN"); + @After + public void tearDownSpark() { + if (spark != null) { + spark.stop(); } } } diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerTableProviderIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerTableProviderIntegrationTestBase.java index 992ca016..529f3fe6 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerTableProviderIntegrationTestBase.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/SparkSpannerTableProviderIntegrationTestBase.java @@ -55,7 +55,7 @@ public void getTableForWrite_withenablePartialRowUpdates_returnsTableWithDataFra } @Test - public void getTableForWrite_withoutenablePartialRowUpdates_returnsTableWithFullSchema() { + public void getTableForWrite_withoutEnablePartialRowUpdates_returnsTableWithFullSchema() { // Arrange T provider = getInstance(); Map props = connectionProperties(); @@ -88,7 +88,7 @@ public void getTableAllowsLowerCaseProperties() { // Arrange T provider = getInstance(); Map props = connectionPropertiesLowerCase(false); - + props.put("enablePartialRowUpdates", "true"); final StructType partialSchema = new StructType().add("long_col", DataTypes.LongType, false); // Act diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java index a70bedaa..69444089 100644 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java +++ b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTest.java @@ -17,6 +17,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.apache.spark.sql.functions.col; import static org.apache.spark.sql.functions.from_json; +import static org.apache.spark.sql.functions.lit; import static org.apache.spark.sql.functions.to_json; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -64,6 +65,12 @@ public static Collection usePostgresSqlValues() { new StructField[] { DataTypes.createStructField("long_col", DataTypes.LongType, false), DataTypes.createStructField("string_col", DataTypes.StringType, true), + DataTypes.createStructField("bool_col", DataTypes.BooleanType, true), + DataTypes.createStructField("double_col", DataTypes.DoubleType, true), + DataTypes.createStructField("timestamp_col", DataTypes.TimestampType, true), + DataTypes.createStructField("date_col", DataTypes.DateType, true), + DataTypes.createStructField("bytes_col", DataTypes.BinaryType, true), + DataTypes.createStructField("numeric_col", DataTypes.createDecimalType(38, 9), true), }); private Map getBaseProps() { @@ -85,24 +92,24 @@ private Dataset setupInitialData(StructType schema, long[] ids) { if (ids.length != 2) Assert.fail("Invalid number of id's provided"); List initialRows = Arrays.asList( - RowFactory.create(ids[0], "original twenty-one"), - RowFactory.create(ids[1], "original twenty-two")); + RowFactory.create(ids[0], "original twenty-one", null, null, null, null, null, null), + RowFactory.create(ids[1], "original twenty-two", null, null, null, null, null, null)); Dataset initialDf = spark.createDataFrame(initialRows, schema); initialDf.write().format("cloud-spanner").options(getBaseProps()).mode(SaveMode.Append).save(); return initialDf; } - private Dataset setupInitialData(StructType schema) { - final long[] ids = new long[] {201L, 202L}; - return setupInitialData(schema, ids); - } - public WriteIntegrationTest(boolean usePostgresSql) { super(); this.usePostgresSql = usePostgresSql; } + @Override + protected boolean getUsePostgreSql() { + return usePostgresSql; + } + @Test public void testWriteWithNulls() { StructType schema = @@ -163,7 +170,10 @@ public void testWriteWithNulls() { @Test public void testIdempotentWrite() { - List rows = Arrays.asList(RowFactory.create(4L, "four"), RowFactory.create(5L, "five")); + List rows = + Arrays.asList( + RowFactory.create(4L, "four", null, null, null, null, null, null), + RowFactory.create(5L, "five", null, null, null, null, null, null)); Dataset df = spark.createDataFrame(rows, SCHEMA); @@ -191,7 +201,8 @@ public void testEmptyDataFrameWrite() { Dataset df = spark.createDataFrame(Collections.emptyList(), SCHEMA); - Map props = getBaseProps(); + Map props = connectionProperties(usePostgresSql); + props.put("table", TestData.WRITE_TABLE_NAME); // Get initial count to ensure no new rows are added long initialCount = spark.read().format("cloud-spanner").options(props).load().count(); @@ -207,30 +218,56 @@ public void testEmptyDataFrameWrite() { @Test public void testUpsert() { // 1. Write initial data using unique keys. - setupInitialData(SCHEMA); - // 2. Write a second DataFrame to update one row and insert another. - List newRows = + java.sql.Timestamp ts = new java.sql.Timestamp(System.currentTimeMillis()); + java.sql.Date dt = new java.sql.Date(System.currentTimeMillis()); + byte[] b = "spanner".getBytes(java.nio.charset.StandardCharsets.UTF_8); + java.math.BigDecimal num = new java.math.BigDecimal("123.456"); + + // 2. Write the initial data (all columns populated) + List initialRows = Arrays.asList( - RowFactory.create(201L, "new twenty-one"), // Update 201 - RowFactory.create(203L, "new twenty-three") // Insert 203 - ); - spark - .createDataFrame(newRows, SCHEMA) - .write() - .format("cloud-spanner") - .options(getBaseProps()) - .mode(SaveMode.Append) - .save(); + RowFactory.create(201L, "original twenty-one", true, 1.5, ts, dt, b, num), + RowFactory.create(202L, "original twenty-two", false, 2.5, ts, dt, b, num)); + Dataset initialDf = spark.createDataFrame(initialRows, SCHEMA); + + Map props = connectionProperties(usePostgresSql); + props.put("table", TestData.WRITE_TABLE_NAME); + props.put("enablePartialRowUpdates", "true"); // Enable the Spanner-side upsert logic + + initialDf.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save(); + + // 3. Create the updates DataFrame. + // 1. Prepare the UPDATE: Filter for row 201 and update its string_col + // By filtering first, we can just use lit() to overwrite the column. + // This preserves the original values for bool_col, double_col, etc. + Dataset updatedRow201 = + initialDf + .filter(col("long_col").equalTo(201L)) + .withColumn("string_col", lit("new twenty-one")); + + // 2. Prepare the INSERT: Create row 301 from scratch using the full schema + // We pad the omitted columns with nulls to match the DSv2 table requirements. + List insertRows = + Collections.singletonList( + RowFactory.create(301L, "new thirty-one", null, null, null, null, null, null)); + Dataset insertedRow301 = spark.createDataFrame(insertRows, SCHEMA); + + // 3. Combine them into your final updatesDs + // unionByName is generally safer than union() as it matches columns by name + // rather than order, preventing accidental data shifts. + Dataset updatesDs = updatedRow201.unionByName(insertedRow301); + // Write the updates + updatesDs.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save(); - // 3. Verify the final state of the rows involved in this test. + // 5. Verify the final state Dataset finalDf = spark .read() .format("cloud-spanner") .options(getBaseProps()) .load() - .filter("long_col IN (201, 202, 203)"); + .filter("long_col IN (201, 202, 301)"); assertEquals(3, finalDf.count()); @@ -238,12 +275,23 @@ public void testUpsert() { finalDf.collectAsList().stream() .collect(java.util.stream.Collectors.toMap(r -> r.getLong(0), r -> r)); - // Check that row 201 was updated. - assertThat(finalRows.get(201L).getString(1)).isEqualTo("new twenty-one"); - // Check that row 202 was not touched. - assertThat(finalRows.get(202L).getString(1)).isEqualTo("original twenty-two"); - // Check that row 203 was inserted. - assertThat(finalRows.get(203L).getString(1)).isEqualTo("new twenty-three"); + Row row201 = finalRows.get(201L); + Row row202 = finalRows.get(202L); + Row row301 = finalRows.get(301L); + + // Verify row 201: string_col was updated, but all other columns remained untouched (Partial + // Update) + assertThat(row201.getString(1)).isEqualTo("new twenty-one"); + assertThat(row201.getBoolean(2)).isEqualTo(true); + assertThat(row201.getDouble(3)).isEqualTo(1.5); + + // Verify row 202: Completely untouched + assertThat(row202.getString(1)).isEqualTo("original twenty-two"); + assertThat(row202.getBoolean(2)).isEqualTo(false); + + // Verify row 203: Inserted, and Spanner/Spark correctly assigned NULLs to the omitted columns + assertThat(row301.getString(1)).isEqualTo("new thirty-one"); + assertTrue(row301.isNullAt(2)); // bool_col should be null } @Test @@ -254,8 +302,10 @@ public void testInsert() { // 2. Write a second DataFrame to update one row and insert another. List newRows = Arrays.asList( - RowFactory.create(211L, "new twenty-one"), // Update 211 - RowFactory.create(213L, "new twenty-three") // Insert 213 + RowFactory.create( + 211L, "new twenty-one", null, null, null, null, null, null), // Update 211 + RowFactory.create( + 213L, "new twenty-three", null, null, null, null, null, null) // Insert 213 ); // To make tests reproducable, set repartition to 1 so all rows are handled in an atomic @@ -276,7 +326,8 @@ public void testInsert() { // 4. Insert 213, happy path List successfulInsertRows = Collections.singletonList( - RowFactory.create(213L, "new twenty-three") // Insert 213 + RowFactory.create( + 213L, "new twenty-three", null, null, null, null, null, null) // Insert 213 ); Dataset successfulDf = spark.createDataFrame(successfulInsertRows, SCHEMA); successfulDf.write().format("cloud-spanner").options(insertProps).mode(SaveMode.Append).save(); @@ -312,8 +363,10 @@ public void testUpdate() { // 2. Write a second DataFrame to update an existing row and update a non-existent row. List errorRows = Arrays.asList( - RowFactory.create(221L, "new twenty-one"), // Update 221 - RowFactory.create(223L, "new twenty-three") // Update 223 + RowFactory.create( + 221L, "new twenty-one", null, null, null, null, null, null), // Update 221 + RowFactory.create( + 223L, "new twenty-three", null, null, null, null, null, null) // Update 223 ); // To make tests reproducable, set repartition to 1 so all rows are handled in atomic @@ -334,7 +387,8 @@ public void testUpdate() { // 4. Update existing 222, happy path List successfulUpdateRows = Collections.singletonList( - RowFactory.create(222L, "new twenty-two") // Update 222 + RowFactory.create( + 222L, "new twenty-two", null, null, null, null, null, null) // Update 222 ); Dataset successfulDf = spark.createDataFrame(successfulUpdateRows, SCHEMA); @@ -364,20 +418,27 @@ public void testUpdate() { @Test public void testReplace() { // 1. Write initial data using unique keys. - StructType schema = - new StructType( - new StructField[] { - DataTypes.createStructField("long_col", DataTypes.LongType, false), - DataTypes.createStructField("string_col", DataTypes.StringType, true), - DataTypes.createStructField("timestamp_col", DataTypes.TimestampType, true), - }); List initialRows = Arrays.asList( RowFactory.create( - 201L, "original twenty-one", java.sql.Timestamp.valueOf("2023-01-01 10:10:10")), + 201L, + "original twenty-one", + null, + null, + java.sql.Timestamp.valueOf("2023-01-01 10:10:10"), + null, + null, + null), RowFactory.create( - 202L, "original twenty-two", java.sql.Timestamp.valueOf("2024-01-01 10:10:10"))); - Dataset initialDf = spark.createDataFrame(initialRows, schema); + 202L, + "original twenty-two", + null, + null, + java.sql.Timestamp.valueOf("2024-01-01 10:10:10"), + null, + null, + null)); + Dataset initialDf = spark.createDataFrame(initialRows, SCHEMA); Map upsertProps = getBaseProps(); @@ -389,13 +450,23 @@ public void testReplace() { RowFactory.create( 201L, "new twenty-one", - java.sql.Timestamp.valueOf("2025-01-01 10:10:10")), // Replace 201 with 3 columns + null, + null, + java.sql.Timestamp.valueOf("2025-01-01 10:10:10"), + null, + null, + null), // Replace 201 RowFactory.create( 202L, "new twenty-two", - java.sql.Timestamp.valueOf("2026-01-01 10:10:10")) // Replace 202 with 3 columns + null, + null, + java.sql.Timestamp.valueOf("2026-01-01 10:10:10"), + null, + null, + null) // Replace 202 ); - Dataset newDf = spark.createDataFrame(newRows, schema); + Dataset newDf = spark.createDataFrame(newRows, SCHEMA); Map replaceProps = getBaseProps(); replaceProps.put("mutationType", "replace"); @@ -429,8 +500,8 @@ public void testReplace() { // 4. Write two rows with one column missing List shortRows = Arrays.asList( - RowFactory.create(201L, "short twenty-one"), - RowFactory.create(202L, "short twenty-two")); + RowFactory.create(201L, "short twenty-one", null, null, null, null, null, null), + RowFactory.create(202L, "short twenty-two", null, null, null, null, null, null)); Dataset shortDf = spark.createDataFrame(shortRows, SCHEMA); shortDf.write().format("cloud-spanner").options(replaceProps).mode(SaveMode.Append).save(); @@ -461,7 +532,9 @@ public void testReplace() { @Test public void testUpdateSetColumnToNull() { // 1. Write initial data with a non-null string value - List initialRows = Collections.singletonList(RowFactory.create(20L, "originalValue")); + List initialRows = + Collections.singletonList( + RowFactory.create(20L, "originalValue", null, null, null, null, null, null)); Dataset initialDf = spark.createDataFrame(initialRows, SCHEMA); Map props = connectionProperties(usePostgresSql); @@ -477,7 +550,8 @@ public void testUpdateSetColumnToNull() { assertThat(dfAfterInitialWrite.first().getString(1)).isEqualTo("originalValue"); // 2. Update the existing row, setting string_col to null - List updateRows = Collections.singletonList(RowFactory.create(20L, null)); + List updateRows = + Collections.singletonList(RowFactory.create(20L, null, null, null, null, null, null, null)); Dataset updateDf = spark.createDataFrame(updateRows, SCHEMA); updateDf.write().format("cloud-spanner").options(props).mode(SaveMode.Append).save(); @@ -487,6 +561,12 @@ public void testUpdateSetColumnToNull() { spark.read().format("cloud-spanner").options(props).load().filter("long_col = 20"); assertEquals(1, dfAfterUpdate.count()); assertNull(dfAfterUpdate.first().get(1)); + assertNull(dfAfterUpdate.first().get(2)); + assertNull(dfAfterUpdate.first().get(3)); + assertNull(dfAfterUpdate.first().get(4)); + assertNull(dfAfterUpdate.first().get(5)); + assertNull(dfAfterUpdate.first().get(6)); + assertNull(dfAfterUpdate.first().get(7)); } private void checkChildRow(Row row) { diff --git a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTestBase.java b/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTestBase.java deleted file mode 100644 index 475fb367..00000000 --- a/spark-3.1-spanner-lib/src/test/java/com/google/cloud/spark/spanner/integration/WriteIntegrationTestBase.java +++ /dev/null @@ -1,58 +0,0 @@ -// Copyright 2023 Google LLC -// -// 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. - -package com.google.cloud.spark.spanner.integration; - -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; - -import java.util.Map; -import org.apache.spark.SparkException; -import org.apache.spark.sql.DataFrameWriter; -import org.apache.spark.sql.Dataset; -import org.apache.spark.sql.Row; -import org.junit.Test; - -public class WriteIntegrationTestBase extends SparkSpannerIntegrationTestBase { - - private static ReadIntegrationTestBase readIntegration = new ReadIntegrationTestBase(); - private Map props = this.connectionProperties(); - - public DataFrameWriter writerToTable(Dataset df, String table) { - return df.write() - .format("cloud-spanner") - .option("viewsEnabled", true) - .option("projectId", props.get("projectId")) - .option("instanceId", props.get("instanceId")) - .option("databaseId", props.get("databaseId")) - .option("emulatorHost", props.get("emulatorHost")) - .option("table", table); - } - - @Test - public void testWritesToTableFail() { - String table = "compositeTable"; - Dataset drf = readIntegration.readFromTable(table); - SparkException e = - assertThrows( - SparkException.class, - () -> { - DataFrameWriter dwf = writerToTable(drf.select("id"), table); - dwf.saveAsTable(table); - }); - assertThat(e) - .hasMessageThat() - .isEqualTo("Table implementation does not support writes: default.compositeTable"); - } -} diff --git a/spark-3.1-spanner-lib/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/spark-3.1-spanner-lib/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister index ee1ae910..a99c9ec1 100644 --- a/spark-3.1-spanner-lib/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister +++ b/spark-3.1-spanner-lib/src/test/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister @@ -1 +1 @@ -com.google.cloud.spark.spanner.Spark31SpannerTableProvider +com.google.cloud.spark.spanner.Spark31SpannerTableProvider \ No newline at end of file diff --git a/spark-3.1-spanner/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/spark-3.1-spanner/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister index ee1ae910..a99c9ec1 100644 --- a/spark-3.1-spanner/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister +++ b/spark-3.1-spanner/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister @@ -1 +1 @@ -com.google.cloud.spark.spanner.Spark31SpannerTableProvider +com.google.cloud.spark.spanner.Spark31SpannerTableProvider \ No newline at end of file