This is an automated email from the ASF dual-hosted git repository.
Amar3tto pushed a commit to branch snowflakeio-yaml
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/snowflakeio-yaml by this push:
new 5bd11daeca4 Add Snowflake Read and Streaming Write
5bd11daeca4 is described below
commit 5bd11daeca49930698d14f331390d5726c00b1a3
Author: Vitaly Terentyev <[email protected]>
AuthorDate: Thu Aug 13 14:25:31 2026 +0400
Add Snowflake Read and Streaming Write
---
sdks/java/io/snowflake/build.gradle | 1 +
.../SnowflakeReadSchemaTransformProvider.java | 285 ++++++++++++++++
.../snowflake/SnowflakeSchemaTransformUtils.java | 323 ++++++++++++++++++
.../io/snowflake/SnowflakeWriteConfiguration.java | 224 ++++++++++++
.../SnowflakeWriteSchemaTransformProvider.java | 317 ++++-------------
.../SnowflakeReadSchemaTransformProviderTest.java | 311 +++++++++++++++++
.../SnowflakeWriteSchemaTransformProviderTest.java | 374 ++++++++++++++++-----
sdks/python/apache_beam/yaml/standard_io.yaml | 31 +-
8 files changed, 1523 insertions(+), 343 deletions(-)
diff --git a/sdks/java/io/snowflake/build.gradle
b/sdks/java/io/snowflake/build.gradle
index 8d9a9a46557..951af328fed 100644
--- a/sdks/java/io/snowflake/build.gradle
+++ b/sdks/java/io/snowflake/build.gradle
@@ -30,6 +30,7 @@ dependencies {
implementation project(path:
":sdks:java:extensions:google-cloud-platform-core")
permitUnusedDeclared project(path:
":sdks:java:extensions:google-cloud-platform-core")
implementation library.java.slf4j_api
+ implementation library.java.everit_json_schema
implementation group: 'net.snowflake', name: 'snowflake-jdbc', version:
'4.0.2'
implementation group: 'com.opencsv', name: 'opencsv', version: '5.12.0'
implementation 'net.snowflake:snowflake-ingest-sdk:4.4.2'
diff --git
a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProvider.java
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProvider.java
new file mode 100644
index 00000000000..f914b0f7dd9
--- /dev/null
+++
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProvider.java
@@ -0,0 +1,285 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.snowflake;
+
+import static
org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.toRow;
+
+import com.google.auto.service.AutoService;
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import java.util.Collections;
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.RowCoder;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
+import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
+import org.apache.beam.sdk.schemas.utils.JsonUtils;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
+
+/** A {@link SchemaTransformProvider} for reading rows from Snowflake. */
+@SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+@AutoService(SchemaTransformProvider.class)
+public class SnowflakeReadSchemaTransformProvider
+ extends
TypedSchemaTransformProvider<SnowflakeReadSchemaTransformProvider.Configuration>
{
+
+ static final String OUTPUT_TAG = "output";
+
+ public static final String IDENTIFIER =
"beam:schematransform:org.apache.beam:snowflake_read:v1";
+
+ @Override
+ public String identifier() {
+ return IDENTIFIER;
+ }
+
+ @Override
+ public String description() {
+ return "Reads rows from a Snowflake table or query using staged CSV
files.";
+ }
+
+ @Override
+ protected Class<Configuration> configurationClass() {
+ return Configuration.class;
+ }
+
+ @Override
+ protected SchemaTransform from(Configuration configuration) {
+ configuration.validate();
+ return new SnowflakeReadSchemaTransform(configuration);
+ }
+
+ @Override
+ public List<String> inputCollectionNames() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public List<String> outputCollectionNames() {
+ return Collections.singletonList(OUTPUT_TAG);
+ }
+
+ private static class SnowflakeReadSchemaTransform extends SchemaTransform
+ implements Serializable {
+
+ private final Configuration configuration;
+
+ private SnowflakeReadSchemaTransform(Configuration configuration) {
+ this.configuration = configuration;
+ }
+
+ @Override
+ public PCollectionRowTuple expand(PCollectionRowTuple input) {
+ Schema outputSchema =
JsonUtils.beamSchemaFromJsonSchema(configuration.getSchema());
+
+ SnowflakeIO.DataSourceConfiguration dataSourceConfiguration =
+ SnowflakeSchemaTransformUtils.createDataSourceConfiguration(
+ configuration.getServerName(),
+ configuration.getUsername(),
+ configuration.getPassword(),
+ configuration.getOauthToken(),
+ configuration.getPrivateKey(),
+ configuration.getPrivateKeyPassphrase(),
+ configuration.getDatabase(),
+ configuration.getSnowflakeSchema(),
+ configuration.getWarehouse(),
+ configuration.getRole());
+
+ SnowflakeIO.Read<Row> read =
+ SnowflakeIO.<Row>read()
+ .withDataSourceConfiguration(dataSourceConfiguration)
+ .withStagingBucketName(configuration.getStagingBucketName())
+
.withStorageIntegrationName(configuration.getStorageIntegrationName())
+ .withCsvMapper(parts -> toRow(parts, outputSchema))
+ .withCoder(RowCoder.of(outputSchema));
+
+ if (configuration.getTable() != null) {
+ read = read.fromTable(configuration.getTable());
+ } else {
+ read = read.fromQuery(configuration.getQuery());
+ }
+
+ if (configuration.getQuotationMark() != null) {
+ read = read.withQuotationMark(configuration.getQuotationMark());
+ }
+
+ PCollection<Row> rows =
+ input.getPipeline().apply("ReadFromSnowflake",
read).setRowSchema(outputSchema);
+
+ return PCollectionRowTuple.of(OUTPUT_TAG, rows);
+ }
+ }
+
+ @AutoValue
+ @DefaultSchema(AutoValueSchema.class)
+ public abstract static class Configuration implements Serializable {
+
+ @SchemaFieldDescription("Snowflake server name.")
+ public abstract String getServerName();
+
+ @SchemaFieldDescription(
+ "Snowflake username. Required for password and private key
authentication.")
+ @Nullable
+ public abstract String getUsername();
+
+ @SchemaFieldDescription(
+ "Snowflake password. Mutually exclusive with OAuth token and private
key.")
+ @Nullable
+ public abstract String getPassword();
+
+ @SchemaFieldDescription(
+ "Snowflake OAuth token. Mutually exclusive with password and private
key.")
+ @Nullable
+ public abstract String getOauthToken();
+
+ @SchemaFieldDescription(
+ "Raw Snowflake private key. Mutually exclusive with password and OAuth
token.")
+ @Nullable
+ public abstract String getPrivateKey();
+
+ @SchemaFieldDescription("Passphrase for the Snowflake private key.")
+ @Nullable
+ public abstract String getPrivateKeyPassphrase();
+
+ @SchemaFieldDescription("Snowflake database name.")
+ public abstract String getDatabase();
+
+ @SchemaFieldDescription("Snowflake schema name.")
+ public abstract String getSnowflakeSchema();
+
+ @SchemaFieldDescription("Snowflake warehouse name.")
+ @Nullable
+ public abstract String getWarehouse();
+
+ @SchemaFieldDescription("Snowflake role.")
+ @Nullable
+ public abstract String getRole();
+
+ @SchemaFieldDescription("Snowflake table to read from.")
+ @Nullable
+ public abstract String getTable();
+
+ @SchemaFieldDescription("Snowflake query to read from.")
+ @Nullable
+ public abstract String getQuery();
+
+ @SchemaFieldDescription("GCS path used to stage CSV files. The path must
end with '/'.")
+ public abstract String getStagingBucketName();
+
+ @SchemaFieldDescription("Snowflake storage integration name.")
+ public abstract String getStorageIntegrationName();
+
+ @SchemaFieldDescription("Output schema encoded using JSON Schema syntax.")
+ public abstract String getSchema();
+
+ @SchemaFieldDescription("Quotation mark used when parsing staged CSV
files.")
+ @Nullable
+ public abstract String getQuotationMark();
+
+ public static Builder builder() {
+ return new
AutoValue_SnowflakeReadSchemaTransformProvider_Configuration.Builder();
+ }
+
+ public abstract Builder toBuilder();
+
+ void validate() {
+ requireNonEmpty(getServerName(), "serverName");
+ requireNonEmpty(getDatabase(), "database");
+ requireNonEmpty(getSnowflakeSchema(), "snowflakeSchema");
+ requireNonEmpty(getStagingBucketName(), "stagingBucketName");
+ requireNonEmpty(getStorageIntegrationName(), "storageIntegrationName");
+ requireNonEmpty(getSchema(), "schema");
+
+ SnowflakeSchemaTransformUtils.validateAuthentication(
+ getUsername(),
+ getPassword(),
+ getOauthToken(),
+ getPrivateKey(),
+ getPrivateKeyPassphrase());
+
+ boolean tablePresent = getTable() != null && !getTable().isEmpty();
+ boolean queryPresent = getQuery() != null && !getQuery().isEmpty();
+
+ if (!tablePresent && !queryPresent) {
+ throw new IllegalArgumentException("Either table or query must be
specified.");
+ }
+
+ if (tablePresent && queryPresent) {
+ throw new IllegalArgumentException("table and query are mutually
exclusive.");
+ }
+
+ if (!getStagingBucketName().endsWith("/")) {
+ throw new IllegalArgumentException("stagingBucketName must end with
'/'");
+ }
+
+ // Validate JSON schema early.
+ JsonUtils.beamSchemaFromJsonSchema(getSchema());
+ }
+
+ private static void requireNonEmpty(String value, String name) {
+ if (value == null || value.isEmpty()) {
+ throw new IllegalArgumentException(name + " cannot be empty");
+ }
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ public abstract Builder setServerName(String value);
+
+ public abstract Builder setUsername(@Nullable String value);
+
+ public abstract Builder setPassword(@Nullable String value);
+
+ public abstract Builder setOauthToken(@Nullable String value);
+
+ public abstract Builder setPrivateKey(@Nullable String value);
+
+ public abstract Builder setPrivateKeyPassphrase(@Nullable String value);
+
+ public abstract Builder setDatabase(String value);
+
+ public abstract Builder setSnowflakeSchema(String value);
+
+ public abstract Builder setWarehouse(@Nullable String value);
+
+ public abstract Builder setRole(@Nullable String value);
+
+ public abstract Builder setTable(@Nullable String value);
+
+ public abstract Builder setQuery(@Nullable String value);
+
+ public abstract Builder setStagingBucketName(String value);
+
+ public abstract Builder setStorageIntegrationName(String value);
+
+ public abstract Builder setSchema(String value);
+
+ public abstract Builder setQuotationMark(@Nullable String value);
+
+ public abstract Configuration build();
+ }
+ }
+}
diff --git
a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeSchemaTransformUtils.java
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeSchemaTransformUtils.java
new file mode 100644
index 00000000000..84d451552c3
--- /dev/null
+++
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeSchemaTransformUtils.java
@@ -0,0 +1,323 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.snowflake;
+
+import java.nio.charset.StandardCharsets;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.io.snowflake.data.SnowflakeColumn;
+import org.apache.beam.sdk.io.snowflake.data.SnowflakeDataType;
+import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
+import org.apache.beam.sdk.io.snowflake.data.datetime.SnowflakeTimestamp;
+import org.apache.beam.sdk.io.snowflake.data.logical.SnowflakeBoolean;
+import org.apache.beam.sdk.io.snowflake.data.numeric.SnowflakeDouble;
+import org.apache.beam.sdk.io.snowflake.data.numeric.SnowflakeNumber;
+import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeBinary;
+import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeVarchar;
+import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
+import org.apache.beam.sdk.io.snowflake.enums.StreamingLogLevel;
+import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+import org.joda.time.Instant;
+
+/** Utilities shared by Snowflake schema transform providers. */
+@SuppressWarnings({
+ "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+public class SnowflakeSchemaTransformUtils {
+
+ public static SnowflakeIO.DataSourceConfiguration
createDataSourceConfiguration(
+ String serverName,
+ @Nullable String username,
+ @Nullable String password,
+ @Nullable String oauthToken,
+ @Nullable String privateKey,
+ @Nullable String privateKeyPassphrase,
+ String database,
+ String snowflakeSchema,
+ @Nullable String warehouse,
+ @Nullable String role) {
+
+ SnowflakeIO.DataSourceConfiguration configuration =
+ SnowflakeIO.DataSourceConfiguration.create();
+
+ if (isNotEmpty(password)) {
+ configuration = configuration.withUsernamePasswordAuth(username,
password);
+ } else if (isNotEmpty(oauthToken)) {
+ configuration = configuration.withOAuth(oauthToken);
+ } else if (isNotEmpty(privateKey)) {
+ if (isNotEmpty(privateKeyPassphrase)) {
+ configuration =
+ configuration.withKeyPairRawAuth(username, privateKey,
privateKeyPassphrase);
+ } else {
+ configuration = configuration.withKeyPairRawAuth(username, privateKey);
+ }
+ }
+
+ configuration =
+
configuration.withServerName(serverName).withDatabase(database).withSchema(snowflakeSchema);
+
+ if (isNotEmpty(warehouse)) {
+ configuration = configuration.withWarehouse(warehouse);
+ }
+
+ if (isNotEmpty(role)) {
+ configuration = configuration.withRole(role);
+ }
+
+ return configuration;
+ }
+
+ public static void validateAuthentication(
+ @Nullable String username,
+ @Nullable String password,
+ @Nullable String oauthToken,
+ @Nullable String privateKey,
+ @Nullable String privateKeyPassphrase) {
+
+ int authenticationMethods = 0;
+
+ if (isNotEmpty(password)) {
+ authenticationMethods++;
+ }
+
+ if (isNotEmpty(oauthToken)) {
+ authenticationMethods++;
+ }
+
+ if (isNotEmpty(privateKey)) {
+ authenticationMethods++;
+ }
+
+ if (authenticationMethods != 1) {
+ throw new IllegalArgumentException(
+ "Exactly one authentication method must be configured: "
+ + "password, oauthToken, or privateKey.");
+ }
+
+ if ((isNotEmpty(password) || isNotEmpty(privateKey)) &&
!isNotEmpty(username)) {
+ throw new IllegalArgumentException(
+ "username is required for password and private key authentication.");
+ }
+
+ if (isNotEmpty(privateKeyPassphrase) && !isNotEmpty(privateKey)) {
+ throw new IllegalArgumentException("privateKeyPassphrase requires
privateKey.");
+ }
+ }
+
+ public static boolean isNotEmpty(@Nullable String value) {
+ return value != null && !value.isEmpty();
+ }
+
+ public static StreamingLogLevel parseStreamingLogLevel(String value) {
+ try {
+ return StreamingLogLevel.valueOf(value);
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ "Unsupported debugMode '" + value + "'. Supported values are ERROR
and INFO.", e);
+ }
+ }
+
+ public static CreateDisposition parseCreateDisposition(String value) {
+ try {
+ return CreateDisposition.valueOf(value);
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ "Unsupported createDisposition '"
+ + value
+ + "'. Supported values are CREATE_IF_NEEDED and CREATE_NEVER.",
+ e);
+ }
+ }
+
+ public static WriteDisposition parseWriteDisposition(String value) {
+ try {
+ return WriteDisposition.valueOf(value);
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ "Unsupported writeDisposition '"
+ + value
+ + "'. Supported values are APPEND, TRUNCATE, and EMPTY.",
+ e);
+ }
+ }
+
+ public static SnowflakeTableSchema toSnowflakeTableSchema(Schema schema) {
+ SnowflakeColumn[] columns =
+ schema.getFields().stream()
+ .map(SnowflakeSchemaTransformUtils::toSnowflakeColumn)
+ .toArray(SnowflakeColumn[]::new);
+
+ return SnowflakeTableSchema.of(columns);
+ }
+
+ public static SnowflakeColumn toSnowflakeColumn(Schema.Field field) {
+ SnowflakeDataType snowflakeType = toSnowflakeDataType(field);
+
+ return SnowflakeColumn.of(field.getName(), snowflakeType,
field.getType().getNullable());
+ }
+
+ public static Row toRow(String[] parts, Schema schema) {
+ if (parts.length != schema.getFieldCount()) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Snowflake row contains %d values, but the configured schema
contains %d fields.",
+ parts.length, schema.getFieldCount()));
+ }
+
+ Row.Builder builder = Row.withSchema(schema);
+
+ for (int i = 0; i < schema.getFieldCount(); i++) {
+ Schema.Field field = schema.getField(i);
+ builder.addValue(toBeamValue(parts[i], field));
+ }
+
+ return builder.build();
+ }
+
+ public static Object toBeamValue(String value, Schema.Field field) {
+ if (value == null || value.isEmpty()) {
+ if (field.getType().getNullable()) {
+ return null;
+ }
+
+ /*
+ * Snowflake COPY encodes NULL as an empty CSV value. Therefore an empty
+ * value cannot be represented for a required non-string type.
+ *
+ * For STRING, preserve the empty string.
+ */
+ if (field.getType().getTypeName() == Schema.TypeName.STRING) {
+ return "";
+ }
+
+ throw new IllegalArgumentException(
+ String.format(
+ "Received an empty value for non-nullable Snowflake field
'%s'.", field.getName()));
+ }
+
+ try {
+ switch (field.getType().getTypeName()) {
+ case BYTE:
+ return Byte.valueOf(value);
+
+ case INT16:
+ return Short.valueOf(value);
+
+ case INT32:
+ return Integer.valueOf(value);
+
+ case INT64:
+ return Long.valueOf(value);
+
+ case FLOAT:
+ return Float.valueOf(value);
+
+ case DOUBLE:
+ return Double.valueOf(value);
+
+ case STRING:
+ return value;
+
+ case BOOLEAN:
+ if ("true".equalsIgnoreCase(value)) {
+ return true;
+ }
+
+ if ("false".equalsIgnoreCase(value)) {
+ return false;
+ }
+
+ throw new IllegalArgumentException(String.format("Invalid boolean
value '%s'.", value));
+
+ case BYTES:
+ return value.getBytes(StandardCharsets.UTF_8);
+
+ case DATETIME:
+ return Instant.parse(value);
+
+ case DECIMAL:
+ case ARRAY:
+ case ITERABLE:
+ case MAP:
+ case ROW:
+ case LOGICAL_TYPE:
+ default:
+ throw unsupportedFieldType(field, null);
+ }
+ } catch (IllegalArgumentException e) {
+ throw new IllegalArgumentException(
+ String.format(
+ "Unable to parse value '%s' as %s for Snowflake field '%s'.",
+ value, field.getType().getTypeName(), field.getName()),
+ e);
+ }
+ }
+
+ public static SnowflakeDataType toSnowflakeDataType(Schema.Field field) {
+ switch (field.getType().getTypeName()) {
+ case BYTE:
+ case INT16:
+ case INT32:
+ case INT64:
+ return SnowflakeNumber.of();
+
+ case FLOAT:
+ case DOUBLE:
+ return SnowflakeDouble.of();
+
+ case STRING:
+ return SnowflakeVarchar.of();
+
+ case BOOLEAN:
+ return SnowflakeBoolean.of();
+
+ case BYTES:
+ return SnowflakeBinary.of();
+
+ case DATETIME:
+ return SnowflakeTimestamp.of();
+
+ case DECIMAL:
+ throw unsupportedFieldType(
+ field, "Beam DECIMAL does not include Snowflake precision and
scale information.");
+
+ case ARRAY:
+ case ITERABLE:
+ case MAP:
+ case ROW:
+ case LOGICAL_TYPE:
+ default:
+ throw unsupportedFieldType(field, null);
+ }
+ }
+
+ public static IllegalArgumentException unsupportedFieldType(
+ Schema.Field field, @Nullable String details) {
+ String message =
+ String.format(
+ "Unsupported Beam field type %s for Snowflake column '%s'.",
+ field.getType().getTypeName(), field.getName());
+
+ if (details != null) {
+ message += " " + details;
+ }
+
+ return new IllegalArgumentException(message);
+ }
+}
diff --git
a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteConfiguration.java
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteConfiguration.java
new file mode 100644
index 00000000000..43fe52b1245
--- /dev/null
+++
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteConfiguration.java
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.snowflake;
+
+import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.*;
+
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
+
+@AutoValue
+@DefaultSchema(AutoValueSchema.class)
+public abstract class SnowflakeWriteConfiguration implements Serializable {
+
+ @SchemaFieldDescription("Snowflake server name.")
+ public abstract String getServerName();
+
+ @SchemaFieldDescription(
+ "Snowflake username. Required for password and private key
authentication.")
+ @Nullable
+ public abstract String getUsername();
+
+ @SchemaFieldDescription(
+ "Snowflake password. Mutually exclusive with OAuth token and private
key.")
+ @Nullable
+ public abstract String getPassword();
+
+ @SchemaFieldDescription(
+ "Snowflake OAuth token. Mutually exclusive with password and private
key.")
+ @Nullable
+ public abstract String getOauthToken();
+
+ @SchemaFieldDescription(
+ "Raw Snowflake private key. Mutually exclusive with password and OAuth
token.")
+ @Nullable
+ public abstract String getPrivateKey();
+
+ @SchemaFieldDescription("Passphrase for the Snowflake private key.")
+ @Nullable
+ public abstract String getPrivateKeyPassphrase();
+
+ @SchemaFieldDescription("Snowflake database name.")
+ public abstract String getDatabase();
+
+ @SchemaFieldDescription("Snowflake schema name.")
+ public abstract String getSchema();
+
+ @SchemaFieldDescription("Snowflake warehouse name.")
+ @Nullable
+ public abstract String getWarehouse();
+
+ @SchemaFieldDescription("Snowflake role.")
+ @Nullable
+ public abstract String getRole();
+
+ @SchemaFieldDescription("Destination Snowflake table. Required for batch
writes.")
+ @Nullable
+ public abstract String getTable();
+
+ @SchemaFieldDescription("Snowflake Snowpipe name. Required for streaming
writes.")
+ @Nullable
+ public abstract String getSnowPipe();
+
+ @SchemaFieldDescription("GCS path used to stage CSV files. The path must end
with '/'.")
+ public abstract String getStagingBucketName();
+
+ @SchemaFieldDescription("Snowflake storage integration name.")
+ public abstract String getStorageIntegrationName();
+
+ @SchemaFieldDescription(
+ "Table creation behavior for batch writes. "
+ + "Supported values are CREATE_IF_NEEDED and CREATE_NEVER.")
+ @Nullable
+ public abstract String getCreateDisposition();
+
+ @SchemaFieldDescription(
+ "Write behavior for batch writes. " + "Supported values are APPEND,
TRUNCATE, and EMPTY.")
+ @Nullable
+ public abstract String getWriteDisposition();
+
+ @SchemaFieldDescription("Quotation mark used when writing values to staged
CSV files.")
+ @Nullable
+ public abstract String getQuotationMark();
+
+ @SchemaFieldDescription("Maximum number of rows to stage before flushing in
streaming mode.")
+ @Nullable
+ public abstract Integer getFlushRowLimit();
+
+ @SchemaFieldDescription(
+ "Maximum time in milliseconds before flushing staged rows in streaming
mode.")
+ @Nullable
+ public abstract Long getFlushTimeLimitMillis();
+
+ @SchemaFieldDescription("Number of output shards used when staging files.")
+ @Nullable
+ public abstract Integer getShardsNumber();
+
+ @SchemaFieldDescription("Streaming log level. Supported values are ERROR and
INFO.")
+ @Nullable
+ public abstract String getDebugMode();
+
+ public static Builder builder() {
+ return new AutoValue_SnowflakeWriteConfiguration.Builder();
+ }
+
+ public abstract Builder toBuilder();
+
+ void validate() {
+ requireNonEmpty(getServerName(), "serverName");
+ requireNonEmpty(getDatabase(), "database");
+ requireNonEmpty(getSchema(), "schema");
+ requireNonEmpty(getStagingBucketName(), "stagingBucketName");
+ requireNonEmpty(getStorageIntegrationName(), "storageIntegrationName");
+
+ SnowflakeSchemaTransformUtils.validateAuthentication(
+ getUsername(), getPassword(), getOauthToken(), getPrivateKey(),
getPrivateKeyPassphrase());
+
+ if (!getStagingBucketName().endsWith("/")) {
+ throw new IllegalArgumentException("stagingBucketName must end with
'/'");
+ }
+
+ String createDisposition = getCreateDisposition();
+ if (createDisposition != null) {
+ parseCreateDisposition(createDisposition);
+ }
+
+ String writeDisposition = getWriteDisposition();
+ if (writeDisposition != null) {
+ parseWriteDisposition(writeDisposition);
+ }
+
+ String debugMode = getDebugMode();
+ if (debugMode != null) {
+ parseStreamingLogLevel(debugMode);
+ }
+
+ Integer flushRowLimit = getFlushRowLimit();
+ if (flushRowLimit != null && flushRowLimit <= 0) {
+ throw new IllegalArgumentException("flushRowLimit must be greater than
0.");
+ }
+
+ Long flushTimeLimitMillis = getFlushTimeLimitMillis();
+ if (flushTimeLimitMillis != null && flushTimeLimitMillis <= 0) {
+ throw new IllegalArgumentException("flushTimeLimitMillis must be greater
than 0.");
+ }
+
+ Integer shardsNumber = getShardsNumber();
+ if (shardsNumber != null && shardsNumber <= 0) {
+ throw new IllegalArgumentException("shardsNumber must be greater than
0.");
+ }
+ }
+
+ private static void requireNonEmpty(String value, String name) {
+ if (value == null || value.isEmpty()) {
+ throw new IllegalArgumentException(name + " cannot be empty");
+ }
+ }
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ public abstract Builder setServerName(String value);
+
+ public abstract Builder setUsername(@Nullable String value);
+
+ public abstract Builder setPassword(@Nullable String value);
+
+ public abstract Builder setOauthToken(@Nullable String value);
+
+ public abstract Builder setPrivateKey(@Nullable String value);
+
+ public abstract Builder setPrivateKeyPassphrase(@Nullable String value);
+
+ public abstract Builder setDatabase(String value);
+
+ public abstract Builder setSchema(String value);
+
+ public abstract Builder setWarehouse(@Nullable String value);
+
+ public abstract Builder setRole(@Nullable String value);
+
+ public abstract Builder setTable(@Nullable String value);
+
+ public abstract Builder setSnowPipe(@Nullable String value);
+
+ public abstract Builder setStagingBucketName(String value);
+
+ public abstract Builder setStorageIntegrationName(String value);
+
+ public abstract Builder setCreateDisposition(@Nullable String value);
+
+ public abstract Builder setWriteDisposition(@Nullable String value);
+
+ public abstract Builder setQuotationMark(@Nullable String value);
+
+ public abstract Builder setFlushRowLimit(@Nullable Integer value);
+
+ public abstract Builder setFlushTimeLimitMillis(@Nullable Long value);
+
+ public abstract Builder setShardsNumber(@Nullable Integer value);
+
+ public abstract Builder setDebugMode(@Nullable String value);
+
+ public abstract SnowflakeWriteConfiguration build();
+ }
+}
diff --git
a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java
index 623bb134f58..3010d88b010 100644
---
a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java
+++
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java
@@ -17,33 +17,20 @@
*/
package org.apache.beam.sdk.io.snowflake;
+import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.*;
+
import com.google.auto.service.AutoService;
-import com.google.auto.value.AutoValue;
import java.io.Serializable;
import java.util.Collections;
import java.util.List;
-import javax.annotation.Nullable;
-import org.apache.beam.sdk.io.snowflake.data.SnowflakeColumn;
-import org.apache.beam.sdk.io.snowflake.data.SnowflakeDataType;
-import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
-import org.apache.beam.sdk.io.snowflake.data.datetime.SnowflakeTimestamp;
-import org.apache.beam.sdk.io.snowflake.data.logical.SnowflakeBoolean;
-import org.apache.beam.sdk.io.snowflake.data.numeric.SnowflakeDouble;
-import org.apache.beam.sdk.io.snowflake.data.numeric.SnowflakeNumber;
-import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeBinary;
-import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeVarchar;
import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition;
-import org.apache.beam.sdk.io.snowflake.enums.WriteDisposition;
-import org.apache.beam.sdk.schemas.AutoValueSchema;
-import org.apache.beam.sdk.schemas.Schema;
-import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
-import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription;
import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider;
import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionRowTuple;
import org.apache.beam.sdk.values.Row;
+import org.joda.time.Duration;
/** A {@link SchemaTransformProvider} for writing Beam rows to Snowflake. */
@SuppressWarnings({
@@ -51,7 +38,7 @@ import org.apache.beam.sdk.values.Row;
})
@AutoService(SchemaTransformProvider.class)
public class SnowflakeWriteSchemaTransformProvider
- extends
TypedSchemaTransformProvider<SnowflakeWriteSchemaTransformProvider.Configuration>
{
+ extends TypedSchemaTransformProvider<SnowflakeWriteConfiguration> {
static final String INPUT_TAG = "input";
@@ -64,16 +51,16 @@ public class SnowflakeWriteSchemaTransformProvider
@Override
public String description() {
- return "Writes Beam Rows to a Snowflake table using staged CSV files.";
+ return "Writes Beam Rows to Snowflake using batch COPY or streaming
Snowpipe.";
}
@Override
- protected Class<Configuration> configurationClass() {
- return Configuration.class;
+ protected Class<SnowflakeWriteConfiguration> configurationClass() {
+ return SnowflakeWriteConfiguration.class;
}
@Override
- protected SchemaTransform from(Configuration configuration) {
+ protected SchemaTransform from(SnowflakeWriteConfiguration configuration) {
configuration.validate();
return new SnowflakeWriteSchemaTransform(configuration);
}
@@ -88,13 +75,12 @@ public class SnowflakeWriteSchemaTransformProvider
return Collections.emptyList();
}
- /** Schema transform that configures and applies {@link SnowflakeIO.Write}.
*/
private static class SnowflakeWriteSchemaTransform extends SchemaTransform
implements Serializable {
- private final Configuration configuration;
+ private final SnowflakeWriteConfiguration configuration;
- private SnowflakeWriteSchemaTransform(Configuration configuration) {
+ private SnowflakeWriteSchemaTransform(SnowflakeWriteConfiguration
configuration) {
this.configuration = configuration;
}
@@ -103,260 +89,91 @@ public class SnowflakeWriteSchemaTransformProvider
PCollection<Row> rows = input.get(INPUT_TAG);
SnowflakeIO.DataSourceConfiguration dataSourceConfiguration =
- SnowflakeIO.DataSourceConfiguration.create()
- .withUsernamePasswordAuth(configuration.getUsername(),
configuration.getPassword())
- .withServerName(configuration.getServerName())
- .withDatabase(configuration.getDatabase())
- .withSchema(configuration.getSchema());
-
- if (configuration.getWarehouse() != null) {
- dataSourceConfiguration =
-
dataSourceConfiguration.withWarehouse(configuration.getWarehouse());
- }
-
- if (configuration.getRole() != null) {
- dataSourceConfiguration =
dataSourceConfiguration.withRole(configuration.getRole());
- }
+ SnowflakeSchemaTransformUtils.createDataSourceConfiguration(
+ configuration.getServerName(),
+ configuration.getUsername(),
+ configuration.getPassword(),
+ configuration.getOauthToken(),
+ configuration.getPrivateKey(),
+ configuration.getPrivateKeyPassphrase(),
+ configuration.getDatabase(),
+ configuration.getSchema(),
+ configuration.getWarehouse(),
+ configuration.getRole());
SnowflakeIO.Write<Row> write =
SnowflakeIO.<Row>write()
.withDataSourceConfiguration(dataSourceConfiguration)
.withStagingBucketName(configuration.getStagingBucketName())
.withStorageIntegrationName(configuration.getStorageIntegrationName())
- .withUserDataMapper(row -> row.getValues().toArray())
- .to(configuration.getTable());
+ .withUserDataMapper(row -> row.getValues().toArray());
- if (configuration.getCreateDisposition() != null) {
- CreateDisposition createDisposition =
- parseCreateDisposition(configuration.getCreateDisposition());
+ boolean streaming = rows.isBounded() == PCollection.IsBounded.UNBOUNDED;
- write = write.withCreateDisposition(createDisposition);
+ if (streaming) {
+ String snowPipe = configuration.getSnowPipe();
- if (createDisposition == CreateDisposition.CREATE_IF_NEEDED) {
- write =
write.withTableSchema(toSnowflakeTableSchema(rows.getSchema()));
+ if (snowPipe == null || snowPipe.isEmpty()) {
+ throw new IllegalArgumentException("snowPipe is required for
streaming writes.");
}
- }
-
- if (configuration.getWriteDisposition() != null) {
- write =
-
write.withWriteDisposition(parseWriteDisposition(configuration.getWriteDisposition()));
- }
-
- if (configuration.getQuotationMark() != null) {
- write = write.withQuotationMark(configuration.getQuotationMark());
- }
-
- rows.apply("WriteToSnowflake", write);
-
- return PCollectionRowTuple.empty(input.getPipeline());
- }
- }
-
- @AutoValue
- @DefaultSchema(AutoValueSchema.class)
- public abstract static class Configuration implements Serializable {
-
- @SchemaFieldDescription("Snowflake server name.")
- public abstract String getServerName();
-
- @SchemaFieldDescription("Snowflake username.")
- public abstract String getUsername();
- @SchemaFieldDescription("Snowflake password.")
- public abstract String getPassword();
+ write = write.withSnowPipe(snowPipe);
- @SchemaFieldDescription("Snowflake database name.")
- public abstract String getDatabase();
-
- @SchemaFieldDescription("Snowflake schema name.")
- public abstract String getSchema();
-
- @SchemaFieldDescription("Snowflake warehouse name.")
- @Nullable
- public abstract String getWarehouse();
-
- @SchemaFieldDescription("Snowflake role.")
- @Nullable
- public abstract String getRole();
+ Integer flushRowLimit = configuration.getFlushRowLimit();
+ if (flushRowLimit != null) {
+ write = write.withFlushRowLimit(flushRowLimit);
+ }
- @SchemaFieldDescription("Destination Snowflake table.")
- public abstract String getTable();
+ Long flushTimeLimitMillis = configuration.getFlushTimeLimitMillis();
+ if (flushTimeLimitMillis != null) {
+ write =
write.withFlushTimeLimit(Duration.millis(flushTimeLimitMillis));
+ }
- @SchemaFieldDescription("GCS path used to stage CSV files. The path must
end with '/'.")
- public abstract String getStagingBucketName();
+ Integer shardsNumber = configuration.getShardsNumber();
+ if (shardsNumber != null) {
+ write = write.withShardsNumber(shardsNumber);
+ }
- @SchemaFieldDescription("Snowflake storage integration name.")
- public abstract String getStorageIntegrationName();
+ String debugMode = configuration.getDebugMode();
+ if (debugMode != null) {
+ write = write.withDebugMode(parseStreamingLogLevel(debugMode));
+ }
+ } else {
+ String table = configuration.getTable();
- @SchemaFieldDescription(
- "Table creation behavior. Supported values are CREATE_IF_NEEDED and
CREATE_NEVER.")
- @Nullable
- public abstract String getCreateDisposition();
+ if (table == null || table.isEmpty()) {
+ throw new IllegalArgumentException("table is required for batch
writes.");
+ }
- @SchemaFieldDescription("Write behavior. Supported values are APPEND,
TRUNCATE, and EMPTY.")
- @Nullable
- public abstract String getWriteDisposition();
+ write = write.to(table);
- @SchemaFieldDescription("Quotation mark used when writing values to staged
CSV files.")
- @Nullable
- public abstract String getQuotationMark();
+ String createDispositionValue = configuration.getCreateDisposition();
- public static Builder builder() {
- return new
AutoValue_SnowflakeWriteSchemaTransformProvider_Configuration.Builder();
- }
+ if (createDispositionValue != null) {
+ CreateDisposition createDisposition =
parseCreateDisposition(createDispositionValue);
- public abstract Builder toBuilder();
+ write = write.withCreateDisposition(createDisposition);
- void validate() {
- requireNonEmpty(getServerName(), "serverName");
- requireNonEmpty(getUsername(), "username");
- requireNonEmpty(getPassword(), "password");
- requireNonEmpty(getDatabase(), "database");
- requireNonEmpty(getSchema(), "schema");
- requireNonEmpty(getTable(), "table");
- requireNonEmpty(getStagingBucketName(), "stagingBucketName");
- requireNonEmpty(getStorageIntegrationName(), "storageIntegrationName");
-
- if (!getStagingBucketName().endsWith("/")) {
- throw new IllegalArgumentException("stagingBucketName must end with
'/'");
- }
+ if (createDisposition == CreateDisposition.CREATE_IF_NEEDED) {
+ write =
write.withTableSchema(toSnowflakeTableSchema(rows.getSchema()));
+ }
+ }
- if (getCreateDisposition() != null) {
- parseCreateDisposition(getCreateDisposition());
- }
+ String writeDispositionValue = configuration.getWriteDisposition();
- if (getWriteDisposition() != null) {
- parseWriteDisposition(getWriteDisposition());
+ if (writeDispositionValue != null) {
+ write =
write.withWriteDisposition(parseWriteDisposition(writeDispositionValue));
+ }
}
- }
- private static void requireNonEmpty(String value, String name) {
- if (value == null || value.isEmpty()) {
- throw new IllegalArgumentException(name + " cannot be empty");
+ String quotationMark = configuration.getQuotationMark();
+ if (quotationMark != null) {
+ write = write.withQuotationMark(quotationMark);
}
- }
-
- @AutoValue.Builder
- public abstract static class Builder {
- public abstract Builder setServerName(String value);
-
- public abstract Builder setUsername(String value);
-
- public abstract Builder setPassword(String value);
-
- public abstract Builder setDatabase(String value);
-
- public abstract Builder setSchema(String value);
-
- public abstract Builder setWarehouse(String value);
-
- public abstract Builder setRole(String value);
-
- public abstract Builder setTable(String value);
-
- public abstract Builder setStagingBucketName(String value);
-
- public abstract Builder setStorageIntegrationName(String value);
-
- public abstract Builder setCreateDisposition(String value);
-
- public abstract Builder setWriteDisposition(String value);
-
- public abstract Builder setQuotationMark(String value);
-
- public abstract Configuration build();
- }
- }
-
- private static CreateDisposition parseCreateDisposition(String value) {
- try {
- return CreateDisposition.valueOf(value);
- } catch (IllegalArgumentException e) {
- throw new IllegalArgumentException(
- "Unsupported createDisposition '"
- + value
- + "'. Supported values are CREATE_IF_NEEDED and CREATE_NEVER.",
- e);
- }
- }
-
- private static WriteDisposition parseWriteDisposition(String value) {
- try {
- return WriteDisposition.valueOf(value);
- } catch (IllegalArgumentException e) {
- throw new IllegalArgumentException(
- "Unsupported writeDisposition '"
- + value
- + "'. Supported values are APPEND, TRUNCATE, and EMPTY.",
- e);
- }
- }
-
- static SnowflakeTableSchema toSnowflakeTableSchema(Schema schema) {
- SnowflakeColumn[] columns =
- schema.getFields().stream()
- .map(SnowflakeWriteSchemaTransformProvider::toSnowflakeColumn)
- .toArray(SnowflakeColumn[]::new);
-
- return SnowflakeTableSchema.of(columns);
- }
-
- private static SnowflakeColumn toSnowflakeColumn(Schema.Field field) {
- SnowflakeDataType snowflakeType = toSnowflakeDataType(field);
-
- return SnowflakeColumn.of(field.getName(), snowflakeType,
field.getType().getNullable());
- }
-
- private static SnowflakeDataType toSnowflakeDataType(Schema.Field field) {
- switch (field.getType().getTypeName()) {
- case BYTE:
- case INT16:
- case INT32:
- case INT64:
- return SnowflakeNumber.of();
-
- case FLOAT:
- case DOUBLE:
- return SnowflakeDouble.of();
-
- case STRING:
- return SnowflakeVarchar.of();
-
- case BOOLEAN:
- return SnowflakeBoolean.of();
-
- case BYTES:
- return SnowflakeBinary.of();
-
- case DATETIME:
- return SnowflakeTimestamp.of();
-
- case DECIMAL:
- throw unsupportedFieldType(
- field, "Beam DECIMAL does not include Snowflake precision and
scale information.");
-
- case ARRAY:
- case ITERABLE:
- case MAP:
- case ROW:
- case LOGICAL_TYPE:
- default:
- throw unsupportedFieldType(field, null);
- }
- }
-
- private static IllegalArgumentException unsupportedFieldType(
- Schema.Field field, @Nullable String details) {
- String message =
- String.format(
- "Unsupported Beam field type %s for Snowflake column '%s'.",
- field.getType().getTypeName(), field.getName());
+ rows.apply("WriteToSnowflake", write);
- if (details != null) {
- message += " " + details;
+ return PCollectionRowTuple.empty(input.getPipeline());
}
-
- return new IllegalArgumentException(message);
}
}
diff --git
a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProviderTest.java
b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProviderTest.java
new file mode 100644
index 00000000000..02f8e2e7b66
--- /dev/null
+++
b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProviderTest.java
@@ -0,0 +1,311 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.snowflake;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertThrows;
+
+import java.nio.charset.StandardCharsets;
+import
org.apache.beam.sdk.io.snowflake.SnowflakeReadSchemaTransformProvider.Configuration;
+import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.values.Row;
+import org.joda.time.Instant;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class SnowflakeReadSchemaTransformProviderTest {
+
+ private static final String SIMPLE_SCHEMA =
+ "{"
+ + "\"type\":\"object\","
+ + "\"properties\":{"
+ + "\"id\":{\"type\":\"integer\"},"
+ + "\"name\":{\"type\":\"string\"}"
+ + "},"
+ + "\"required\":[\"id\",\"name\"]"
+ + "}";
+
+ private final SnowflakeReadSchemaTransformProvider provider =
+ new SnowflakeReadSchemaTransformProvider();
+
+ @Test
+ public void testIdentifier() {
+ assertThat(
+ provider.identifier(),
equalTo("beam:schematransform:org.apache.beam:snowflake_read:v1"));
+ }
+
+ @Test
+ public void testInputCollectionNames() {
+ assertThat(provider.inputCollectionNames(), empty());
+ }
+
+ @Test
+ public void testOutputCollectionNames() {
+ assertThat(provider.outputCollectionNames(), contains("output"));
+ }
+
+ @Test
+ public void testValidTableConfiguration() {
+ provider.from(validConfiguration().setTable("table").build());
+ }
+
+ @Test
+ public void testValidQueryConfiguration() {
+ provider.from(validConfiguration().setQuery("SELECT * FROM
table").build());
+ }
+
+ @Test
+ public void testTableAndQueryAreMutuallyExclusive() {
+ Configuration configuration =
+ validConfiguration().setTable("table").setQuery("SELECT * FROM
table").build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("table and query are mutually
exclusive."));
+ }
+
+ @Test
+ public void testTableOrQueryIsRequired() {
+ Configuration configuration = validConfiguration().build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("Either table or query must be
specified."));
+ }
+
+ @Test
+ public void testStagingBucketMustEndWithSlash() {
+ Configuration configuration =
+
validConfiguration().setTable("table").setStagingBucketName("gs://bucket/staging").build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("stagingBucketName must end
with '/'"));
+ }
+
+ @Test
+ public void testConvertsCsvValuesToBeamRow() {
+ Schema schema =
+ Schema.builder()
+ .addByteField("byte_value")
+ .addInt16Field("short_value")
+ .addInt32Field("int_value")
+ .addInt64Field("long_value")
+ .addFloatField("float_value")
+ .addDoubleField("double_value")
+ .addStringField("string_value")
+ .addBooleanField("boolean_value")
+ .addByteArrayField("bytes_value")
+ .addDateTimeField("datetime_value")
+ .build();
+
+ String[] values = {
+ "1", "2", "3", "4", "5.5", "6.5", "hello", "true", "abc",
"2026-08-13T09:00:00.000Z"
+ };
+
+ Row row = SnowflakeSchemaTransformUtils.toRow(values, schema);
+
+ assertThat(row.getByte("byte_value"), equalTo((byte) 1));
+ assertThat(row.getInt16("short_value"), equalTo((short) 2));
+ assertThat(row.getInt32("int_value"), equalTo(3));
+ assertThat(row.getInt64("long_value"), equalTo(4L));
+ assertThat(row.getFloat("float_value"), equalTo(5.5F));
+ assertThat(row.getDouble("double_value"), equalTo(6.5D));
+ assertThat(row.getString("string_value"), equalTo("hello"));
+ assertThat(row.getBoolean("boolean_value"), equalTo(true));
+ assertArrayEquals("abc".getBytes(StandardCharsets.UTF_8),
row.getBytes("bytes_value"));
+ assertThat(
+ row.getDateTime("datetime_value"),
equalTo(Instant.parse("2026-08-13T09:00:00.000Z")));
+ }
+
+ @Test
+ public void testNullableEmptyValueBecomesNull() {
+ Schema schema = Schema.builder().addNullableField("value",
Schema.FieldType.INT64).build();
+
+ Row row = SnowflakeSchemaTransformUtils.toRow(new String[] {""}, schema);
+
+ assertThat(row.getValue("value"), equalTo(null));
+ }
+
+ @Test
+ public void testEmptyRequiredStringIsPreserved() {
+ Schema schema = Schema.builder().addStringField("value").build();
+
+ Row row = SnowflakeSchemaTransformUtils.toRow(new String[] {""}, schema);
+
+ assertThat(row.getString("value"), equalTo(""));
+ }
+
+ @Test
+ public void testEmptyRequiredNonStringIsRejected() {
+ Schema schema = Schema.builder().addInt64Field("value").build();
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> SnowflakeSchemaTransformUtils.toRow(new String[] {""},
schema));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("Received an empty value for non-nullable Snowflake field
'value'."));
+ }
+
+ @Test
+ public void testWrongNumberOfFieldsIsRejected() {
+ Schema schema =
Schema.builder().addInt64Field("id").addStringField("name").build();
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> SnowflakeSchemaTransformUtils.toRow(new String[] {"1"},
schema));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("Snowflake row contains 1 values, but the configured schema
contains 2 fields."));
+ }
+
+ @Test
+ public void testArrayIsRejected() {
+ Schema schema = Schema.builder().addArrayField("values",
Schema.FieldType.STRING).build();
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> SnowflakeSchemaTransformUtils.toRow(new String[] {"value"},
schema));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("Unable to parse value 'value' as ARRAY for Snowflake field
'values'."));
+ }
+
+ @Test
+ public void testOauthAuthenticationIsValid() {
+ provider.from(
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setOauthToken("token")
+ .setTable("table")
+ .build());
+ }
+
+ @Test
+ public void testPrivateKeyAuthenticationIsValid() {
+ provider.from(
+ validConfiguration()
+ .setPassword(null)
+ .setPrivateKey("private-key")
+ .setPrivateKeyPassphrase("passphrase")
+ .setTable("table")
+ .build());
+ }
+
+ @Test
+ public void testAuthenticationMethodIsRequired() {
+ Configuration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setOauthToken(null)
+ .setPrivateKey(null)
+ .setTable("table")
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo(
+ "Exactly one authentication method must be configured: "
+ + "password, oauthToken, or privateKey."));
+ }
+
+ @Test
+ public void testMultipleAuthenticationMethodsAreRejected() {
+ Configuration configuration =
+ validConfiguration().setOauthToken("token").setTable("table").build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo(
+ "Exactly one authentication method must be configured: "
+ + "password, oauthToken, or privateKey."));
+ }
+
+ @Test
+ public void testUsernameIsRequiredForPrivateKeyAuthentication() {
+ Configuration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setPrivateKey("private-key")
+ .setTable("table")
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("username is required for password and private key
authentication."));
+ }
+
+ @Test
+ public void testPrivateKeyPassphraseRequiresPrivateKey() {
+ Configuration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setOauthToken("token")
+ .setPrivateKeyPassphrase("passphrase")
+ .setTable("table")
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("privateKeyPassphrase requires
privateKey."));
+ }
+
+ private static Configuration.Builder validConfiguration() {
+ return Configuration.builder()
+ .setServerName("account.snowflakecomputing.com")
+ .setUsername("username")
+ .setPassword("password")
+ .setDatabase("database")
+ .setSnowflakeSchema("public")
+ .setWarehouse("warehouse")
+ .setRole("role")
+ .setStagingBucketName("gs://bucket/staging/")
+ .setStorageIntegrationName("storage_integration")
+ .setSchema(SIMPLE_SCHEMA);
+ }
+}
diff --git
a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java
b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java
index 336d7f01f03..bdd5aae0e75 100644
---
a/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java
+++
b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java
@@ -24,7 +24,8 @@ import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.assertThrows;
-import
org.apache.beam.sdk.io.snowflake.SnowflakeWriteSchemaTransformProvider.Configuration;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.RowCoder;
import org.apache.beam.sdk.io.snowflake.data.SnowflakeColumn;
import org.apache.beam.sdk.io.snowflake.data.SnowflakeTableSchema;
import org.apache.beam.sdk.io.snowflake.data.datetime.SnowflakeTimestamp;
@@ -34,6 +35,12 @@ import
org.apache.beam.sdk.io.snowflake.data.numeric.SnowflakeNumber;
import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeBinary;
import org.apache.beam.sdk.io.snowflake.data.text.SnowflakeVarchar;
import org.apache.beam.sdk.schemas.Schema;
+import org.apache.beam.sdk.schemas.transforms.SchemaTransform;
+import org.apache.beam.sdk.testing.TestStream;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.PCollectionRowTuple;
+import org.apache.beam.sdk.values.Row;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -41,6 +48,8 @@ import org.junit.runners.JUnit4;
@RunWith(JUnit4.class)
public class SnowflakeWriteSchemaTransformProviderTest {
+ private final transient Pipeline pipeline = Pipeline.create();
+
private final SnowflakeWriteSchemaTransformProvider provider =
new SnowflakeWriteSchemaTransformProvider();
@@ -62,31 +71,17 @@ public class SnowflakeWriteSchemaTransformProviderTest {
@Test
public void testValidConfiguration() {
- Configuration configuration = validConfiguration().build();
-
- provider.from(configuration);
+ provider.from(validConfiguration().build());
}
@Test
public void testBlankQuotationMarkIsAllowed() {
- Configuration configuration =
validConfiguration().setQuotationMark("").build();
-
- provider.from(configuration);
+ provider.from(validConfiguration().setQuotationMark("").build());
}
@Test
- public void testMissingServerName() {
- Configuration configuration =
- Configuration.builder()
- .setUsername("username")
- .setPassword("password")
- .setDatabase("database")
- .setSchema("schema")
- .setTable("table")
- .setServerName("")
- .setStagingBucketName("gs://bucket/staging/")
- .setStorageIntegrationName("storage_integration")
- .build();
+ public void testMissingServerNameIsRejected() {
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setServerName("").build();
IllegalArgumentException exception =
assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
@@ -95,28 +90,14 @@ public class SnowflakeWriteSchemaTransformProviderTest {
}
@Test
- public void testMissingTable() {
- Configuration configuration =
- Configuration.builder()
- .setServerName("account.snowflakecomputing.com")
- .setUsername("username")
- .setPassword("password")
- .setDatabase("database")
- .setSchema("schema")
- .setTable("")
- .setStagingBucketName("gs://bucket/staging/")
- .setStorageIntegrationName("storage_integration")
- .build();
-
- IllegalArgumentException exception =
- assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
-
- assertThat(exception.getMessage(), equalTo("table cannot be empty"));
+ public void testMissingTableIsAllowedAtConfigurationTime() {
+ // Whether table is required depends on whether the input is bounded.
+ provider.from(validConfiguration().setTable(null).build());
}
@Test
public void testStagingBucketMustEndWithSlash() {
- Configuration configuration =
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setStagingBucketName("gs://bucket/staging").build();
IllegalArgumentException exception =
@@ -126,8 +107,9 @@ public class SnowflakeWriteSchemaTransformProviderTest {
}
@Test
- public void testInvalidCreateDisposition() {
- Configuration configuration =
validConfiguration().setCreateDisposition("INVALID").build();
+ public void testInvalidCreateDispositionIsRejected() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration().setCreateDisposition("INVALID").build();
IllegalArgumentException exception =
assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
@@ -135,13 +117,24 @@ public class SnowflakeWriteSchemaTransformProviderTest {
assertThat(
exception.getMessage(),
equalTo(
- "Unsupported createDisposition 'INVALID'. Supported values are "
- + "CREATE_IF_NEEDED and CREATE_NEVER."));
+ "Unsupported createDisposition 'INVALID'. "
+ + "Supported values are CREATE_IF_NEEDED and CREATE_NEVER."));
}
@Test
- public void testInvalidWriteDisposition() {
- Configuration configuration =
validConfiguration().setWriteDisposition("INVALID").build();
+ public void testCreateIfNeededIsAccepted() {
+
provider.from(validConfiguration().setCreateDisposition("CREATE_IF_NEEDED").build());
+ }
+
+ @Test
+ public void testCreateNeverIsAccepted() {
+
provider.from(validConfiguration().setCreateDisposition("CREATE_NEVER").build());
+ }
+
+ @Test
+ public void testInvalidWriteDispositionIsRejected() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration().setWriteDisposition("INVALID").build();
IllegalArgumentException exception =
assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
@@ -149,46 +142,249 @@ public class SnowflakeWriteSchemaTransformProviderTest {
assertThat(
exception.getMessage(),
equalTo(
- "Unsupported writeDisposition 'INVALID'. Supported values are "
- + "APPEND, TRUNCATE, and EMPTY."));
+ "Unsupported writeDisposition 'INVALID'. "
+ + "Supported values are APPEND, TRUNCATE, and EMPTY."));
}
@Test
public void testSupportedWriteDispositions() {
provider.from(validConfiguration().setWriteDisposition("APPEND").build());
+
provider.from(validConfiguration().setWriteDisposition("TRUNCATE").build());
+
provider.from(validConfiguration().setWriteDisposition("EMPTY").build());
}
@Test
- public void testCreateNeverIsSupported() {
- Configuration configuration =
validConfiguration().setCreateDisposition("CREATE_NEVER").build();
+ public void testOauthAuthenticationIsValid() {
+ provider.from(
+
validConfiguration().setUsername(null).setPassword(null).setOauthToken("token").build());
+ }
- provider.from(configuration);
+ @Test
+ public void testPrivateKeyAuthenticationIsValid() {
+ provider.from(
+ validConfiguration()
+ .setPassword(null)
+ .setPrivateKey("private-key")
+ .setPrivateKeyPassphrase("passphrase")
+ .build());
}
- private static Configuration.Builder validConfiguration() {
- return Configuration.builder()
- .setServerName("account.snowflakecomputing.com")
- .setUsername("username")
- .setPassword("password")
- .setDatabase("database")
- .setSchema("schema")
- .setWarehouse("warehouse")
- .setRole("role")
- .setTable("table")
- .setStagingBucketName("gs://bucket/staging/")
- .setStorageIntegrationName("storage_integration");
+ @Test
+ public void testPrivateKeyWithoutPassphraseIsValid() {
+
provider.from(validConfiguration().setPassword(null).setPrivateKey("private-key").build());
}
@Test
- public void testConvertsBeamSchemaToSnowflakeSchema() {
+ public void testAuthenticationMethodIsRequired() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setOauthToken(null)
+ .setPrivateKey(null)
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo(
+ "Exactly one authentication method must be configured: "
+ + "password, oauthToken, or privateKey."));
+ }
+
+ @Test
+ public void testMultipleAuthenticationMethodsAreRejected() {
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setOauthToken("token").build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo(
+ "Exactly one authentication method must be configured: "
+ + "password, oauthToken, or privateKey."));
+ }
+
+ @Test
+ public void testUsernameIsRequiredForPasswordAuthentication() {
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setUsername(null).build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("username is required for password and private key
authentication."));
+ }
+
+ @Test
+ public void testUsernameIsRequiredForPrivateKeyAuthentication() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setPrivateKey("private-key")
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("username is required for password and private key
authentication."));
+ }
+
+ @Test
+ public void testPrivateKeyPassphraseRequiresPrivateKey() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration()
+ .setUsername(null)
+ .setPassword(null)
+ .setOauthToken("token")
+ .setPrivateKeyPassphrase("passphrase")
+ .build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("privateKeyPassphrase requires
privateKey."));
+ }
+
+ @Test
+ public void testInvalidDebugModeIsRejected() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration().setDebugMode("INVALID").build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(
+ exception.getMessage(),
+ equalTo("Unsupported debugMode 'INVALID'. " + "Supported values are
ERROR and INFO."));
+ }
+
+ @Test
+ public void testSupportedDebugModes() {
+ provider.from(validConfiguration().setDebugMode("ERROR").build());
+
+ provider.from(validConfiguration().setDebugMode("INFO").build());
+ }
+
+ @Test
+ public void testFlushRowLimitMustBePositive() {
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setFlushRowLimit(0).build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("flushRowLimit must be greater
than 0."));
+ }
+
+ @Test
+ public void testFlushTimeLimitMustBePositive() {
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration().setFlushTimeLimitMillis(0L).build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("flushTimeLimitMillis must be
greater than 0."));
+ }
+
+ @Test
+ public void testShardsNumberMustBePositive() {
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setShardsNumber(0).build();
+
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () ->
provider.from(configuration));
+
+ assertThat(exception.getMessage(), equalTo("shardsNumber must be greater
than 0."));
+ }
+
+ @Test
+ public void testBatchWriteRequiresTable() {
+ Schema schema =
Schema.builder().addInt64Field("id").addStringField("name").build();
+
+ Row row = Row.withSchema(schema).addValues(1L, "Alice").build();
+
+ PCollection<Row> rows =
+
pipeline.apply(Create.of(row).withCoder(RowCoder.of(schema))).setRowSchema(schema);
+
+ SnowflakeWriteConfiguration configuration =
validConfiguration().setTable(null).build();
+
+ SchemaTransform transform = provider.from(configuration);
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> transform.expand(PCollectionRowTuple.of("input", rows)));
+
+ assertThat(exception.getMessage(), equalTo("table is required for batch
writes."));
+ }
+
+ @Test
+ public void testStreamingWriteRequiresSnowPipe() {
+ Schema schema =
Schema.builder().addInt64Field("id").addStringField("name").build();
+
+ Row row = Row.withSchema(schema).addValues(1L, "Alice").build();
+
+ TestStream<Row> stream =
+
TestStream.create(RowCoder.of(schema)).addElements(row).advanceWatermarkToInfinity();
+
+ PCollection<Row> rows = pipeline.apply(stream).setRowSchema(schema);
+
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration().setTable(null).setSnowPipe(null).build();
+
+ SchemaTransform transform = provider.from(configuration);
+
+ IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> transform.expand(PCollectionRowTuple.of("input", rows)));
+
+ assertThat(exception.getMessage(), equalTo("snowPipe is required for
streaming writes."));
+ }
+
+ @Test
+ public void testStreamingConfigurationIsAccepted() {
+ Schema schema =
Schema.builder().addInt64Field("id").addStringField("name").build();
+
+ Row row = Row.withSchema(schema).addValues(1L, "Alice").build();
+
+ TestStream<Row> stream =
+
TestStream.create(RowCoder.of(schema)).addElements(row).advanceWatermarkToInfinity();
+
+ PCollection<Row> rows = pipeline.apply(stream).setRowSchema(schema);
+
+ SnowflakeWriteConfiguration configuration =
+ validConfiguration()
+ .setTable(null)
+ .setSnowPipe("MY_PIPE")
+ .setFlushRowLimit(50000)
+ .setFlushTimeLimitMillis(18000L)
+ .setShardsNumber(1)
+ .setDebugMode("ERROR")
+ .build();
+
+ SchemaTransform transform = provider.from(configuration);
+
+ transform.expand(PCollectionRowTuple.of("input", rows));
+ }
+
+ @Test
+ public void testScalarSchemaMapping() {
Schema schema =
Schema.builder()
.addByteField("byte_value")
- .addInt16Field("short_value")
- .addInt32Field("int_value")
- .addInt64Field("long_value")
+ .addInt16Field("int16_value")
+ .addInt32Field("int32_value")
+ .addInt64Field("int64_value")
.addFloatField("float_value")
.addDoubleField("double_value")
.addStringField("string_value")
@@ -198,20 +394,16 @@ public class SnowflakeWriteSchemaTransformProviderTest {
.build();
SnowflakeTableSchema snowflakeSchema =
- SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema);
+ SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema);
SnowflakeColumn[] columns = snowflakeSchema.getColumns();
- assertThat(columns.length, equalTo(10));
-
assertThat(columns[0].getDataType(), instanceOf(SnowflakeNumber.class));
assertThat(columns[1].getDataType(), instanceOf(SnowflakeNumber.class));
assertThat(columns[2].getDataType(), instanceOf(SnowflakeNumber.class));
assertThat(columns[3].getDataType(), instanceOf(SnowflakeNumber.class));
-
assertThat(columns[4].getDataType(), instanceOf(SnowflakeDouble.class));
assertThat(columns[5].getDataType(), instanceOf(SnowflakeDouble.class));
-
assertThat(columns[6].getDataType(), instanceOf(SnowflakeVarchar.class));
assertThat(columns[7].getDataType(), instanceOf(SnowflakeBoolean.class));
assertThat(columns[8].getDataType(), instanceOf(SnowflakeBinary.class));
@@ -219,31 +411,23 @@ public class SnowflakeWriteSchemaTransformProviderTest {
}
@Test
- public void testPreservesColumnNamesAndNullability() {
+ public void testSchemaMappingPreservesNameAndNullability() {
Schema schema =
Schema.builder()
- .addStringField("required_value")
- .addNullableField("optional_value", Schema.FieldType.INT64)
+ .addInt64Field("id")
+ .addNullableField("name", Schema.FieldType.STRING)
.build();
SnowflakeTableSchema snowflakeSchema =
- SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema);
+ SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema);
SnowflakeColumn[] columns = snowflakeSchema.getColumns();
- assertThat(columns[0].getName(), equalTo("required_value"));
+ assertThat(columns[0].getName(), equalTo("id"));
assertThat(columns[0].isNullable(), equalTo(false));
- assertThat(columns[1].getName(), equalTo("optional_value"));
+ assertThat(columns[1].getName(), equalTo("name"));
assertThat(columns[1].isNullable(), equalTo(true));
- }
-
- @Test
- public void testSnowflakeSchemaSql() {
- Schema schema =
Schema.builder().addInt64Field("id").addNullableStringField("name").build();
-
- SnowflakeTableSchema snowflakeSchema =
- SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema);
assertThat(snowflakeSchema.sql(), equalTo("id NUMBER(38,0), name VARCHAR
NULL"));
}
@@ -255,7 +439,7 @@ public class SnowflakeWriteSchemaTransformProviderTest {
IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
- () ->
SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema));
+ () ->
SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema));
assertThat(
exception.getMessage(),
@@ -271,7 +455,7 @@ public class SnowflakeWriteSchemaTransformProviderTest {
IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
- () ->
SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema));
+ () ->
SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema));
assertThat(
exception.getMessage(),
@@ -287,18 +471,24 @@ public class SnowflakeWriteSchemaTransformProviderTest {
IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
- () ->
SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema));
+ () ->
SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema));
assertThat(
exception.getMessage(),
equalTo("Unsupported Beam field type ROW for Snowflake column
'nested'."));
}
- @Test
- public void testCreateIfNeededIsSupported() {
- Configuration configuration =
- validConfiguration().setCreateDisposition("CREATE_IF_NEEDED").build();
-
- provider.from(configuration);
+ private static SnowflakeWriteConfiguration.Builder validConfiguration() {
+ return SnowflakeWriteConfiguration.builder()
+ .setServerName("account.snowflakecomputing.com")
+ .setUsername("username")
+ .setPassword("password")
+ .setDatabase("database")
+ .setSchema("public")
+ .setWarehouse("warehouse")
+ .setRole("role")
+ .setTable("table")
+ .setStagingBucketName("gs://bucket/staging/")
+ .setStorageIntegrationName("storage_integration");
}
}
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml
b/sdks/python/apache_beam/yaml/standard_io.yaml
index 9fc25b49fc6..4b5cb82b22f 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -208,27 +208,56 @@
# Snowflake
- type: renaming
transforms:
+ 'ReadFromSnowflake': 'ReadFromSnowflake'
'WriteToSnowflake': 'WriteToSnowflake'
config:
mappings:
+ 'ReadFromSnowflake':
+ server_name: 'server_name'
+ username: 'username'
+ password: 'password'
+ oauth_token: 'oauth_token'
+ private_key: 'private_key'
+ private_key_passphrase: 'private_key_passphrase'
+ database: 'database'
+ snowflake_schema: 'snowflake_schema'
+ warehouse: 'warehouse'
+ role: 'role'
+ table: 'table'
+ query: 'query'
+ staging_bucket_name: 'staging_bucket_name'
+ storage_integration_name: 'storage_integration_name'
+ schema: 'schema'
+ quotation_mark: 'quotation_mark'
'WriteToSnowflake':
server_name: 'server_name'
username: 'username'
password: 'password'
+ oauth_token: 'oauth_token'
+ private_key: 'private_key'
+ private_key_passphrase: 'private_key_passphrase'
database: 'database'
schema: 'schema'
warehouse: 'warehouse'
role: 'role'
table: 'table'
+ snow_pipe: 'snow_pipe'
staging_bucket_name: 'staging_bucket_name'
storage_integration_name: 'storage_integration_name'
create_disposition: 'create_disposition'
write_disposition: 'write_disposition'
quotation_mark: 'quotation_mark'
+ flush_row_limit: 'flush_row_limit'
+ flush_time_limit_millis: 'flush_time_limit_millis'
+ shards_number: 'shards_number'
+ debug_mode: 'debug_mode'
underlying_provider:
type: beamJar
transforms:
- 'WriteToSnowflake':
'beam:schematransform:org.apache.beam:snowflake_write:v1'
+ 'ReadFromSnowflake':
+ 'beam:schematransform:org.apache.beam:snowflake_read:v1'
+ 'WriteToSnowflake':
+ 'beam:schematransform:org.apache.beam:snowflake_write:v1'
config:
gradle_target: 'sdks:java:io:snowflake:expansion-service:shadowJar'
