diff --git a/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json b/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json index 37dd25bf9029..34a6e02150e7 100644 --- a/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json +++ b/.github/trigger_files/beam_CloudML_Benchmarks_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 3 + "modification": 4 } diff --git a/.test-infra/tools/stale_cleaner.py b/.test-infra/tools/stale_cleaner.py index d59af061ce9b..313c9490ba54 100644 --- a/.test-infra/tools/stale_cleaner.py +++ b/.test-infra/tools/stale_cleaner.py @@ -365,6 +365,7 @@ def clean_pubsub_topics(): prefixes = [ "psit_topic_input", "psit_topic_output", + "psit_topic_ordering", "wc_topic_input", "wc_topic_output", "leader_board_it_input_topic", @@ -421,6 +422,12 @@ def clean_pubsub_subscriptions(): # Restrict subscription cleanup to the NYC taxi prefix only. prefixes = [ "taxirides-realtime_beam_", + "pubsub_io_performance", + "psit_sub_input", + "psit_sub_output", + "psit_sub_ordering", + "wc_subscription_input", + "wc_subscription_output", ] # Create a PubSubSubscriptionCleaner instance diff --git a/infra/enforcement/README.md b/infra/enforcement/README.md index b92e5f7e1802..90fc76974758 100644 --- a/infra/enforcement/README.md +++ b/infra/enforcement/README.md @@ -138,7 +138,7 @@ The enforcement tools are consolidated into a single daily workflow (`.github/wo This unified workflow executes both security domains sequentially: - **IAM Policy Enforcement:** Validates user bindings against the defined policies. -- **Unmanaged Keys Audit:** Detects rogue service account keys generated outside the official rotation system. +- **Unmanaged Keys Audit:** Detects rogue service account keys generated outside the official rotation system and reports them to the `[IAC_DRIFT_SA_KEY]` issue. **Note**: - **Manual trigger**: The workflow can also be triggered manually via `workflow_dispatch`. @@ -175,7 +175,7 @@ python account_keys.py --action generate - **check**: Validates service account keys and their permissions against defined policies and reports any differences (default behavior) - **announce**: Creates or updates a GitHub issue and sends an email notification when service account keys policies differ from the defined ones. - For general configuration errors, it updates the main compliance issue. - - **For unmanaged/rogue keys (Security Alerts)**, it consolidates alerts into a dedicated `[SECURITY]` issue acting as a live dashboard. It updates the issue by placing the newest audit report at the top and moving the previous reports into a collapsed `
` history section. If the keys are revoked and the infrastructure becomes healthy, the system automatically resolves and closes the issue. + - **For unmanaged/rogue keys**, it consolidates alerts into a dedicated `[IAC_DRIFT_SA_KEY]` issue acting as a live dashboard. It updates the issue by placing the newest audit report at the top and moving the previous reports into a collapsed `
` history section. If the keys are revoked and the infrastructure becomes healthy, the system automatically resolves and closes the issue. - **print**: Prints announcement details for testing purposes without creating actual GitHub issues or sending emails - **generate**: Updates the compliance file to match the current GCP service account keys and Secret Manager permissions diff --git a/infra/enforcement/account_keys.py b/infra/enforcement/account_keys.py index a1248e3ee09c..3c173d2a8afc 100644 --- a/infra/enforcement/account_keys.py +++ b/infra/enforcement/account_keys.py @@ -28,6 +28,7 @@ SECRET_MANAGER_LABEL = "beam-infra-secret-manager" IAC_DRIFT_SA_KEY = "IAC_DRIFT_SA_KEY" +ACCOUNT_KEYS_POLICY = "ACCOUNT_KEYS_POLICY" class AuthorizedUser(TypedDict): email: str @@ -385,8 +386,8 @@ def create_announcement(self, recipient: str) -> None: if general_issues: self.logger.info(f"Found {len(general_issues)} general compliance issues. Triggering announcement...") - title = f"[{IAC_DRIFT_SA_KEY}] Action Required: Unauthorized Service Accounts Detected" - body = f"Unauthorized Service Accounts Report\n\n" + title = f"[{ACCOUNT_KEYS_POLICY}] Action Required: Service Account Policy Drift" + body = f"Service Account Policy Drift Report\n\n" body += f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" for issue in general_issues: body += f"- {issue}\n" @@ -423,8 +424,8 @@ def print_announcement(self, recipient: str) -> None: if general_issues: self.logger.info("Printing general compliance announcement...") - title = f"[IAC_DRIFT_SA_KEY] Action Required: Unauthorized Service Accounts Detected" - body = f"Unauthorized Service Accounts Report\n\n" + title = f"[{ACCOUNT_KEYS_POLICY}] Action Required: Service Account Policy Drift" + body = f"Service Account Policy Drift Report\n\n" body += f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" for issue in general_issues: body += f"- {issue}\n" diff --git a/infra/enforcement/sending.py b/infra/enforcement/sending.py index 9d24a816fbc3..67857145aea0 100644 --- a/infra/enforcement/sending.py +++ b/infra/enforcement/sending.py @@ -228,7 +228,7 @@ def report_unmanaged_keys(self, project_id: str, compilance_issues: List[str]) - self.logger.info("No compliance issues to report to Github.") return - issue_title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected" + issue_title = "[IAC_DRIFT_SA_KEY] Action Required: Unmanaged Service Account Keys Detected" #markdown body timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") new_report = f"### Unmanaged Keys Audit Report ({timestamp})\n" @@ -238,6 +238,12 @@ def report_unmanaged_keys(self, project_id: str, compilance_issues: List[str]) - new_report += f"- {issue_text}\n" new_report += "\n*Please investigate and revoke these keys if they are not part of the official rotation system.*" + remediation = "\n\n### Remediation\n" + remediation += "1. Delete all reported keys as soon as possible.\n" + remediation += "2. Replace the deleted keys using the official Beam key rotation system. It creates the service account key and registers its key ID and private key in the corresponding managed Secret Manager secret (`-key`). Do not create replacement keys manually in IAM.\n" + remediation += "3. Run the audit again to confirm that the reported keys have been removed and the replacement keys are managed by the rotation system.\n" + remediation += "\nFor more information, consult `infra/keys/README.md`." + new_report += remediation open_issues = self._get_open_issues(issue_title) if open_issues: @@ -272,7 +278,7 @@ def resolve_unmanaged_keys(self) -> None: Finds any open security issues regarding rogue keys and automatically closes them if the infrastructure is now healthy. """ - issue_title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected" + issue_title = "[IAC_DRIFT_SA_KEY] Action Required: Unmanaged Service Account Keys Detected" open_issues = self._get_open_issues(issue_title) if open_issues: target_issue = open_issues[0] diff --git a/infra/enforcement/test_sending.py b/infra/enforcement/test_sending.py index 26d4080adec5..90104e7bb13e 100644 --- a/infra/enforcement/test_sending.py +++ b/infra/enforcement/test_sending.py @@ -51,7 +51,7 @@ def test_get_open_issues_flaky_retry(self, mock_request): "items": [ { "number": 1234, - "title": "[SECURITY] Action Required: Unmanaged Service Account Keys Detected", + "title": "[IAC_DRIFT_SA_KEY] Action Required: Unmanaged Service Account Keys Detected", "body": "Test body", "state": "open", "html_url": "https://github.com/apache/beam/issues/1234", @@ -64,7 +64,7 @@ def test_get_open_issues_flaky_retry(self, mock_request): mock_request.side_effect = [mock_response_fail, mock_response_success] # Call get_open_issues - issues = self.client._get_open_issues("[SECURITY] Action Required: Unmanaged Service Account Keys Detected") + issues = self.client._get_open_issues("[IAC_DRIFT_SA_KEY] Action Required: Unmanaged Service Account Keys Detected") # Verify that two requests were made (one retry) self.assertEqual(mock_request.call_count, 2) @@ -79,7 +79,7 @@ def test_get_open_issues_query_format(self, mock_request): mock_response.json.return_value = {"items": []} mock_request.return_value = mock_response - title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected" + title = "[IAC_DRIFT_SA_KEY] Action Required: Unmanaged Service Account Keys Detected" self.client._get_open_issues(title) # Verify that the query parameter was passed correctly to requests diff --git a/sdks/java/io/snowflake/build.gradle b/sdks/java/io/snowflake/build.gradle index 8d9a9a46557f..4286c9fb8e11 100644 --- a/sdks/java/io/snowflake/build.gradle +++ b/sdks/java/io/snowflake/build.gradle @@ -30,6 +30,8 @@ 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 + provided library.java.everit_json_schema + permitUnusedDeclared 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/SnowflakeIO.java b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java index d25a0d31c919..4267ff7324cc 100644 --- a/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java +++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeIO.java @@ -34,7 +34,6 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.stream.Collectors; -import javax.annotation.Nullable; import javax.sql.DataSource; import net.snowflake.client.api.datasource.SnowflakeDataSource; import net.snowflake.client.api.datasource.SnowflakeDataSourceFactory; @@ -87,6 +86,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Joiner; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.Duration; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -231,6 +231,7 @@ public interface CsvMapper extends Serializable { */ @FunctionalInterface public interface UserDataMapper extends Serializable { + @Nullable Object[] mapRow(T element); } 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 000000000000..affd1fda651b --- /dev/null +++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeReadSchemaTransformProvider.java @@ -0,0 +1,289 @@ +/* + * 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. */ +@AutoService(SchemaTransformProvider.class) +public class SnowflakeReadSchemaTransformProvider + extends TypedSchemaTransformProvider { + + 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 configurationClass() { + return Configuration.class; + } + + @Override + protected SchemaTransform from(Configuration configuration) { + configuration.validate(); + return new SnowflakeReadSchemaTransform(configuration); + } + + @Override + public List inputCollectionNames() { + return Collections.emptyList(); + } + + @Override + public List 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 read = + SnowflakeIO.read() + .withDataSourceConfiguration(dataSourceConfiguration) + .withStagingBucketName(configuration.getStagingBucketName()) + .withStorageIntegrationName(configuration.getStorageIntegrationName()) + .withCsvMapper(parts -> toRow(parts, outputSchema)) + .withCoder(RowCoder.of(outputSchema)); + + String table = configuration.getTable(); + if (table != null) { + read = read.fromTable(table); + } else { + String query = configuration.getQuery(); + if (query != null) { + read = read.fromQuery(query); + } + } + + String quotationMark = configuration.getQuotationMark(); + if (quotationMark != null) { + read = read.withQuotationMark(quotationMark); + } + + PCollection 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()); + + String table = getTable(); + boolean tablePresent = table != null && !table.isEmpty(); + String query = getQuery(); + boolean queryPresent = query != null && !query.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 000000000000..694775879ee0 --- /dev/null +++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeSchemaTransformUtils.java @@ -0,0 +1,342 @@ +/* + * 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 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.checkerframework.checker.nullness.qual.EnsuresNonNullIf; +import org.joda.time.Instant; + +/** Utilities shared by Snowflake schema transform providers. */ +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) && isNotEmpty(username)) { + configuration = configuration.withUsernamePasswordAuth(username, password); + } else if (isNotEmpty(oauthToken)) { + configuration = configuration.withOAuth(oauthToken); + } else if (isNotEmpty(privateKey) && isNotEmpty(username)) { + 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."); + } + } + + @EnsuresNonNullIf(expression = "#1", result = true) + 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 @Nullable 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 decodeHex(value); + + 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); + } + } + + private static byte[] decodeHex(String value) { + if ((value.length() & 1) != 0) { + throw new IllegalArgumentException("Invalid hexadecimal Snowflake binary value."); + } + + byte[] result = new byte[value.length() / 2]; + + for (int i = 0; i < value.length(); i += 2) { + int high = Character.digit(value.charAt(i), 16); + int low = Character.digit(value.charAt(i + 1), 16); + + if (high == -1 || low == -1) { + throw new IllegalArgumentException("Invalid hexadecimal Snowflake binary value."); + } + + result[i / 2] = (byte) ((high << 4) | low); + } + + return result; + } + + 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 000000000000..0225367b45ed --- /dev/null +++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteConfiguration.java @@ -0,0 +1,227 @@ +/* + * 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.parseCreateDisposition; +import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.parseStreamingLogLevel; +import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.parseWriteDisposition; + +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 '/'"); + } + + // Parse configured enum values to validate that they are supported. + 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 new file mode 100644 index 000000000000..72f1700e7671 --- /dev/null +++ b/sdks/java/io/snowflake/src/main/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProvider.java @@ -0,0 +1,179 @@ +/* + * 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.parseCreateDisposition; +import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.parseStreamingLogLevel; +import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.parseWriteDisposition; +import static org.apache.beam.sdk.io.snowflake.SnowflakeSchemaTransformUtils.toSnowflakeTableSchema; + +import com.google.auto.service.AutoService; +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import org.apache.beam.sdk.io.snowflake.enums.CreateDisposition; +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. */ +@AutoService(SchemaTransformProvider.class) +public class SnowflakeWriteSchemaTransformProvider + extends TypedSchemaTransformProvider { + + 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 Snowflake using batch COPY or streaming Snowpipe."; + } + + @Override + protected Class configurationClass() { + return SnowflakeWriteConfiguration.class; + } + + @Override + protected SchemaTransform from(SnowflakeWriteConfiguration configuration) { + configuration.validate(); + return new SnowflakeWriteSchemaTransform(configuration); + } + + @Override + public List inputCollectionNames() { + return Collections.singletonList(INPUT_TAG); + } + + @Override + public List outputCollectionNames() { + return Collections.emptyList(); + } + + private static class SnowflakeWriteSchemaTransform extends SchemaTransform + implements Serializable { + + private final SnowflakeWriteConfiguration configuration; + + private SnowflakeWriteSchemaTransform(SnowflakeWriteConfiguration configuration) { + this.configuration = configuration; + } + + @Override + public PCollectionRowTuple expand(PCollectionRowTuple input) { + PCollection rows = input.get(INPUT_TAG); + + SnowflakeIO.DataSourceConfiguration dataSourceConfiguration = + SnowflakeSchemaTransformUtils.createDataSourceConfiguration( + configuration.getServerName(), + configuration.getUsername(), + configuration.getPassword(), + configuration.getOauthToken(), + configuration.getPrivateKey(), + configuration.getPrivateKeyPassphrase(), + configuration.getDatabase(), + configuration.getSchema(), + configuration.getWarehouse(), + configuration.getRole()); + + SnowflakeIO.Write write = + SnowflakeIO.write() + .withDataSourceConfiguration(dataSourceConfiguration) + .withStagingBucketName(configuration.getStagingBucketName()) + .withStorageIntegrationName(configuration.getStorageIntegrationName()) + .withUserDataMapper(row -> row.getValues().toArray()); + + boolean streaming = rows.isBounded() == PCollection.IsBounded.UNBOUNDED; + + if (streaming) { + String snowPipe = configuration.getSnowPipe(); + + if (snowPipe == null || snowPipe.isEmpty()) { + throw new IllegalArgumentException("snowPipe is required for streaming writes."); + } + + write = write.withSnowPipe(snowPipe); + + Integer flushRowLimit = configuration.getFlushRowLimit(); + if (flushRowLimit != null) { + write = write.withFlushRowLimit(flushRowLimit); + } + + Long flushTimeLimitMillis = configuration.getFlushTimeLimitMillis(); + if (flushTimeLimitMillis != null) { + write = write.withFlushTimeLimit(Duration.millis(flushTimeLimitMillis)); + } + + Integer shardsNumber = configuration.getShardsNumber(); + if (shardsNumber != null) { + write = write.withShardsNumber(shardsNumber); + } + + String debugMode = configuration.getDebugMode(); + if (debugMode != null) { + write = write.withDebugMode(parseStreamingLogLevel(debugMode)); + } + } else { + String table = configuration.getTable(); + + if (table == null || table.isEmpty()) { + throw new IllegalArgumentException("table is required for batch writes."); + } + + write = write.to(table); + + String createDispositionValue = configuration.getCreateDisposition(); + + if (createDispositionValue != null) { + CreateDisposition createDisposition = parseCreateDisposition(createDispositionValue); + + write = write.withCreateDisposition(createDisposition); + + if (createDisposition == CreateDisposition.CREATE_IF_NEEDED) { + write = write.withTableSchema(toSnowflakeTableSchema(rows.getSchema())); + } + } + + String writeDispositionValue = configuration.getWriteDisposition(); + + if (writeDispositionValue != null) { + write = write.withWriteDisposition(parseWriteDisposition(writeDispositionValue)); + } + } + + String quotationMark = configuration.getQuotationMark(); + if (quotationMark != null) { + write = write.withQuotationMark(quotationMark); + } + + rows.apply("WriteToSnowflake", write); + + return PCollectionRowTuple.empty(input.getPipeline()); + } + } +} 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 000000000000..440103b7cfb4 --- /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", "616263", "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 new file mode 100644 index 000000000000..bdd5aae0e75f --- /dev/null +++ b/sdks/java/io/snowflake/src/test/java/org/apache/beam/sdk/io/snowflake/SnowflakeWriteSchemaTransformProviderTest.java @@ -0,0 +1,494 @@ +/* + * 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.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; +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.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; + +@RunWith(JUnit4.class) +public class SnowflakeWriteSchemaTransformProviderTest { + + private final transient Pipeline pipeline = Pipeline.create(); + + 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() { + provider.from(validConfiguration().build()); + } + + @Test + public void testBlankQuotationMarkIsAllowed() { + provider.from(validConfiguration().setQuotationMark("").build()); + } + + @Test + public void testMissingServerNameIsRejected() { + SnowflakeWriteConfiguration configuration = validConfiguration().setServerName("").build(); + + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> provider.from(configuration)); + + assertThat(exception.getMessage(), equalTo("serverName cannot be empty")); + } + + @Test + public void testMissingTableIsAllowedAtConfigurationTime() { + // Whether table is required depends on whether the input is bounded. + provider.from(validConfiguration().setTable(null).build()); + } + + @Test + public void testStagingBucketMustEndWithSlash() { + SnowflakeWriteConfiguration 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 testInvalidCreateDispositionIsRejected() { + SnowflakeWriteConfiguration 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 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)); + + 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 testOauthAuthenticationIsValid() { + provider.from( + validConfiguration().setUsername(null).setPassword(null).setOauthToken("token").build()); + } + + @Test + public void testPrivateKeyAuthenticationIsValid() { + provider.from( + validConfiguration() + .setPassword(null) + .setPrivateKey("private-key") + .setPrivateKeyPassphrase("passphrase") + .build()); + } + + @Test + public void testPrivateKeyWithoutPassphraseIsValid() { + provider.from(validConfiguration().setPassword(null).setPrivateKey("private-key").build()); + } + + @Test + 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 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 stream = + TestStream.create(RowCoder.of(schema)).addElements(row).advanceWatermarkToInfinity(); + + PCollection 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 stream = + TestStream.create(RowCoder.of(schema)).addElements(row).advanceWatermarkToInfinity(); + + PCollection 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("int16_value") + .addInt32Field("int32_value") + .addInt64Field("int64_value") + .addFloatField("float_value") + .addDoubleField("double_value") + .addStringField("string_value") + .addBooleanField("boolean_value") + .addByteArrayField("bytes_value") + .addDateTimeField("datetime_value") + .build(); + + SnowflakeTableSchema snowflakeSchema = + SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema); + + SnowflakeColumn[] columns = snowflakeSchema.getColumns(); + + 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 testSchemaMappingPreservesNameAndNullability() { + Schema schema = + Schema.builder() + .addInt64Field("id") + .addNullableField("name", Schema.FieldType.STRING) + .build(); + + SnowflakeTableSchema snowflakeSchema = + SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema); + + SnowflakeColumn[] columns = snowflakeSchema.getColumns(); + + assertThat(columns[0].getName(), equalTo("id")); + assertThat(columns[0].isNullable(), equalTo(false)); + + assertThat(columns[1].getName(), equalTo("name")); + assertThat(columns[1].isNullable(), equalTo(true)); + + 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, + () -> SnowflakeSchemaTransformUtils.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, + () -> SnowflakeSchemaTransformUtils.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, + () -> SnowflakeSchemaTransformUtils.toSnowflakeTableSchema(schema)); + + assertThat( + exception.getMessage(), + equalTo("Unsupported Beam field type ROW for Snowflake column 'nested'.")); + } + + 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/io/iobase.py b/sdks/python/apache_beam/io/iobase.py index aa03280050fa..b7be85935997 100644 --- a/sdks/python/apache_beam/io/iobase.py +++ b/sdks/python/apache_beam/io/iobase.py @@ -1716,15 +1716,19 @@ def total_work(self) -> float: def fraction_completed(self) -> float: if self._fraction is not None: return self._fraction - else: - return float(self._completed) / self.total_work + total_work = self.total_work + if total_work == 0.: + return 1.0 + return float(self._completed) / total_work @property def fraction_remaining(self) -> float: if self._fraction is not None: return 1 - self._fraction - else: - return float(self._remaining) / self.total_work + total_work = self.total_work + if total_work == 0.: + return 0.0 + return float(self._remaining) / total_work def with_completed(self, completed: int) -> 'RestrictionProgress': return RestrictionProgress( diff --git a/sdks/python/apache_beam/io/iobase_test.py b/sdks/python/apache_beam/io/iobase_test.py index 60c261563155..4621fa45bb15 100644 --- a/sdks/python/apache_beam/io/iobase_test.py +++ b/sdks/python/apache_beam/io/iobase_test.py @@ -261,5 +261,26 @@ def test_read_unbounded_serializes_as_expanded_composite(self): self.assertTrue(read_transforms[0].subtransforms) +class RestrictionProgressTest(unittest.TestCase): + def test_restriction_progress(self): + # Total work == 0 edge cases (avoids ZeroDivisionError) + progress_zero_int = iobase.RestrictionProgress(completed=0, remaining=0) + self.assertEqual(progress_zero_int.fraction_completed, 1.0) + self.assertEqual(progress_zero_int.fraction_remaining, 0.0) + + # Progress with completed and remaining + progress_work = iobase.RestrictionProgress(completed=25, remaining=75) + self.assertEqual(progress_work.completed_work, 25) + self.assertEqual(progress_work.remaining_work, 75) + self.assertEqual(progress_work.total_work, 100) + self.assertEqual(progress_work.fraction_completed, 0.25) + self.assertEqual(progress_work.fraction_remaining, 0.75) + + # Progress with fraction + progress_frac = iobase.RestrictionProgress(fraction=0.4) + self.assertEqual(progress_frac.fraction_completed, 0.4) + self.assertAlmostEqual(progress_frac.fraction_remaining, 0.6) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/yaml/extended_tests/databases/snowflake.yaml b/sdks/python/apache_beam/yaml/extended_tests/databases/snowflake.yaml new file mode 100644 index 000000000000..7ddf5c337de7 --- /dev/null +++ b/sdks/python/apache_beam/yaml/extended_tests/databases/snowflake.yaml @@ -0,0 +1,90 @@ +# +# 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. +# + +fixtures: + - name: SNOWFLAKE + type: apache_beam.yaml.integration_tests.snowflake_fixture + +pipelines: + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - number_column: 1 + boolean_column: true + string_column: one + - number_column: 2 + boolean_column: false + string_column: two + + - type: WriteToSnowflake + config: + server_name: "{SNOWFLAKE[SERVER_NAME]}" + username: "{SNOWFLAKE[USERNAME]}" + password: "{SNOWFLAKE[PASSWORD]}" + database: "{SNOWFLAKE[DATABASE]}" + schema: "{SNOWFLAKE[SCHEMA]}" + table: "{SNOWFLAKE[TABLE]}" + staging_bucket_name: "{SNOWFLAKE[STAGING_BUCKET_NAME]}" + storage_integration_name: "{SNOWFLAKE[STORAGE_INTEGRATION_NAME]}" + create_disposition: CREATE_IF_NEEDED + write_disposition: TRUNCATE + + - pipeline: + type: chain + transforms: + - type: ReadFromSnowflake + config: + server_name: "{SNOWFLAKE[SERVER_NAME]}" + username: "{SNOWFLAKE[USERNAME]}" + password: "{SNOWFLAKE[PASSWORD]}" + database: "{SNOWFLAKE[DATABASE]}" + snowflake_schema: "{SNOWFLAKE[SCHEMA]}" + table: "{SNOWFLAKE[TABLE]}" + staging_bucket_name: "{SNOWFLAKE[STAGING_BUCKET_NAME]}" + storage_integration_name: "{SNOWFLAKE[STORAGE_INTEGRATION_NAME]}" + schema: | + { + "type": "object", + "properties": { + "number_column": { + "type": "integer" + }, + "boolean_column": { + "type": "boolean" + }, + "string_column": { + "type": "string" + } + }, + "required": [ + "number_column", + "boolean_column", + "string_column" + ] + } + - type: AssertEqual + config: + elements: + - number_column: 1 + boolean_column: true + string_column: one + - number_column: 2 + boolean_column: false + string_column: two diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index 649e77f1513a..32794f4588d2 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -397,6 +397,41 @@ def temp_mysql_database(): yield jdbc_url +@contextlib.contextmanager +def snowflake_fixture(): + options = PipelineOptions().get_all_options() + + required = [ + 'server_name', + 'username', + 'password', + 'staging_bucket_name', + 'storage_integration_name', + 'database', + 'schema', + 'table', + ] + + missing = [name for name in required if not options.get(name)] + if missing: + raise unittest.SkipTest( + 'Snowflake YAML integration test requires external configuration: ' + + ', '.join(missing)) + + yield { + 'SERVER_NAME': options['server_name'], + 'USERNAME': options['username'], + 'PASSWORD': options['password'], + 'STAGING_BUCKET_NAME': options['staging_bucket_name'], + 'STORAGE_INTEGRATION_NAME': options['storage_integration_name'], + 'DATABASE': options['database'], + 'SCHEMA': options['schema'], + 'TABLE': options['table'], + 'WAREHOUSE': options.get('warehouse'), + 'ROLE': options.get('role'), + } + + @contextlib.contextmanager def temp_debezium_postgres_database(): """Provides a temporary PostgreSQL database configured for Debezium CDC.""" diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index f1ba727cc609..253573c8069a 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -308,6 +308,62 @@ config: gradle_target: 'sdks:java:io:amazon-web-services2:expansion-service:shadowJar' +# 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: + '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' + # Databases - type: renaming transforms: diff --git a/sdks/python/build.gradle b/sdks/python/build.gradle index a15d4719ff7b..75afa5d7c968 100644 --- a/sdks/python/build.gradle +++ b/sdks/python/build.gradle @@ -151,7 +151,6 @@ tasks.register("yamlIntegrationTests") { dependsOn ":sdks:java:extensions:sql:expansion-service:shadowJar" dependsOn ":sdks:java:io:expansion-service:build" dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build" - dependsOn ":sdks:java:io:debezium:expansion-service:shadowJar" dependsOn ":sdks:java:io:messaging-expansion-service:shadowJar" doLast { @@ -173,6 +172,7 @@ tasks.register("postCommitYamlIntegrationTests") { dependsOn ":sdks:java:io:expansion-service:build" dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build" dependsOn ":sdks:java:io:debezium:expansion-service:shadowJar" + dependsOn ":sdks:java:io:snowflake:expansion-service:shadowJar" dependsOn ":sdks:java:io:amazon-web-services2:expansion-service:shadowJar" doLast { diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 3bdbd4df41b6..abe7867181ec 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -592,10 +592,12 @@ task anthropicInferenceTest { task installTFTRequirements { dependsOn 'initializeForDataflowJob' doLast { + // Keep the SDK under test: TFT requirements can pull a different Beam from PyPI. + def sdk = project.ext.sdkLocation exec { workingDir "$rootProject.projectDir/sdks/python/apache_beam/testing/benchmarks/cloudml/" executable 'sh' - args '-c', ". ${envdir}/bin/activate && pip install -r requirements.txt" + args '-c', ". ${envdir}/bin/activate && pip install -r requirements.txt && pip install --force-reinstall --no-deps '${sdk}[gcp]'" } } } @@ -607,6 +609,8 @@ task tftTests { doLast { def opts = project.findProperty('opts') opts += " --sdk_location=${project.ext.sdkLocation}" + // Reinstall the staged SDK after requirements on workers (runs after requirements install). + opts += " --extra_package=${project.ext.sdkLocation}" def testOpts = basicPytestOpts + ["--numprocesses=8", "--dist=loadfile"] def argMap = [ "test_opts": testOpts,