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

commit 3f39865016f66da0301aea88d15451f04f4023aa
Author: Vitaly Terentyev <[email protected]>
AuthorDate: Thu Aug 6 16:31:39 2026 +0400

    Add Snowflake YAML write transform
---
 .../SnowflakeWriteSchemaTransformProvider.java     | 362 +++++++++++++++++++++
 .../SnowflakeWriteSchemaTransformProviderTest.java | 304 +++++++++++++++++
 sdks/python/apache_beam/yaml/standard_io.yaml      |  27 ++
 3 files changed, 693 insertions(+)

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
new file mode 100644
index 00000000000..623bb134f58
--- /dev/null
+++ 
b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java
@@ -0,0 +1,362 @@
+/*
+ * 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 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;
+
+/** A {@link SchemaTransformProvider} for writing Beam rows to Snowflake. */
+@SuppressWarnings({
+  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
+})
+@AutoService(SchemaTransformProvider.class)
+public class SnowflakeWriteSchemaTransformProvider
+    extends 
TypedSchemaTransformProvider<SnowflakeWriteSchemaTransformProvider.Configuration>
 {
+
+  static final String INPUT_TAG = "input";
+
+  public static final String IDENTIFIER = 
"beam:schematransform:org.apache.beam:snowflake_write:v1";
+
+  @Override
+  public String identifier() {
+    return IDENTIFIER;
+  }
+
+  @Override
+  public String description() {
+    return "Writes Beam Rows to a Snowflake table using staged CSV files.";
+  }
+
+  @Override
+  protected Class<Configuration> configurationClass() {
+    return Configuration.class;
+  }
+
+  @Override
+  protected SchemaTransform from(Configuration configuration) {
+    configuration.validate();
+    return new SnowflakeWriteSchemaTransform(configuration);
+  }
+
+  @Override
+  public List<String> inputCollectionNames() {
+    return Collections.singletonList(INPUT_TAG);
+  }
+
+  @Override
+  public List<String> outputCollectionNames() {
+    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 SnowflakeWriteSchemaTransform(Configuration configuration) {
+      this.configuration = configuration;
+    }
+
+    @Override
+    public PCollectionRowTuple expand(PCollectionRowTuple input) {
+      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());
+      }
+
+      SnowflakeIO.Write<Row> write =
+          SnowflakeIO.<Row>write()
+              .withDataSourceConfiguration(dataSourceConfiguration)
+              .withStagingBucketName(configuration.getStagingBucketName())
+              
.withStorageIntegrationName(configuration.getStorageIntegrationName())
+              .withUserDataMapper(row -> row.getValues().toArray())
+              .to(configuration.getTable());
+
+      if (configuration.getCreateDisposition() != null) {
+        CreateDisposition createDisposition =
+            parseCreateDisposition(configuration.getCreateDisposition());
+
+        write = write.withCreateDisposition(createDisposition);
+
+        if (createDisposition == CreateDisposition.CREATE_IF_NEEDED) {
+          write = 
write.withTableSchema(toSnowflakeTableSchema(rows.getSchema()));
+        }
+      }
+
+      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();
+
+    @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.")
+    public abstract String getTable();
+
+    @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. Supported values are CREATE_IF_NEEDED and 
CREATE_NEVER.")
+    @Nullable
+    public abstract String getCreateDisposition();
+
+    @SchemaFieldDescription("Write behavior. 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();
+
+    public static Builder builder() {
+      return new 
AutoValue_SnowflakeWriteSchemaTransformProvider_Configuration.Builder();
+    }
+
+    public abstract Builder toBuilder();
+
+    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 (getCreateDisposition() != null) {
+        parseCreateDisposition(getCreateDisposition());
+      }
+
+      if (getWriteDisposition() != null) {
+        parseWriteDisposition(getWriteDisposition());
+      }
+    }
+
+    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(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());
+
+    if (details != null) {
+      message += " " + details;
+    }
+
+    return new IllegalArgumentException(message);
+  }
+}
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
new file mode 100644
index 00000000000..336d7f01f03
--- /dev/null
+++ 
b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java
@@ -0,0 +1,304 @@
+/*
+ * 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.hamcrest.Matchers.instanceOf;
+import static org.junit.Assert.assertThrows;
+
+import 
org.apache.beam.sdk.io.snowflake.SnowflakeWriteSchemaTransformProvider.Configuration;
+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;
+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.schemas.Schema;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class SnowflakeWriteSchemaTransformProviderTest {
+
+  private final SnowflakeWriteSchemaTransformProvider provider =
+      new SnowflakeWriteSchemaTransformProvider();
+
+  @Test
+  public void testIdentifier() {
+    assertThat(
+        provider.identifier(), 
equalTo("beam:schematransform:org.apache.beam:snowflake_write:v1"));
+  }
+
+  @Test
+  public void testInputCollectionNames() {
+    assertThat(provider.inputCollectionNames(), contains("input"));
+  }
+
+  @Test
+  public void testOutputCollectionNames() {
+    assertThat(provider.outputCollectionNames(), empty());
+  }
+
+  @Test
+  public void testValidConfiguration() {
+    Configuration configuration = validConfiguration().build();
+
+    provider.from(configuration);
+  }
+
+  @Test
+  public void testBlankQuotationMarkIsAllowed() {
+    Configuration configuration = 
validConfiguration().setQuotationMark("").build();
+
+    provider.from(configuration);
+  }
+
+  @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();
+
+    IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> 
provider.from(configuration));
+
+    assertThat(exception.getMessage(), equalTo("serverName cannot be empty"));
+  }
+
+  @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"));
+  }
+
+  @Test
+  public void testStagingBucketMustEndWithSlash() {
+    Configuration configuration =
+        
validConfiguration().setStagingBucketName("gs://bucket/staging").build();
+
+    IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> 
provider.from(configuration));
+
+    assertThat(exception.getMessage(), equalTo("stagingBucketName must end 
with '/'"));
+  }
+
+  @Test
+  public void testInvalidCreateDisposition() {
+    Configuration configuration = 
validConfiguration().setCreateDisposition("INVALID").build();
+
+    IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> 
provider.from(configuration));
+
+    assertThat(
+        exception.getMessage(),
+        equalTo(
+            "Unsupported createDisposition 'INVALID'. Supported values are "
+                + "CREATE_IF_NEEDED and CREATE_NEVER."));
+  }
+
+  @Test
+  public void testInvalidWriteDisposition() {
+    Configuration configuration = 
validConfiguration().setWriteDisposition("INVALID").build();
+
+    IllegalArgumentException exception =
+        assertThrows(IllegalArgumentException.class, () -> 
provider.from(configuration));
+
+    assertThat(
+        exception.getMessage(),
+        equalTo(
+            "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();
+
+    provider.from(configuration);
+  }
+
+  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 testConvertsBeamSchemaToSnowflakeSchema() {
+    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();
+
+    SnowflakeTableSchema snowflakeSchema =
+        SnowflakeWriteSchemaTransformProvider.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));
+    assertThat(columns[9].getDataType(), instanceOf(SnowflakeTimestamp.class));
+  }
+
+  @Test
+  public void testPreservesColumnNamesAndNullability() {
+    Schema schema =
+        Schema.builder()
+            .addStringField("required_value")
+            .addNullableField("optional_value", Schema.FieldType.INT64)
+            .build();
+
+    SnowflakeTableSchema snowflakeSchema =
+        SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema);
+
+    SnowflakeColumn[] columns = snowflakeSchema.getColumns();
+
+    assertThat(columns[0].getName(), equalTo("required_value"));
+    assertThat(columns[0].isNullable(), equalTo(false));
+
+    assertThat(columns[1].getName(), equalTo("optional_value"));
+    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"));
+  }
+
+  @Test
+  public void testDecimalIsRejected() {
+    Schema schema = Schema.builder().addDecimalField("amount").build();
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> 
SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema));
+
+    assertThat(
+        exception.getMessage(),
+        equalTo(
+            "Unsupported Beam field type DECIMAL for Snowflake column 
'amount'. "
+                + "Beam DECIMAL does not include Snowflake precision and scale 
information."));
+  }
+
+  @Test
+  public void testArrayIsRejected() {
+    Schema schema = Schema.builder().addArrayField("values", 
Schema.FieldType.STRING).build();
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> 
SnowflakeWriteSchemaTransformProvider.toSnowflakeTableSchema(schema));
+
+    assertThat(
+        exception.getMessage(),
+        equalTo("Unsupported Beam field type ARRAY for Snowflake column 
'values'."));
+  }
+
+  @Test
+  public void testNestedRowIsRejected() {
+    Schema nestedSchema = Schema.builder().addStringField("value").build();
+
+    Schema schema = Schema.builder().addRowField("nested", 
nestedSchema).build();
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> 
SnowflakeWriteSchemaTransformProvider.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);
+  }
+}
diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml 
b/sdks/python/apache_beam/yaml/standard_io.yaml
index 796429b5bdf..9fc25b49fc6 100644
--- a/sdks/python/apache_beam/yaml/standard_io.yaml
+++ b/sdks/python/apache_beam/yaml/standard_io.yaml
@@ -205,6 +205,33 @@
       config:
         gradle_target: 'sdks:java:io:debezium:expansion-service:shadowJar'
 
+# Snowflake
+- type: renaming
+  transforms:
+    'WriteToSnowflake': 'WriteToSnowflake'
+  config:
+    mappings:
+      'WriteToSnowflake':
+        server_name: 'server_name'
+        username: 'username'
+        password: 'password'
+        database: 'database'
+        schema: 'schema'
+        warehouse: 'warehouse'
+        role: 'role'
+        table: 'table'
+        staging_bucket_name: 'staging_bucket_name'
+        storage_integration_name: 'storage_integration_name'
+        create_disposition: 'create_disposition'
+        write_disposition: 'write_disposition'
+        quotation_mark: 'quotation_mark'
+    underlying_provider:
+      type: beamJar
+      transforms:
+        'WriteToSnowflake': 
'beam:schematransform:org.apache.beam:snowflake_write:v1'
+      config:
+        gradle_target: 'sdks:java:io:snowflake:expansion-service:shadowJar'
+
 # Databases
 - type: renaming
   transforms:

Reply via email to