diff --git a/.github/workflows/build-report.yml b/.github/workflows/build-report.yml new file mode 100644 index 0000000..2271363 --- /dev/null +++ b/.github/workflows/build-report.yml @@ -0,0 +1,50 @@ +# Copyright © 2024 Cask Data, Inc. +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +# This workflow will build a Java project with Maven +# For more information see: https://help.github.com/actions/language-and-framework-guides/building-and-testing-java-with-maven +# Note: Any changes to this workflow would be used only after merging into develop +name: Build Unit Tests Report + +on: + workflow_run: + workflows: + - Build with unit tests + types: + - completed + +jobs: + build: + runs-on: ubuntu-latest + + if: ${{ github.event.workflow_run.conclusion != 'skipped' }} + + steps: + # Pinned 1.0.0 version + - uses: marocchino/action-workflow_run-status@54b6e87d6cb552fc5f36dbe9a722a6048725917a + + - name: Download artifact + uses: actions/download-artifact@v4 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + run-id: ${{ github.event.workflow_run.id }} + path: artifacts/ + + - name: Surefire Report + # Pinned 3.5.2 version + uses: mikepenz/action-junit-report@16a9560bd02f11e7e3bf6b3e2ef6bba6c9d07c32 + if: always() + with: + report_paths: '**/target/surefire-reports/TEST-*.xml' + github_token: ${{ secrets.GITHUB_TOKEN }} + detailed_summary: true + commit: ${{ github.event.workflow_run.head_sha }} + check_name: Build Test Report \ No newline at end of file diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml new file mode 100644 index 0000000..ce71df5 --- /dev/null +++ b/.github/workflows/main.yml @@ -0,0 +1,59 @@ +# Copyright © 2024 Cask Data, Inc. +# Licensed under the Apache License, Version 2.0 (the "License"); you may not +# use this file except in compliance with the License. You may obtain a copy of +# the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations under +# the License. + +# This workflow will build a Java project with Maven +# For more information see: https://help.github.com/actions/language-and-framework-guides/building-and-testing-java-with-maven +# Note: Any changes to this workflow would be used only after merging into develop +name: Build with unit tests + +on: + push: + branches: [ develop, release/** ] + pull_request: + branches: [ develop, release/** ] + types: [opened, synchronize, reopened, labeled] + +jobs: + build: + runs-on: k8s-runner-build + + # We allow builds: + # 1) When it's a merge into a branch + # 2) For PRs that are labeled as build and + # - It's a code change + # - A build label was just added + # A bit complex, but prevents builds when other labels are manipulated + if: > + github.event_name == 'push' + || (contains(github.event.pull_request.labels.*.name, 'build') + && (github.event.action != 'labeled' || github.event.label.name == 'build') + ) + steps: + - uses: actions/checkout@v3 + with: + ref: ${{ github.event.workflow_run.head_sha }} + - name: Cache + uses: actions/cache@v3 + with: + path: ~/.m2/repository + key: ${{ runner.os }}-maven-${{ github.workflow }}-${{ hashFiles('**/pom.xml') }} + restore-keys: | + ${{ runner.os }}-maven-${{ github.workflow }} + - name: Build with Maven + run: mvn clean test -fae -T 2 -B -V -DcloudBuild -Dmaven.wagon.http.retryHandler.count=3 -Dmaven.wagon.httpconnectionManager.ttlSeconds=25 + - name: Archive build artifacts + uses: actions/upload-artifact@v4 + if: always() + with: + name: reports-${{ github.run_id }} + path: | + **/target/rat.txt + **/target/surefire-reports/* diff --git a/docs/Snowflake-batchsource.md b/docs/Snowflake-batchsource.md index 76d7a1d..9b810fe 100644 --- a/docs/Snowflake-batchsource.md +++ b/docs/Snowflake-batchsource.md @@ -25,7 +25,9 @@ log in to Snowflake, minus the "snowflakecomputing.com"). E.g. "myaccount.us-cen **Role:** Role to use (e.g. `ACCOUNTADMIN`). -**Import Query:** Query for data import. +**Import Query Type** - Method used to retrieve schema from the source. +* **Table Name**: The name of the table to retrieve the schema. +* **Import Query**: Query for data import. ### Credentials @@ -68,7 +70,7 @@ Data Types Mapping ---------- | Snowflake Data Types | CDAP Schema Data Type | Comment | -| ------------------------------ | --------------------- | --------------------------------------------------------- | +| ------------------------------ |-----------------------| --------------------------------------------------------- | | NUMBER | decimal | Default precision and scale are (38,0). | | DECIMAL | decimal | Synonymous with NUMBER. | | NUMERIC | decimal | Synonymous with NUMBER. | @@ -87,10 +89,10 @@ Data Types Mapping | DATE | date | | | DATETIME | timestamp | Alias for TIMESTAMP_NTZ | | TIME | time | | -| TIMESTAMP | timestamp/string | Alias for one of the TIMESTAMP variations (TIMESTAMP_NTZ by default).| +| TIMESTAMP | timestamp | Alias for one of the TIMESTAMP variations (TIMESTAMP_NTZ by default).| | TIMESTAMP_LTZ | timestamp | TIMESTAMP with local time zone; time zone, if provided, is not stored.| | TIMESTAMP_NTZ | timestamp | TIMESTAMP with no time zone; time zone, if provided, is not stored.| -| TIMESTAMP_TZ | string | TIMESTAMP with time zone. | +| TIMESTAMP_TZ | timestamp | TIMESTAMP with time zone. | | VARIANT | string | A tagged universal type, which can store values of any other type, including OBJECT and ARRAY, up to a maximum size of 16 MB compressed. | | OBJECT | string | This will return a json with the data | | ARRAY | string | This will return a json with the data | | \ No newline at end of file diff --git a/pom.xml b/pom.xml index 4f72c1f..43f6ad6 100644 --- a/pom.xml +++ b/pom.xml @@ -21,7 +21,7 @@ io.cdap.plugin snowflake-plugins - 1.1.0-SNAPSHOT + 1.1.6-SNAPSHOT jar Snowflake plugins @@ -29,8 +29,8 @@ true UTF-8 - 6.1.0-SNAPSHOT - 2.3.0-SNAPSHOT + 6.4.1 + 2.6.1 1.6 2.8.0 2.3.1 @@ -47,17 +47,13 @@ 1.7.4 4.0 2.4 - 3.13.24 + 4.3.4 - - sonatype - https://oss.sonatype.org/content/groups/public - sonatype-snapshots - https://oss.sonatype.org/content/repositories/snapshots + https://central.sonatype.com/repository/maven-snapshots @@ -398,8 +394,8 @@ 1.1.0 - system:cdap-data-pipeline[6.1.0-SNAPSHOT,7.0.0-SNAPSHOT) - system:cdap-data-streams[6.1.0-SNAPSHOT,7.0.0-SNAPSHOT) + system:cdap-data-pipeline[6.4.1-SNAPSHOT,7.0.0-SNAPSHOT) + system:cdap-data-streams[6.4.1-SNAPSHOT,7.0.0-SNAPSHOT) diff --git a/src/main/java/io/cdap/plugin/snowflake/common/BaseSnowflakeConfig.java b/src/main/java/io/cdap/plugin/snowflake/common/BaseSnowflakeConfig.java index 718a707..2ccb2c9 100644 --- a/src/main/java/io/cdap/plugin/snowflake/common/BaseSnowflakeConfig.java +++ b/src/main/java/io/cdap/plugin/snowflake/common/BaseSnowflakeConfig.java @@ -278,7 +278,7 @@ public boolean canConnect() { && !containsMacro(PROPERTY_PASSWORD) && !containsMacro(PROPERTY_WAREHOUSE) && !containsMacro(PROPERTY_ROLE) && !containsMacro(PROPERTY_CLIENT_ID) && !containsMacro(PROPERTY_CLIENT_SECRET) && !containsMacro(PROPERTY_REFRESH_TOKEN) - && !containsMacro(PROPERTY_PRIVATE_KEY)); + && !containsMacro(PROPERTY_PRIVATE_KEY) && !containsMacro(PROPERTY_PASSPHRASE)); } protected void validateConnection(FailureCollector collector) { @@ -299,8 +299,8 @@ protected void validateConnection(FailureCollector collector) { .withConfigProperty(PROPERTY_USERNAME); // TODO: for oauth2 - if (keyPairEnabled) { - failure.withConfigProperty(PROPERTY_PRIVATE_KEY); + if (Boolean.TRUE.equals(keyPairEnabled)) { + failure.withConfigProperty(PROPERTY_PRIVATE_KEY).withConfigProperty(PROPERTY_PASSPHRASE); } else { failure.withConfigProperty(PROPERTY_PASSWORD); } diff --git a/src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java b/src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java index 3169bea..38aa738 100644 --- a/src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java +++ b/src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java @@ -23,7 +23,7 @@ import io.cdap.plugin.snowflake.common.OAuthUtil; import io.cdap.plugin.snowflake.common.exception.ConnectionTimeoutException; import io.cdap.plugin.snowflake.common.util.QueryUtil; -import net.snowflake.client.jdbc.SnowflakeBasicDataSource; +import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource; import org.apache.http.impl.client.HttpClients; import java.io.BufferedWriter; @@ -32,6 +32,7 @@ import java.io.IOException; import java.lang.reflect.Field; import java.sql.Connection; +import java.sql.DatabaseMetaData; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.ResultSetMetaData; @@ -66,6 +67,30 @@ public void runSQL(String query) throws IOException { } + /** + * Returns field descriptors for specified table name + * @param schemaName The name of schema containing the table + * @param tableName The name of table whose metadata needs to be retrieved + * @return list of field descriptors + */ + public List getFieldDescriptors(String schemaName, String tableName) throws IOException { + List fieldDescriptors = new ArrayList<>(); + try (Connection connection = dataSource.getConnection()) { + DatabaseMetaData dbMetaData = connection.getMetaData(); + try (ResultSet columns = dbMetaData.getColumns(null, schemaName, tableName, null)) { + while (columns.next()) { + String columnName = columns.getString("COLUMN_NAME"); + int columnType = columns.getInt("DATA_TYPE"); + boolean nullable = columns.getInt("NULLABLE") == DatabaseMetaData.columnNullable; + fieldDescriptors.add(new SnowflakeFieldDescriptor(columnName, columnType, nullable)); + } + } + } catch (SQLException e) { + throw new IOException(); + } + return fieldDescriptors; + } + /** * Returns field descriptors for specified import query. * @@ -110,7 +135,11 @@ private void initDataSource(SnowflakeBasicDataSource dataSource, BaseSnowflakeCo if (config.getOauth2Enabled()) { String accessToken = OAuthUtil.getAccessTokenByRefreshToken(HttpClients.createDefault(), config); - dataSource.setOauthToken(accessToken); + // In JDBC 4.x, setOauthToken() was removed. The recommended approach is to explicitly + // set the authenticator to "oauth" and pass the access token as the password. + // Migration guide: https://docs.snowflake.com/en/developer-guide/jdbc/jdbc-migration + dataSource.setAuthenticator("oauth"); + dataSource.setPassword(accessToken); } else if (config.getKeyPairEnabled()) { dataSource.setUser(config.getUsername()); @@ -171,4 +200,13 @@ private static String writeTextToTmpFile(String text) { throw new RuntimeException("Cannot write key to temporary file", e); } } + + /** + * Retrieves schema name from the configuration + * + * @return The schema name + */ + public String getSchema() { + return config.getSchemaName(); + } } diff --git a/src/main/java/io/cdap/plugin/snowflake/common/util/QueryUtil.java b/src/main/java/io/cdap/plugin/snowflake/common/util/QueryUtil.java index 72d52c9..64f3e22 100644 --- a/src/main/java/io/cdap/plugin/snowflake/common/util/QueryUtil.java +++ b/src/main/java/io/cdap/plugin/snowflake/common/util/QueryUtil.java @@ -16,19 +16,24 @@ package io.cdap.plugin.snowflake.common.util; +import com.google.common.base.Strings; + /** * Transforms import query. */ public class QueryUtil { // Matches "limit ". Also "limit $$$$" and "limit ''" which means unlimited in Snowflake. - private static final String LIMIT_PATTERN = "(?i)LIMIT (''|\\$\\$\\$\\$|\\d+)"; + private static final String LIMIT_PATTERN = "(?i)LIMIT (NULL|''|\\$\\$\\$\\$|\\d+)"; private static final String LIMIT_STRING = "limit %s"; private QueryUtil() { } public static String removeSemicolon(String importQuery) { + if (Strings.isNullOrEmpty(importQuery)) { + return null; + } if (importQuery.endsWith(";")) { importQuery = importQuery.substring(0, importQuery.length() - 1); } diff --git a/src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java b/src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java index 9448f23..161c132 100644 --- a/src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java +++ b/src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java @@ -24,13 +24,16 @@ import io.cdap.plugin.snowflake.common.client.SnowflakeFieldDescriptor; import io.cdap.plugin.snowflake.common.exception.SchemaParseException; import io.cdap.plugin.snowflake.source.batch.SnowflakeBatchSourceConfig; +import io.cdap.plugin.snowflake.source.batch.SnowflakeInputFormatProvider; import io.cdap.plugin.snowflake.source.batch.SnowflakeSourceAccessor; import java.io.IOException; +import java.sql.SQLException; import java.sql.Types; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; +import javax.annotation.Nullable; /** * Resolves schema. @@ -51,31 +54,52 @@ public class SchemaHelper { .put(Types.BOOLEAN, Schema.of(Schema.Type.BOOLEAN)) .put(Types.BIGINT, Schema.decimalOf(38)) .put(Types.SMALLINT, Schema.decimalOf(38)) + .put(Types.TIMESTAMP_WITH_TIMEZONE, Schema.of(Schema.LogicalType.TIMESTAMP_MICROS)) .build(); private SchemaHelper() { } + /** + * Retrieves schema for the Snowflake batch source based on the given configuration. + * + * @param config The configuration for Snowflake batch source + * @param collector The failure collector to capture any schema retrieval errors. + * @return The resolved schema for Snowflake source + */ public static Schema getSchema(SnowflakeBatchSourceConfig config, FailureCollector collector) { if (!config.canConnect()) { - return null; + return getParsedSchema(config.getSchema()); } - SnowflakeSourceAccessor snowflakeSourceAccessor = new SnowflakeSourceAccessor(config); - return getSchema(snowflakeSourceAccessor, config.getSchema(), collector, config.getImportQuery()); + SnowflakeSourceAccessor snowflakeSourceAccessor = + new SnowflakeSourceAccessor(config, SnowflakeInputFormatProvider.PROPERTY_DEFAULT_ESCAPE_CHAR); + return getSchema( + snowflakeSourceAccessor, + config.getSchema(), + collector, + config.getTableName(), + config.getImportQuery() + ); } + /** + * Retrieves schema for a Snowflake source based on the provided parameters. + * + * @param snowflakeAccessor The {@link SnowflakeSourceAccessor} used to connect to Snowflake. + * @param schema A JSON-format schema string + * @param collector The {@link FailureCollector} to collect errors if schema retrieval fails. + * @param tableName The name of the table in Snowflake. + * @param importQuery The query to fetch data from Snowflake, used when `tableName` is not provided. + * @return The parsed {@link Schema} if successful, or {@code null} if an error occurs. + */ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String schema, - FailureCollector collector, String importQuery) { + FailureCollector collector, String tableName, String importQuery) { try { if (!Strings.isNullOrEmpty(schema)) { - try { - return Schema.parseJson(schema); - } catch (IOException | IllegalStateException e) { - throw new SchemaParseException(e); - } + return getParsedSchema(schema); } - return Strings.isNullOrEmpty(importQuery) ? null : getSchema(snowflakeAccessor, importQuery); + return getSchema(snowflakeAccessor, tableName, importQuery); } catch (SchemaParseException e) { collector.addFailure(String.format("Unable to retrieve output schema. Reason: '%s'", e.getMessage()), null) @@ -85,6 +109,37 @@ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String } } + public static Schema getParsedSchema(String schema) { + if (Strings.isNullOrEmpty(schema)) { + return null; + } + try { + return Schema.parseJson(schema); + } catch (IOException | IllegalStateException e) { + throw new SchemaParseException(e); + } + } + + private static Schema getSchema(SnowflakeAccessor snowflakeAccessor, @Nullable String tableName, + @Nullable String importQuery) { + try { + List result; + if (!Strings.isNullOrEmpty(tableName)) { + result = snowflakeAccessor.getFieldDescriptors(snowflakeAccessor.getSchema(), tableName); + } else if (!Strings.isNullOrEmpty(importQuery)) { + result = snowflakeAccessor.describeQuery(importQuery); + } else { + return null; + } + List fields = result.stream() + .map(fieldDescriptor -> Schema.Field.of(fieldDescriptor.getName(), getSchema(fieldDescriptor))) + .collect(Collectors.toList()); + return Schema.recordOf("data", fields); + } catch (IOException e) { + throw new SchemaParseException(e); + } + } + public static Schema getSchema(SnowflakeAccessor snowflakeAccessor, String importQuery) { try { List result = snowflakeAccessor.describeQuery(importQuery); diff --git a/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkAccessor.java b/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkAccessor.java index 0eba82d..7e18837 100644 --- a/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkAccessor.java +++ b/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkAccessor.java @@ -17,7 +17,8 @@ package io.cdap.plugin.snowflake.sink.batch; import io.cdap.plugin.snowflake.common.client.SnowflakeAccessor; -import net.snowflake.client.jdbc.SnowflakeConnection; +import net.snowflake.client.api.connection.SnowflakeConnection; +import net.snowflake.client.api.connection.UploadStreamConfig; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; @@ -51,9 +52,10 @@ public void uploadStream(InputStream inputStream, String stageDir) throws IOExce LOG.info("Uploading file '{}' to table stage", filename); try (Connection connection = dataSource.getConnection()) { + UploadStreamConfig uploadConfig = UploadStreamConfig.builder().setDestPrefix(null).setCompressData(true).build(); connection.unwrap(SnowflakeConnection.class).uploadStream(stageDir, - null, - inputStream, filename, true); + filename, + inputStream, uploadConfig); } catch (SQLException e) { throw new IOException(e); } diff --git a/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java b/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java index a30d28f..4ae2291 100644 --- a/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java +++ b/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java @@ -98,8 +98,8 @@ public void validate(Schema inputSchema, FailureCollector failureCollector) { * @param schema input schema to check */ private void validateInputSchema(Schema schema, FailureCollector failureCollector) { - // schema can be null in case it is a macro - if (schema == null) { + // return if schema is null or if any of the connection properties or table name is macro + if (schema == null || !canConnect() || containsMacro(PROPERTY_TABLE_NAME)) { return; } diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/ImportQueryType.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/ImportQueryType.java new file mode 100644 index 0000000..1f170ab --- /dev/null +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/ImportQueryType.java @@ -0,0 +1,49 @@ +/* + * Copyright © 2025 Cask Data, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ + +package io.cdap.plugin.snowflake.source.batch; + + +/** + * Enum to specify the import query type used in Snowflake Batch Source. + */ +public enum ImportQueryType { + IMPORT_QUERY ("importQuery"), + TABLE_NAME ("tableName"); + + private String value; + + ImportQueryType(String value) { + this.value = value; + } + + public String getValue() { + return value; + } + + public static ImportQueryType fromString(String value) { + if (value == null) { + return ImportQueryType.IMPORT_QUERY; + } + + for (ImportQueryType type : ImportQueryType.values()) { + if (type.value.equalsIgnoreCase(value)) { + return type; + } + } + return ImportQueryType.IMPORT_QUERY; + } +} diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSource.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSource.java index 53abea9..f5fde15 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSource.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSource.java @@ -16,6 +16,7 @@ package io.cdap.plugin.snowflake.source.batch; +import com.google.common.base.Strings; import io.cdap.cdap.api.annotation.Description; import io.cdap.cdap.api.annotation.Name; import io.cdap.cdap.api.annotation.Plugin; @@ -33,6 +34,7 @@ import io.cdap.plugin.snowflake.common.util.SchemaHelper; import org.apache.hadoop.io.NullWritable; +import java.util.HashMap; import java.util.Map; import java.util.stream.Collectors; @@ -68,7 +70,11 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) { public void prepareRun(BatchSourceContext context) { FailureCollector failureCollector = context.getFailureCollector(); config.validate(failureCollector); - + Map arguments = new HashMap<>(context.getArguments().asMap()); + String escapeChar = arguments.containsKey(SnowflakeInputFormatProvider.PROPERTY_ESCAPE_CHAR) && + !Strings.isNullOrEmpty(arguments.get(SnowflakeInputFormatProvider.PROPERTY_ESCAPE_CHAR)) + ? arguments.get(SnowflakeInputFormatProvider.PROPERTY_ESCAPE_CHAR) + : SnowflakeInputFormatProvider.PROPERTY_DEFAULT_ESCAPE_CHAR; Schema schema = SchemaHelper.getSchema(config, failureCollector); failureCollector.getOrThrowException(); @@ -81,14 +87,13 @@ public void prepareRun(BatchSourceContext context) { .collect(Collectors.toList())); } - context.setInput(Input.of(config.getReferenceName(), new SnowflakeInputFormatProvider(config))); + context.setInput(Input.of(config.getReferenceName(), new SnowflakeInputFormatProvider(config, escapeChar))); } @Override public void initialize(BatchRuntimeContext context) throws Exception { super.initialize(context); - SnowflakeSourceAccessor snowflakeAccessor = new SnowflakeSourceAccessor(config); - Schema schema = SchemaHelper.getSchema(snowflakeAccessor, config.getImportQuery()); + Schema schema = SchemaHelper.getSchema(config, context.getFailureCollector()); this.transformer = new SnowflakeMapToRecordTransformer(schema); } diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfig.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfig.java index 9cd9ac8..66e23cc 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfig.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfig.java @@ -16,10 +16,14 @@ package io.cdap.plugin.snowflake.source.batch; +import com.google.common.base.Strings; import io.cdap.cdap.api.annotation.Description; import io.cdap.cdap.api.annotation.Macro; import io.cdap.cdap.api.annotation.Name; +import io.cdap.cdap.etl.api.FailureCollector; import io.cdap.plugin.snowflake.common.BaseSnowflakeConfig; + +import java.util.Objects; import javax.annotation.Nullable; /** @@ -31,6 +35,8 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig { public static final String PROPERTY_IMPORT_QUERY = "importQuery"; public static final String PROPERTY_MAX_SPLIT_SIZE = "maxSplitSize"; public static final String PROPERTY_SCHEMA = "schema"; + public static final String PROPERTY_TABLE_NAME = "tableName"; + public static final String PROPERTY_IMPORT_QUERY_TYPE = "importQueryType"; @Name(PROPERTY_REFERENCE_NAME) @Description("This will be used to uniquely identify this source/sink for lineage, annotating metadata, etc.") @@ -39,6 +45,7 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig { @Name(PROPERTY_IMPORT_QUERY) @Description("Query for import data.") @Macro + @Nullable private String importQuery; @Name(PROPERTY_MAX_SPLIT_SIZE) @@ -52,19 +59,40 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig { @Macro private String schema; + + @Name(PROPERTY_TABLE_NAME) + @Description("The name of the table used to retrieve the schema.") + @Macro + @Nullable + private final String tableName; + + @Name(PROPERTY_IMPORT_QUERY_TYPE) + @Description("Whether to select Table Name or Import Query to extract the data.") + @Macro + @Nullable + private final String importQueryType; + + public SnowflakeBatchSourceConfig(String referenceName, String accountName, String database, - String schemaName, String importQuery, String username, String password, + String schemaName, @Nullable String importQuery, String username, String password, @Nullable Boolean keyPairEnabled, @Nullable String path, @Nullable String passphrase, @Nullable Boolean oauth2Enabled, @Nullable String clientId, @Nullable String clientSecret, @Nullable String refreshToken, Long maxSplitSize, - @Nullable String connectionArguments, @Nullable String schema) { - super(accountName, database, schemaName, username, password, - keyPairEnabled, path, passphrase, oauth2Enabled, clientId, clientSecret, refreshToken, connectionArguments); + @Nullable String connectionArguments, + @Nullable String schema, + @Nullable String tableName, + @Nullable String importQueryType) { + super( + accountName, database, schemaName, username, password, keyPairEnabled, path, passphrase, + oauth2Enabled, clientId, clientSecret, refreshToken, connectionArguments + ); this.referenceName = referenceName; this.importQuery = importQuery; this.maxSplitSize = maxSplitSize; this.schema = schema; + this.tableName = tableName; + this.importQueryType = getImportQueryType(); } public String getImportQuery() { @@ -83,4 +111,47 @@ public String getReferenceName() { public String getSchema() { return schema; } + + @Nullable + public String getTableName() { + return tableName; + } + + @Nullable + public String getImportQueryType() { + return importQueryType == null ? ImportQueryType.IMPORT_QUERY.name() : importQueryType; + } + + public void validate(FailureCollector collector) { + super.validate(collector); + + if (!containsMacro(PROPERTY_MAX_SPLIT_SIZE) && Objects.nonNull(maxSplitSize) + && maxSplitSize < 0) { + collector.addFailure("Maximum Slit Size cannot be a negative number.", null) + .withConfigProperty(PROPERTY_MAX_SPLIT_SIZE); + } + + if (!containsMacro(PROPERTY_IMPORT_QUERY_TYPE)) { + boolean isImportQuerySelected = ImportQueryType.IMPORT_QUERY.getValue().equals(importQueryType); + boolean isTableNameSelected = ImportQueryType.TABLE_NAME.getValue().equals(importQueryType); + + if (isImportQuerySelected && !containsMacro(PROPERTY_IMPORT_QUERY) && Strings.isNullOrEmpty(importQuery)) { + collector.addFailure("Import Query cannot be empty", null) + .withConfigProperty(PROPERTY_IMPORT_QUERY); + + } else if (isTableNameSelected && !containsMacro(PROPERTY_TABLE_NAME) && Strings.isNullOrEmpty(tableName)) { + collector.addFailure("Table Name cannot be empty", null) + .withConfigProperty(PROPERTY_TABLE_NAME); + } + } else { + boolean isImportQueryMissing = !containsMacro(PROPERTY_IMPORT_QUERY) && Strings.isNullOrEmpty(importQuery); + boolean isTableNameMissing = !containsMacro(PROPERTY_TABLE_NAME) && Strings.isNullOrEmpty(tableName); + + if (isImportQueryMissing && isTableNameMissing) { + collector.addFailure("Either 'Import Query' or 'Table Name' must be provided.", null) + .withConfigProperty(PROPERTY_IMPORT_QUERY) + .withConfigProperty(PROPERTY_TABLE_NAME); + } + } + } } diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormat.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormat.java index 5b655e3..c4c1d1b 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormat.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormat.java @@ -60,6 +60,8 @@ private SnowflakeSourceAccessor getSnowflakeAccessor(Configuration configuration SnowflakeInputFormatProvider.PROPERTY_CONFIG_JSON); SnowflakeBatchSourceConfig config = GSON.fromJson( configJson, SnowflakeBatchSourceConfig.class); - return new SnowflakeSourceAccessor(config); + String escapeChar = configuration.get(SnowflakeInputFormatProvider.PROPERTY_ESCAPE_CHAR, + SnowflakeInputFormatProvider.PROPERTY_DEFAULT_ESCAPE_CHAR); + return new SnowflakeSourceAccessor(config, escapeChar); } } diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormatProvider.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormatProvider.java index 9020599..97f6996 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormatProvider.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeInputFormatProvider.java @@ -29,13 +29,17 @@ public class SnowflakeInputFormatProvider implements InputFormatProvider { public static final String PROPERTY_CONFIG_JSON = "cdap.snowflake.source.config"; + public static final String PROPERTY_ESCAPE_CHAR = "cdap.snowflake.source.escape"; + + public static final String PROPERTY_DEFAULT_ESCAPE_CHAR = "\\"; private static final Gson GSON = new Gson(); private final Map conf; - public SnowflakeInputFormatProvider(SnowflakeBatchSourceConfig config) { + public SnowflakeInputFormatProvider(SnowflakeBatchSourceConfig config, String escapeChar) { this.conf = new ImmutableMap.Builder() .put(PROPERTY_CONFIG_JSON, GSON.toJson(config)) + .put(PROPERTY_ESCAPE_CHAR, escapeChar) .build(); } diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformer.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformer.java index 2f75653..d8990c0 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformer.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformer.java @@ -24,6 +24,7 @@ import org.slf4j.LoggerFactory; import java.math.BigDecimal; +import java.math.RoundingMode; import java.time.Instant; import java.time.LocalDate; import java.time.LocalTime; @@ -83,7 +84,8 @@ private Object convertValue(String fieldName, String value, Schema fieldSchema) case TIME_MICROS: return TimeUnit.NANOSECONDS.toMicros(LocalTime.parse(value).toNanoOfDay()); case DECIMAL: - return new BigDecimal(value).setScale(fieldSchema.getScale()).unscaledValue().toByteArray(); + return new BigDecimal(value).setScale(fieldSchema.getScale(), + RoundingMode.HALF_EVEN).unscaledValue().toByteArray(); default: throw new IllegalArgumentException( String.format("Field '%s' is of unsupported type '%s'", fieldSchema.getDisplayName(), diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java index 9bda860..41fdcf1 100644 --- a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java +++ b/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java @@ -17,10 +17,13 @@ package io.cdap.plugin.snowflake.source.batch; import au.com.bytecode.opencsv.CSVReader; +import com.google.common.base.Strings; +import io.cdap.cdap.api.data.schema.Schema; import io.cdap.plugin.snowflake.common.client.SnowflakeAccessor; import io.cdap.plugin.snowflake.common.util.QueryUtil; -import io.cdap.plugin.snowflake.sink.batch.SnowflakeSinkAccessor; -import net.snowflake.client.jdbc.SnowflakeConnection; +import io.cdap.plugin.snowflake.common.util.SchemaHelper; +import net.snowflake.client.api.connection.DownloadStreamConfig; +import net.snowflake.client.api.connection.SnowflakeConnection; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; @@ -33,12 +36,13 @@ import java.util.ArrayList; import java.util.List; import java.util.UUID; +import java.util.stream.Collectors; /** * A class which accesses Snowflake API to do actions used by batch source. */ public class SnowflakeSourceAccessor extends SnowflakeAccessor { - private static final Logger LOG = LoggerFactory.getLogger(SnowflakeSinkAccessor.class); + private static final Logger LOG = LoggerFactory.getLogger(SnowflakeSourceAccessor.class); // Directory should be unique, so that parallel pipelines can run correctly, as well as after failure we don't // have old stage files in the dir. private static final String STAGE_PATH = "@~/cdap_stage/result" + UUID.randomUUID() + "/"; @@ -60,10 +64,12 @@ public class SnowflakeSourceAccessor extends SnowflakeAccessor { "OVERWRITE=TRUE HEADER=TRUE SINGLE=FALSE"; private static final String COMMAND_MAX_FILE_SIZE = " MAX_FILE_SIZE=%s"; private final SnowflakeBatchSourceConfig config; + private final char escapeChar; - public SnowflakeSourceAccessor(SnowflakeBatchSourceConfig config) { + public SnowflakeSourceAccessor(SnowflakeBatchSourceConfig config, String escapeChar) { super(config); this.config = config; + this.escapeChar = escapeChar.charAt(0); } /** @@ -74,7 +80,20 @@ public SnowflakeSourceAccessor(SnowflakeBatchSourceConfig config) { */ public List prepareStageSplits() throws IOException { LOG.info("Loading data into stage: '{}'", STAGE_PATH); - String copy = String.format(COMAND_COPY_INTO, QueryUtil.removeSemicolon(config.getImportQuery())); + String importQuery = config.getImportQuery(); + if (Strings.isNullOrEmpty(importQuery)) { + String tableName = config.getTableName(); + Schema schema = SchemaHelper.getParsedSchema(config.getSchema()); + if (schema != null && schema.getFields() != null && !schema.getFields().isEmpty()) { + String columns = schema.getFields().stream() + .map(Schema.Field::getName) + .collect(Collectors.joining(",")); + importQuery = String.format("SELECT %s FROM %s", columns, tableName); + } else { + importQuery = String.format("SELECT * FROM %s", tableName); + } + } + String copy = String.format(COMAND_COPY_INTO, QueryUtil.removeSemicolon(importQuery)); if (config.getMaxSplitSize() > 0) { copy = copy + String.format(COMMAND_MAX_FILE_SIZE, config.getMaxSplitSize()); } @@ -113,10 +132,11 @@ public void removeStageFile(String stageSplit) throws IOException { */ public CSVReader buildCsvReader(String stageSplit) throws IOException { try (Connection connection = dataSource.getConnection()) { + DownloadStreamConfig downloadStreamConfig = DownloadStreamConfig.builder().setDecompress(true).build(); InputStream downloadStream = connection.unwrap(SnowflakeConnection.class) - .downloadStream("@~", stageSplit, true); + .downloadStream("@~", stageSplit, downloadStreamConfig); InputStreamReader inputStreamReader = new InputStreamReader(downloadStream); - return new CSVReader(inputStreamReader); + return new CSVReader(inputStreamReader, ',', '"', escapeChar); } catch (SQLException e) { throw new IOException(e); } diff --git a/src/test/java/io/cdap/plugin/snowflake/common/BaseSnowflakeTest.java b/src/test/java/io/cdap/plugin/snowflake/common/BaseSnowflakeTest.java index cd79627..b021cc8 100644 --- a/src/test/java/io/cdap/plugin/snowflake/common/BaseSnowflakeTest.java +++ b/src/test/java/io/cdap/plugin/snowflake/common/BaseSnowflakeTest.java @@ -17,13 +17,15 @@ package io.cdap.plugin.snowflake.common; import io.cdap.cdap.etl.mock.test.HydratorTestBase; +import io.cdap.cdap.test.TestConfiguration; import io.cdap.plugin.snowflake.Constants; import io.cdap.plugin.snowflake.common.client.SnowflakeAccessorTest; import io.cdap.plugin.snowflake.source.batch.SnowflakeBatchSourceConfig; import io.cdap.plugin.snowflake.source.batch.SnowflakeBatchSourceConfigBuilder; -import net.snowflake.client.jdbc.SnowflakeBasicDataSource; +import net.snowflake.client.internal.api.implementation.datasource.SnowflakeBasicDataSource; import org.junit.Assume; import org.junit.BeforeClass; +import org.junit.ClassRule; import org.junit.Rule; import org.junit.internal.AssumptionViolatedException; import org.junit.rules.TestName; @@ -55,6 +57,9 @@ public abstract class BaseSnowflakeTest extends HydratorTestBase { private static final Logger LOG = LoggerFactory.getLogger(SnowflakeAccessorTest.class); + @ClassRule + public static final TestConfiguration TEST_CONFIG = new TestConfiguration("explore.enabled", false); + protected static final String ACCOUNT_NAME = System.getProperty("snowflake.test.account.name"); protected static final String DATABASE = System.getProperty("snowflake.test.database"); protected static final String SCHEMA = System.getProperty("snowflake.test.schema"); diff --git a/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java b/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java index 9df01e4..2394d9b 100644 --- a/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java +++ b/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java @@ -18,6 +18,7 @@ import io.cdap.plugin.snowflake.Constants; import io.cdap.plugin.snowflake.common.BaseSnowflakeTest; +import io.cdap.plugin.snowflake.source.batch.SnowflakeInputFormatProvider; import io.cdap.plugin.snowflake.source.batch.SnowflakeSourceAccessor; import org.junit.Assert; import org.junit.Test; @@ -44,7 +45,8 @@ */ public class SnowflakeAccessorTest extends BaseSnowflakeTest { - private SnowflakeSourceAccessor snowflakeAccessor = new SnowflakeSourceAccessor(CONFIG); + private SnowflakeSourceAccessor snowflakeAccessor = + new SnowflakeSourceAccessor(CONFIG, SnowflakeInputFormatProvider.PROPERTY_DEFAULT_ESCAPE_CHAR); @Test public void testDescribeQuery() throws Exception { @@ -88,6 +90,43 @@ public void testDescribeQuery() throws Exception { Assert.assertEquals(expected, actual); } + @Test + public void testDescribeTable() throws Exception { + String schemaName = "TEST_SCHEMA"; + String tableName = "TEST_TABLE"; + List expected = Arrays.asList( + new SnowflakeFieldDescriptor("COLUMN_NUMBER", -5, true), + new SnowflakeFieldDescriptor("COLUMN_DECIMAL", -5, true), + new SnowflakeFieldDescriptor("COLUMN_NUMERIC", -5, true), + new SnowflakeFieldDescriptor("COLUMN_INT", -5, true), + new SnowflakeFieldDescriptor("COLUMN_INTEGER", -5, true), + new SnowflakeFieldDescriptor("COLUMN_BIGINT", -5, true), + new SnowflakeFieldDescriptor("COLUMN_SMALLINT", -5, true), + new SnowflakeFieldDescriptor("COLUMN_FLOAT", 8, true), + new SnowflakeFieldDescriptor("COLUMN_DOUBLE", 8, true), + new SnowflakeFieldDescriptor("COLUMN_REAL", 8, true), + new SnowflakeFieldDescriptor("COLUMN_VARCHAR", 12, true), + new SnowflakeFieldDescriptor("COLUMN_CHAR", 12, true), + new SnowflakeFieldDescriptor("COLUMN_TEXT", 12, true), + new SnowflakeFieldDescriptor("COLUMN_BINARY", -2, true), + new SnowflakeFieldDescriptor("COLUMN_BOOLEAN", 16, true), + new SnowflakeFieldDescriptor("COLUMN_DATE", 91, true), + new SnowflakeFieldDescriptor("COLUMN_TIMESTAMP", 93, true), + new SnowflakeFieldDescriptor("COLUMN_VARIANT", 12, true), + new SnowflakeFieldDescriptor("COLUMN_OBJECT", 12, true), + new SnowflakeFieldDescriptor("COLUMN_ARRAY", 12, true) + ); + + List actual = snowflakeAccessor.getFieldDescriptors( + String.valueOf(Constants.TEST_TABLE_SCHEMA), + Constants.TEST_TABLE); + + Assert.assertNotNull(actual); + Assert.assertFalse(actual.isEmpty()); + Assert.assertEquals(expected, actual); + } + + @Test public void testPrepareStageSplits() throws Exception { Pattern expected = Pattern.compile("cdap_stage/result.*data__0_0_0\\.csv\\.gz"); diff --git a/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java b/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java index d9d0740..c6b03b0 100644 --- a/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java +++ b/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java @@ -25,7 +25,9 @@ import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; + import java.io.IOException; +import java.sql.SQLException; import java.sql.Types; import java.util.ArrayList; import java.util.Arrays; @@ -47,7 +49,7 @@ public void testGetSchema() { ); MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); - Schema actual = SchemaHelper.getSchema(null, expected.toString(), collector, null); + Schema actual = SchemaHelper.getSchema(null, expected.toString(), collector, null, null); Assert.assertTrue(collector.getValidationFailures().isEmpty()); Assert.assertEquals(expected, actual); @@ -56,30 +58,34 @@ public void testGetSchema() { @Test public void testGetSchemaInvalidJson() { MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); - SchemaHelper.getSchema(null, "{}", collector, null); + SchemaHelper.getSchema(null, "{}", collector, null, null); ValidationAssertions.assertValidationFailed( collector, Collections.singletonList(SnowflakeBatchSourceConfig.PROPERTY_SCHEMA)); } @Test - public void testGetSchemaFromSnowflakeUnknownType() throws IOException { + public void testGetSchemaFromSnowflakeUnknownType() throws IOException, SQLException { + String importQuery = "SELECT * FROM someTable"; + String tableName = "USER"; MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); SnowflakeSourceAccessor snowflakeAccessor = Mockito.mock(SnowflakeSourceAccessor.class); List sample = new ArrayList<>(); sample.add(new SnowflakeFieldDescriptor("field1", -1000, false)); - Mockito.when(snowflakeAccessor.describeQuery(null)).thenReturn(sample); + Mockito.when(snowflakeAccessor.getFieldDescriptors(null, tableName)).thenReturn(sample); - SchemaHelper.getSchema(snowflakeAccessor, null, collector, null); + SchemaHelper.getSchema(snowflakeAccessor, null, collector, tableName, importQuery); ValidationAssertions.assertValidationFailed( collector, Collections.singletonList(SnowflakeBatchSourceConfig.PROPERTY_SCHEMA)); } @Test - public void testGetSchemaFromSnowflake() throws IOException { + public void testGetSchemaFromSnowflake() throws IOException, SQLException { + String importQuery = "SELECT * FROM someTable"; + String tableName = "USER"; MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); SnowflakeSourceAccessor snowflakeAccessor = Mockito.mock(SnowflakeSourceAccessor.class); @@ -107,7 +113,9 @@ public void testGetSchemaFromSnowflake() throws IOException { new SnowflakeFieldDescriptor("field121", Types.BIGINT, true), new SnowflakeFieldDescriptor("field122", Types.BIGINT, false), new SnowflakeFieldDescriptor("field131", Types.SMALLINT, true), - new SnowflakeFieldDescriptor("field132", Types.SMALLINT, false)); + new SnowflakeFieldDescriptor("field132", Types.SMALLINT, false), + new SnowflakeFieldDescriptor("field133", Types.TIMESTAMP_WITH_TIMEZONE, false), + new SnowflakeFieldDescriptor("field134", Types.TIMESTAMP_WITH_TIMEZONE, true)); Schema expected = Schema.recordOf( "data", @@ -134,12 +142,65 @@ public void testGetSchemaFromSnowflake() throws IOException { Schema.Field.of("field121", Schema.nullableOf(Schema.decimalOf(38))), Schema.Field.of("field122", Schema.decimalOf(38)), Schema.Field.of("field131", Schema.nullableOf(Schema.decimalOf(38))), - Schema.Field.of("field132", Schema.decimalOf(38)) + Schema.Field.of("field132", Schema.decimalOf(38)), + Schema.Field.of("field133", Schema.of(Schema.LogicalType.TIMESTAMP_MICROS)), + Schema.Field.of("field134", Schema.nullableOf(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS))) + ); + + Mockito.when(snowflakeAccessor.describeQuery(importQuery)).thenReturn(sample); + Mockito.when(snowflakeAccessor.getFieldDescriptors(Mockito.any(), Mockito.eq(tableName))).thenReturn(sample); + + Schema actual = SchemaHelper.getSchema(snowflakeAccessor, null, collector, tableName, importQuery); + + Assert.assertTrue(collector.getValidationFailures().isEmpty()); + Assert.assertEquals(expected, actual); + } + + @Test + public void testGetSchemaWhenMacroIsEnabled() { + Schema expected = Schema.recordOf("test", + Schema.Field.of("test_field", Schema.nullableOf(Schema.of(Schema.Type.LONG))) ); - Mockito.when(snowflakeAccessor.describeQuery(null)).thenReturn(sample); + SnowflakeBatchSourceConfig mockConfig = Mockito.mock(SnowflakeBatchSourceConfig.class); + Mockito.when(mockConfig.canConnect()).thenReturn(false); + Mockito.when(mockConfig.getSchema()).thenReturn(expected.toString()); + + MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); + Schema actual = SchemaHelper.getSchema(mockConfig, collector); + + Assert.assertTrue(collector.getValidationFailures().isEmpty()); + Assert.assertEquals(expected, actual); - Schema actual = SchemaHelper.getSchema(snowflakeAccessor, null, collector, null); + } + + @Test + public void testGetSchemaWhenMacroIsEnabledSchemaIsNull() { + + SnowflakeBatchSourceConfig mockConfig = Mockito.mock(SnowflakeBatchSourceConfig.class); + Mockito.when(mockConfig.canConnect()).thenReturn(false); + Mockito.when(mockConfig.getSchema()).thenReturn(null); + + MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); + Schema actual = SchemaHelper.getSchema(mockConfig, collector); + + Assert.assertTrue(collector.getValidationFailures().isEmpty()); + Assert.assertNull(actual); + + } + + @Test + public void testGetSchemaManuallyUpdatedTheSchema() { + Schema expected = Schema.recordOf("test", + Schema.Field.of("test_field", Schema.nullableOf(Schema.of(Schema.Type.LONG))) + ); + + SnowflakeBatchSourceConfig mockConfig = Mockito.mock(SnowflakeBatchSourceConfig.class); + Mockito.when(mockConfig.canConnect()).thenReturn(false); + Mockito.when(mockConfig.getSchema()).thenReturn(expected.toString()); + + MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); + Schema actual = SchemaHelper.getSchema(mockConfig, collector); Assert.assertTrue(collector.getValidationFailures().isEmpty()); Assert.assertEquals(expected, actual); diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java index 7f2f035..febcd3f 100644 --- a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java +++ b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java @@ -38,7 +38,9 @@ public class SnowflakeBatchSourceConfigBuilder { "", 0L, "", - ""); + "", + "", + ImportQueryType.IMPORT_QUERY.name()); private String referenceName; private String accountName; @@ -57,6 +59,8 @@ public class SnowflakeBatchSourceConfigBuilder { private Long maxSplitSize; private String connectionArguments; private String schema; + private String tableName; + private String importQueryType; public SnowflakeBatchSourceConfigBuilder() { } @@ -79,6 +83,8 @@ public SnowflakeBatchSourceConfigBuilder(SnowflakeBatchSourceConfig config) { this.maxSplitSize = config.getMaxSplitSize(); this.connectionArguments = config.getConnectionArguments(); this.schema = config.getSchema(); + this.tableName = config.getTableName(); + this.importQueryType = config.getImportQueryType(); } public SnowflakeBatchSourceConfigBuilder setReferenceName(String referenceName) { @@ -166,6 +172,16 @@ public SnowflakeBatchSourceConfigBuilder setSchema(String schema) { return this; } + public SnowflakeBatchSourceConfigBuilder setTableName(String tableName) { + this.tableName = tableName; + return this; + } + + public SnowflakeBatchSourceConfigBuilder setImportQueryType(String importQueryType) { + this.importQueryType = importQueryType; + return this; + } + public SnowflakeBatchSourceConfig build() { return new SnowflakeBatchSourceConfig(referenceName, accountName, @@ -183,6 +199,8 @@ public SnowflakeBatchSourceConfig build() { refreshToken, maxSplitSize, connectionArguments, - schema); + schema, + tableName, + importQueryType); } } diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java index a071bea..82c1806 100644 --- a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java +++ b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java @@ -20,6 +20,7 @@ import io.cdap.plugin.snowflake.ValidationAssertions; import org.junit.Assert; import org.junit.Test; + import java.util.Collections; /** @@ -68,4 +69,19 @@ public void validatePassword() { ValidationAssertions.assertValidationFailed( collector, Collections.singletonList(SnowflakeBatchSourceConfig.PROPERTY_PASSWORD)); } + + @Test + public void validateTableNameAndImportQuery() { + SnowflakeBatchSourceConfig config = + new SnowflakeBatchSourceConfigBuilder(SnowflakeBatchSourceConfigBuilder.CONFIG) + .setTableName(null) + .setImportQuery(null) + .build(); + + MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE); + config.validate(collector); + + Assert.assertFalse("Both table name and import query cannot be null", + collector.getValidationFailures().isEmpty()); + } } diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformerTest.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformerTest.java index c68e4ea..f81799d 100644 --- a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformerTest.java +++ b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeMapToRecordTransformerTest.java @@ -50,8 +50,8 @@ public void transform() { row.put("COLUMN_CHARACTER", "2"); row.put("COLUMN_STRING", "text_115"); row.put("COLUMN_TEXT", "text_116"); - row.put("COLUMN_BINARY", "text_117"); - row.put("COLUMN_VARBINARY", "text_118"); + row.put("COLUMN_BINARY", "746578745f313137"); + row.put("COLUMN_VARBINARY", "746578745f313138"); row.put("COLUMN_BOOLEAN", "true"); row.put("COLUMN_DATE", "2019-01-01"); row.put("COLUMN_DATETIME", "2019-01-01T01:01:01+00:00"); diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeS3StageEndpointTest.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeS3StageEndpointTest.java new file mode 100644 index 0000000..05b1a52 --- /dev/null +++ b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeS3StageEndpointTest.java @@ -0,0 +1,78 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); you may not + * use this file except in compliance with the License. You may obtain a copy of + * the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT + * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the + * License for the specific language governing permissions and limitations under + * the License. + */ + +package io.cdap.plugin.snowflake.source.batch; + +import net.snowflake.client.internal.jdbc.cloud.storage.SnowflakeStorageClient; +import net.snowflake.client.internal.jdbc.cloud.storage.StageInfo; +import net.snowflake.client.internal.jdbc.cloud.storage.StorageClientFactory; +import org.junit.Assert; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.net.URI; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +/** + * Tests that the S3 client used by the driver's downloadStream accepts the stage endpoint returned by Snowflake. + * Government accounts return an endpoint without a URI scheme, which JDBC 4.0.2 rejected (snowflake-jdbc#2572). + */ +public class SnowflakeS3StageEndpointTest { + + private static final String GOV_REGION = "us-gov-west-1"; + private static final String GOV_ENDPOINT = "s3-fips.us-gov-west-1.amazonaws.com"; + + @Test + public void testSchemelessStageEndpointIsNormalized() throws Exception { + Assert.assertEquals(Optional.of(URI.create("https://" + GOV_ENDPOINT)), getEndpointOverride(GOV_ENDPOINT)); + } + + @Test + public void testHttpsStageEndpointIsUnchanged() throws Exception { + Assert.assertEquals(Optional.of(URI.create("https://" + GOV_ENDPOINT)), + getEndpointOverride("https://" + GOV_ENDPOINT)); + } + + @Test + public void testNoStageEndpoint() throws Exception { + Assert.assertEquals(Optional.empty(), getEndpointOverride(null)); + } + + /** + * Creates the S3 client the same way the driver's downloadStream does and returns its endpoint override. + */ + private static Object getEndpointOverride(String stageEndPoint) throws Exception { + Map credentials = new HashMap<>(); + credentials.put("AWS_KEY_ID", "test-key-id"); + credentials.put("AWS_SECRET_KEY", "test-secret"); + credentials.put("AWS_TOKEN", "test-token"); + StageInfo stageInfo = StageInfo.createStageInfo("S3", "bucket/stage/", credentials, GOV_REGION, + stageEndPoint, null, true); + SnowflakeStorageClient client = StorageClientFactory.getFactory().createClient(stageInfo, 1, null, null); + + // The AWS SDK is shaded inside the driver, so the endpoint is read reflectively. + Field amazonClientField = client.getClass().getDeclaredField("amazonClient"); + amazonClientField.setAccessible(true); + Object amazonClient = amazonClientField.get(client); + Object clientConfig = amazonClient.getClass().getMethod("serviceClientConfiguration").invoke(amazonClient); + Method endpointOverride = clientConfig.getClass().getMethod("endpointOverride"); + endpointOverride.setAccessible(true); + return endpointOverride.invoke(clientConfig); + } +} diff --git a/widgets/Snowflake-batchsource.json b/widgets/Snowflake-batchsource.json index 18a670e..1e202ff 100644 --- a/widgets/Snowflake-batchsource.json +++ b/widgets/Snowflake-batchsource.json @@ -36,6 +36,40 @@ "label": "Role", "name": "role" }, + { + "name": "importQueryType", + "label": "Import Query Type", + "widget-type": "radio-group", + "widget-attributes": { + "layout": "inline", + "default": "importQuery", + "options": [ + { + "id": "importQuery", + "label": "Native Query" + }, + { + "id": "tableName", + "label": "Named Table" + } + ] + } + }, + { + "widget-type": "textbox", + "label": "Table Name", + "name": "tableName", + "plugin-function": { + "label": "Get Schema", + "widget": "outputSchema", + "output-property": "schema", + "omit-properties": [ + { + "name": "schema" + } + ] + } + }, { "widget-type": "textarea", "label": "Import Query", @@ -200,6 +234,30 @@ } ] }, + { + "name": "ImportQuery", + "condition": { + "expression": "importQueryType != 'tableName'" + }, + "show": [ + { + "type": "property", + "name": "importQuery" + } + ] + }, + { + "name": "NativeTableName", + "condition": { + "expression": "importQueryType == 'tableName'" + }, + "show": [ + { + "type": "property", + "name": "tableName" + } + ] + }, { "name": "OAuth2EnabledFilter", "condition": {