consumer) {
- Bucket sourceBucket = storage.get(sourcePath.getBucket());
+ Bucket sourceBucket = null;
+ try {
+ sourceBucket = storage.get(sourcePath.getBucket());
+ } catch (StorageException e) {
+ // Add more descriptive error message
+ throw new RuntimeException(
+ String.format("Unable to access source bucket %s. ", sourcePath.getBucket())
+ + "Ensure you entered the correct bucket path.", e);
+ }
if (sourceBucket == null) {
throw new IllegalArgumentException(
String.format("Source bucket '%s' does not exist.", sourcePath.getBucket()));
}
- Bucket destBucket = storage.get(destPath.getBucket());
+ Bucket destBucket = null;
+ try {
+ destBucket = storage.get(destPath.getBucket());
+ } catch (StorageException e) {
+ // Add more descriptive error message
+ throw new RuntimeException(
+ String.format("Unable to access destination bucket %s. ", destPath.getBucket())
+ + "Ensure you entered the correct bucket path.", e);
+ }
if (destBucket == null) {
throw new IllegalArgumentException(
String.format("Destination bucket '%s' does not exist. Please create it first.", destPath.getBucket()));
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetter.java b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetter.java
new file mode 100644
index 0000000000..0f8f0bf500
--- /dev/null
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetter.java
@@ -0,0 +1,246 @@
+/*
+ * Copyright © 2020 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.gcp.gcs.actions;
+
+import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageOptions;
+import com.google.common.base.Joiner;
+import com.google.gson.GsonBuilder;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import com.google.gson.JsonSyntaxException;
+import io.cdap.cdap.api.annotation.Description;
+import io.cdap.cdap.api.annotation.Name;
+import io.cdap.cdap.api.annotation.Plugin;
+import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.cdap.etl.api.FailureCollector;
+import io.cdap.cdap.etl.api.PipelineConfigurer;
+import io.cdap.cdap.etl.api.StageConfigurer;
+import io.cdap.cdap.etl.api.action.Action;
+import io.cdap.cdap.etl.api.action.ActionContext;
+import io.cdap.plugin.gcp.common.GCPUtils;
+import io.cdap.plugin.gcp.gcs.GCSPath;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * This class GCSArgumentSetter get json file configuration from GCS.
+ *
+ * The plugin provides the ability to map json properties as pipeline arguments name and columns
+ * values as pipeline arguments
+ * {
+ * "arguments" : [
+ * { "name" : "input.path", "type" : "string", "value" : "/data/sunny_feeds/master"},
+ * { "name" : "parse.schema, "type" : "schema", "value" : [
+ * { "name" : "fname", "type" : "string", "nullable" : true },
+ * { "name" : "age", "type" : "int", "nullable" : true},
+ * { "name" : "salary", "type" : "float", "nullable" : false}
+ * ]},
+ * { "name" : "directives", "type" : "array", "value" : [
+ * "parse-as-json body",
+ * "columns-replace s/body_//g",
+ * "keep f1,f2"
+ * ]}
+ * ]
+ * }
+ *
+ * In case when `nullable` is not specified inside field json, then it defaults to true.
+ *
+ */
+@Plugin(type = Action.PLUGIN_TYPE)
+@Name(GCSArgumentSetter.NAME)
+@Description("Argument setter for dynamically configuring pipeline from GCS")
+public final class GCSArgumentSetter extends Action {
+
+ public static final String NAME = "GCSArgumentSetter";
+ private GCSArgumentSetterConfig config;
+
+ @Override
+ public void configurePipeline(PipelineConfigurer configurer) {
+ super.configurePipeline(configurer);
+ StageConfigurer stageConfigurer = configurer.getStageConfigurer();
+ FailureCollector collector = stageConfigurer.getFailureCollector();
+ config.validate(collector);
+ }
+
+ @Override
+ public void run(ActionContext context) throws Exception {
+ config.validateProperties(context.getFailureCollector());
+ String fileContent = GCSArgumentSetter.getContent(config);
+
+ try {
+ Configuration configuration =
+ new GsonBuilder().create().fromJson(fileContent, Configuration.class);
+ for (Argument argument : configuration.getArguments()) {
+ String name = argument.getName();
+ String value = argument.getValue();
+ if (value != null) {
+ context.getArguments().set(name, value);
+ } else {
+ throw new RuntimeException(
+ "Configuration '" + name + "' is null. Cannot set argument to null.");
+ }
+ }
+ } catch (JsonSyntaxException e) {
+ throw new RuntimeException(
+ String.format(
+ "Could not parse response from '%s': %s", config.getPath(), e.getMessage()));
+ }
+ }
+
+ private static Storage getStorage(GCSArgumentSetterConfig config) throws IOException {
+ return StorageOptions.newBuilder()
+ .setProjectId(config.getProject())
+ .setCredentials(getCredentials(config))
+ .build()
+ .getService();
+ }
+
+ public static ServiceAccountCredentials getCredentials(GCSArgumentSetterConfig config)
+ throws IOException {
+ ServiceAccountCredentials credentials = null;
+ switch (config.getServiceAccountType()) {
+ case FILE_PATH:
+ credentials = GCPUtils.loadServiceAccountCredentials(config.getServiceAccountFilePath());
+ break;
+ case JSON:
+ InputStream jsonInputStream =
+ new ByteArrayInputStream(config.getServiceAccountJSON().getBytes());
+ credentials = ServiceAccountCredentials.fromStream(jsonInputStream);
+ break;
+ }
+ return credentials;
+ }
+
+ public static String getContent(GCSArgumentSetterConfig config) throws IOException {
+ Storage storage = getStorage(config);
+ GCSPath path = config.getPath();
+ Blob blob = storage.get(path.getBucket(), path.getName());
+ return new String(blob.getContent());
+ }
+
+ private static final class Configuration {
+
+ private List arguments;
+
+ public List getArguments() {
+ return arguments;
+ }
+ }
+
+ private static final class Argument {
+
+ private String name;
+ private String type;
+ private JsonElement value;
+
+ Argument() {
+ type = "string";
+ }
+
+ public String getName() {
+ return name;
+ }
+
+ public void setName(String name) {
+ this.name = name;
+ }
+
+ public String getValue() {
+ if (value == null) {
+ throw new IllegalArgumentException("Null Argument value for name '" + name + "'");
+ }
+ if (type.equalsIgnoreCase("schema")) {
+ return createSchema(value).toString();
+ } else if (type.equalsIgnoreCase("int")) {
+ return Integer.toString(value.getAsInt());
+ } else if (type.equalsIgnoreCase("float")) {
+ return Float.toString(value.getAsFloat());
+ } else if (type.equalsIgnoreCase("double")) {
+ return Double.toString(value.getAsDouble());
+ } else if (type.equalsIgnoreCase("short")) {
+ return Short.toString(value.getAsShort());
+ } else if (type.equalsIgnoreCase("string")) {
+ return value.getAsString();
+ } else if (type.equalsIgnoreCase("char")) {
+ return Character.toString(value.getAsCharacter());
+ } else if (type.equalsIgnoreCase("array")) {
+ List values = new ArrayList<>();
+ for (JsonElement v : value.getAsJsonArray()) {
+ values.add(v.getAsString());
+ }
+ return Joiner.on("\n").join(values);
+ } else if (type.equalsIgnoreCase("map")) {
+ List values = new ArrayList<>();
+ for (Map.Entry entry : value.getAsJsonObject().entrySet()) {
+ values.add(String.format("%s=%s", entry.getKey(), entry.getValue().getAsString()));
+ }
+ return Joiner.on(";").join(values);
+ }
+
+ throw new IllegalArgumentException("Invalid argument value '" + value + "'");
+ }
+
+ private Schema createSchema(JsonElement array) {
+ List fields = new ArrayList<>();
+ for (JsonElement field : array.getAsJsonArray()) {
+ Schema.Type type = Schema.Type.STRING;
+
+ JsonObject object = field.getAsJsonObject();
+ String fieldType = object.get("type").getAsString().toLowerCase();
+
+ boolean isNullable = true;
+ if (object.get("nullable") != null) {
+ isNullable = object.get("nullable").getAsBoolean();
+ }
+
+ String name = object.get("name").getAsString();
+
+ if (fieldType.equals("double")) {
+ type = Schema.Type.DOUBLE;
+ } else if (fieldType.equals("float")) {
+ type = Schema.Type.FLOAT;
+ } else if (fieldType.equals("long")) {
+ type = Schema.Type.LONG;
+ } else if (fieldType.equals("int")) {
+ type = Schema.Type.INT;
+ } else if (fieldType.equals("short")) {
+ type = Schema.Type.INT;
+ } else if (fieldType.equals("string")) {
+ type = Schema.Type.STRING;
+ }
+
+ Schema schema;
+ if (isNullable) {
+ schema = Schema.nullableOf(Schema.of(type));
+ } else {
+ schema = Schema.of(type);
+ }
+ Schema.Field fld = Schema.Field.of(name, schema);
+ fields.add(fld);
+ }
+
+ return Schema.recordOf("record", fields);
+ }
+ }
+}
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetterConfig.java b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetterConfig.java
new file mode 100644
index 0000000000..573b9c75bf
--- /dev/null
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSArgumentSetterConfig.java
@@ -0,0 +1,125 @@
+/*
+ * Copyright © 2020 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.gcp.gcs.actions;
+
+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.gcp.common.GCPConfig;
+import io.cdap.plugin.gcp.gcs.GCSPath;
+import javax.annotation.Nullable;
+
+/** Holds configuration required for configuring {@link GCSArgumentSetter}. */
+public final class GCSArgumentSetterConfig extends GCPConfig {
+
+ public static final String NAME_PATH = "path";
+ public static final String NAME_SERVICE_ACCOUNT_TYPE = "serviceAccountType";
+ public static final String NAME_SERVICE_ACCOUNT_JSON = "serviceAccountJSON";
+
+ @Name(NAME_PATH)
+ @Macro
+ @Description("GCS Path to the file containing the arguments")
+ private String path;
+
+ @Name(NAME_SERVICE_ACCOUNT_TYPE)
+ @Macro
+ @Nullable
+ @Description(
+ "Provide service account as JSON. When it is set to 'Yes', "
+ + "the content of service account key needs to be copied, whereas when it is set to 'No' "
+ + "the service Account file path needs to be specified. The default value is 'No'")
+ private String serviceAccountType;
+
+ @Name(NAME_SERVICE_ACCOUNT_JSON)
+ @Macro
+ @Nullable
+ @Description("The content of the service account.")
+ private String serviceAccountJSON;
+
+ public void validate(FailureCollector collector) {
+ validateProperties(collector);
+
+ if (canConnect()) {
+ try {
+ GCSArgumentSetter.getContent(this);
+ } catch (Exception e) {
+ collector.addFailure("Can not get content from GCP!", null);
+ }
+ }
+ collector.getOrThrowException();
+ }
+
+ public void validateProperties(FailureCollector collector) {
+ if (!containsMacro(NAME_PATH)) {
+ try {
+ getPath();
+ } catch (IllegalArgumentException e) {
+ collector.addFailure(e.getMessage(), null).withConfigProperty(NAME_PATH);
+ }
+ }
+
+ if (getServiceAccountType() == ServiceAccountType.JSON
+ && !containsMacro(NAME_SERVICE_ACCOUNT_JSON)
+ && Strings.isNullOrEmpty(getServiceAccountJSON())) {
+ collector
+ .addFailure("Required property 'Service Account JSON' has no value.", "")
+ .withConfigProperty(NAME_SERVICE_ACCOUNT_JSON);
+ }
+ }
+
+ private boolean canConnect() {
+ boolean canConnect =
+ !containsMacro(NAME_PATH)
+ && !(containsMacro(NAME_PROJECT) || AUTO_DETECT.equals(project))
+ && !(containsMacro(NAME_SERVICE_ACCOUNT_TYPE));
+
+ if (!canConnect) {
+ return false;
+ }
+
+ ServiceAccountType serviceAccountType = getServiceAccountType();
+
+ if (serviceAccountType == ServiceAccountType.FILE_PATH) {
+ return !containsMacro(NAME_SERVICE_ACCOUNT_FILE_PATH)
+ && !Strings.isNullOrEmpty(getServiceAccountFilePath());
+ }
+ return !containsMacro(NAME_SERVICE_ACCOUNT_JSON)
+ && !Strings.isNullOrEmpty(getServiceAccountJSON());
+ }
+
+ public GCSPath getPath() {
+ return GCSPath.from(path);
+ }
+
+ public ServiceAccountType getServiceAccountType() {
+ return "JSON".equalsIgnoreCase(serviceAccountType)
+ ? ServiceAccountType.JSON
+ : ServiceAccountType.FILE_PATH;
+ }
+
+ public String getServiceAccountJSON() {
+ return serviceAccountJSON;
+ }
+
+ /** The type of service account. */
+ public enum ServiceAccountType {
+ FILE_PATH,
+ JSON;
+ }
+}
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketCreate.java b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketCreate.java
index b4894db306..d645544400 100644
--- a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketCreate.java
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketCreate.java
@@ -1,5 +1,5 @@
/*
- * Copyright © 2015 Cask Data, Inc.
+ * Copyright © 2015-2020 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
@@ -17,7 +17,9 @@
package io.cdap.plugin.gcp.gcs.actions;
import com.google.auth.Credentials;
+import com.google.cloud.storage.Bucket;
import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
import io.cdap.cdap.api.annotation.Description;
import io.cdap.cdap.api.annotation.Macro;
import io.cdap.cdap.api.annotation.Name;
@@ -100,7 +102,16 @@ public void run(ActionContext context) throws Exception {
}
// create the gcs buckets if not exist
- if (storage.get(gcsPath.getBucket()) == null) {
+ Bucket bucket = null;
+ try {
+ bucket = storage.get(gcsPath.getBucket());
+ } catch (StorageException e) {
+ // Add more descriptive error message
+ throw new RuntimeException(
+ String.format("Unable to access or create bucket %s. ", gcsPath.getBucket())
+ + "Ensure you entered the correct bucket path and have permissions for it.", e);
+ }
+ if (bucket == null) {
GCPUtils.createBucket(storage, gcsPath.getBucket(), config.location,
context.getArguments().get(GCPUtils.CMEK_KEY));
undoBucket.add(bucketPath);
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketDelete.java b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketDelete.java
index 412b983b19..72d51f0421 100644
--- a/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketDelete.java
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/actions/GCSBucketDelete.java
@@ -1,5 +1,5 @@
/*
- * Copyright © 2015 Cask Data, Inc.
+ * Copyright © 2015-2020 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
@@ -16,6 +16,9 @@
package io.cdap.plugin.gcp.gcs.actions;
+import com.google.auth.Credentials;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
import io.cdap.cdap.api.annotation.Description;
import io.cdap.cdap.api.annotation.Macro;
import io.cdap.cdap.api.annotation.Name;
@@ -25,6 +28,7 @@
import io.cdap.cdap.etl.api.action.Action;
import io.cdap.cdap.etl.api.action.ActionContext;
import io.cdap.plugin.gcp.common.GCPConfig;
+import io.cdap.plugin.gcp.common.GCPUtils;
import io.cdap.plugin.gcp.gcs.GCSPath;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
@@ -61,6 +65,8 @@ public void run(ActionContext context) throws Exception {
Configuration configuration = new Configuration();
String serviceAccountFilePath = config.getServiceAccountFilePath();
+ Credentials credentials = serviceAccountFilePath == null ?
+ null : GCPUtils.loadServiceAccountCredentials(serviceAccountFilePath);
if (serviceAccountFilePath != null) {
configuration.set("google.cloud.auth.service.account.json.keyfile", serviceAccountFilePath);
}
@@ -74,8 +80,19 @@ public void run(ActionContext context) throws Exception {
configuration.setBoolean("fs.gs.impl.disable.cache", true);
List gcsPaths = new ArrayList<>();
+ Storage storage = GCPUtils.getStorage(config.getProject(), credentials);
for (String path : config.getPaths()) {
- gcsPaths.add(new Path(GCSPath.from(path).getUri()));
+ GCSPath gcsPath = GCSPath.from(path);
+ // Check if the bucket is accessible
+ try {
+ storage.get(gcsPath.getBucket());
+ } catch (StorageException e) {
+ // Add more descriptive error message
+ throw new RuntimeException(
+ String.format("Unable to access or create bucket %s. ", gcsPath.getBucket())
+ + "Ensure you entered the correct bucket path and have permissions for it.", e);
+ }
+ gcsPaths.add(new Path(gcsPath.getUri()));
}
FileSystem fs;
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSink.java b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSink.java
index b194ef272b..f11b3ddec7 100644
--- a/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSink.java
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSink.java
@@ -1,5 +1,5 @@
/*
- * Copyright © 2015-2016 Cask Data, Inc.
+ * Copyright © 2015-2020 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
@@ -17,8 +17,13 @@
package io.cdap.plugin.gcp.gcs.sink;
import com.google.auth.Credentials;
+import com.google.cloud.storage.Bucket;
import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
import com.google.common.base.Strings;
+import com.google.gson.Gson;
+import com.google.gson.JsonSyntaxException;
+import com.google.gson.reflect.TypeToken;
import io.cdap.cdap.api.annotation.Description;
import io.cdap.cdap.api.annotation.Macro;
import io.cdap.cdap.api.annotation.Name;
@@ -26,8 +31,10 @@
import io.cdap.cdap.api.data.schema.Schema;
import io.cdap.cdap.etl.api.FailureCollector;
import io.cdap.cdap.etl.api.PipelineConfigurer;
+import io.cdap.cdap.etl.api.StageMetrics;
import io.cdap.cdap.etl.api.batch.BatchSink;
import io.cdap.cdap.etl.api.batch.BatchSinkContext;
+import io.cdap.cdap.etl.api.validation.ValidatingOutputFormat;
import io.cdap.plugin.common.LineageRecorder;
import io.cdap.plugin.format.FileFormat;
import io.cdap.plugin.format.plugin.AbstractFileSink;
@@ -35,9 +42,13 @@
import io.cdap.plugin.gcp.common.GCPReferenceSinkConfig;
import io.cdap.plugin.gcp.common.GCPUtils;
import io.cdap.plugin.gcp.gcs.GCSPath;
+import io.cdap.plugin.gcp.gcs.StorageClient;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.text.SimpleDateFormat;
+import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -50,7 +61,16 @@
@Name("GCS")
@Description("Writes records to one or more files in a directory on Google Cloud Storage.")
public class GCSBatchSink extends AbstractFileSink {
+
+ private static final Logger LOG = LoggerFactory.getLogger(GCSBatchSink.class);
+ public static final String RECORD_COUNT = "recordcount";
+ private static final String RECORDS_UPDATED_METRIC = "records.updated";
+ public static final String AVRO_NAMED_OUTPUT = "avro.mo.config.namedOutput";
+ public static final String COMMON_NAMED_OUTPUT = "mapreduce.output.basename";
+
private final GCSBatchSinkConfig config;
+ private String outputPath;
+
public GCSBatchSink(GCSBatchSinkConfig config) {
super(config);
@@ -62,6 +82,18 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
super.configurePipeline(pipelineConfigurer);
}
+ @Override
+ public ValidatingOutputFormat getValidatingOutputFormat(PipelineConfigurer pipelineConfigurer) {
+ ValidatingOutputFormat delegate = super.getValidatingOutputFormat(pipelineConfigurer);
+ return new GCSOutputFormatProvider(delegate);
+ }
+
+ @Override
+ public ValidatingOutputFormat getOutputFormatForRun(BatchSinkContext context) throws InstantiationException {
+ ValidatingOutputFormat outputFormatForRun = super.getOutputFormatForRun(context);
+ return new GCSOutputFormatProvider(outputFormatForRun);
+ }
+
@Override
public void prepareRun(BatchSinkContext context) throws Exception {
super.prepareRun(context);
@@ -69,14 +101,37 @@ public void prepareRun(BatchSinkContext context) throws Exception {
Credentials credentials = config.getServiceAccountFilePath() == null ?
null : GCPUtils.loadServiceAccountCredentials(config.getServiceAccountFilePath());
Storage storage = GCPUtils.getStorage(config.getProject(), credentials);
- if (storage.get(config.getBucket()) == null) {
+ Bucket bucket;
+ try {
+ bucket = storage.get(config.getBucket());
+ } catch (StorageException e) {
+ throw new RuntimeException(
+ String.format("Unable to access or create bucket %s. ", config.getBucket())
+ + "Ensure you entered the correct bucket path and have permissions for it.", e);
+ }
+ if (bucket == null) {
GCPUtils.createBucket(storage, config.getBucket(), config.getLocation(), cmekKey);
}
}
@Override
protected Map getFileSystemProperties(BatchSinkContext context) {
- return GCPUtils.getFileSystemProperties(config, config.getPath(), new HashMap<>());
+ Map properties = GCPUtils.getFileSystemProperties(config, config.getPath(), new HashMap<>());
+ properties.putAll(config.getFileSystemProperties());
+ String outputFileBaseName = config.getOutputFileNameBase();
+ if (outputFileBaseName == null || outputFileBaseName.isEmpty()) {
+ return properties;
+ }
+
+ properties.put(AVRO_NAMED_OUTPUT, outputFileBaseName);
+ properties.put(COMMON_NAMED_OUTPUT, outputFileBaseName);
+ return properties;
+ }
+
+ @Override
+ protected String getOutputDir(long logicalStartTime) {
+ this.outputPath = super.getOutputDir(logicalStartTime);
+ return this.outputPath;
}
@Override
@@ -84,6 +139,88 @@ protected void recordLineage(LineageRecorder lineageRecorder, List outpu
lineageRecorder.recordWrite("Write", "Wrote to Google Cloud Storage.", outputFields);
}
+ @Override
+ public void onRunFinish(boolean succeeded, BatchSinkContext context) {
+ super.onRunFinish(succeeded, context);
+ emitMetrics(succeeded, context);
+ }
+
+ private void emitMetrics(boolean succeeded, BatchSinkContext context) {
+ if (!succeeded) {
+ return;
+ }
+
+ try {
+ StorageClient storageClient = StorageClient.create(config.getProject(), config.getServiceAccountFilePath());
+ storageClient.mapMetaDataForAllBlobs(getPrefixPath(), new MetricsEmitter(context.getMetrics())::emitMetrics);
+ } catch (Exception e) {
+ LOG.warn("Metrics for the number of affected rows in GCS Sink maybe incorrect.", e);
+ }
+ }
+
+ private String getPrefixPath() {
+ String filenameBase = getFilenameBase();
+ if (filenameBase == null) {
+ return outputPath;
+ }
+ //Following code is for handling a special case for saving files in same output directory.
+ //The specified file prefix from Filesystem Properties/Output file base name can be used to make filename unique.
+ //The code is based on assumptions from the internal implementation of
+ //org.apache.avro.mapreduce.AvroOutputFormatBase and org.apache.hadoop.mapreduce.lib.output.FileOutputFormat
+ String outputPathPrefix = outputPath.endsWith("/") ? outputPath.substring(0, outputPath.length() - 1) : outputPath;
+ return String.format("%s/%s-", outputPathPrefix, filenameBase);
+ }
+
+ @Nullable
+ private String getFilenameBase() {
+ String outputFileBaseName = config.getOutputFileNameBase();
+ if (outputFileBaseName != null && !outputFileBaseName.isEmpty()) {
+ return outputFileBaseName;
+ }
+
+ Map fileSystemProperties = config.getFileSystemProperties();
+ if (fileSystemProperties.containsKey(AVRO_NAMED_OUTPUT) && config.getFormat() == FileFormat.AVRO) {
+ return fileSystemProperties.get(AVRO_NAMED_OUTPUT);
+ }
+ if (fileSystemProperties.containsKey(COMMON_NAMED_OUTPUT)) {
+ return fileSystemProperties.get(COMMON_NAMED_OUTPUT);
+ }
+ return null;
+ }
+
+ private static class MetricsEmitter {
+ private StageMetrics stageMetrics;
+
+ private MetricsEmitter(StageMetrics stageMetrics) {
+ this.stageMetrics = stageMetrics;
+ }
+
+ public void emitMetrics(Map metaData) {
+ long totalRows = extractRecordCount(metaData);
+ if (totalRows == 0) {
+ return;
+ }
+
+ //work around since StageMetrics count() only takes int as of now
+ int cap = 10000; // so the loop will not cause significant delays
+ long count = totalRows / Integer.MAX_VALUE;
+ if (count > cap) {
+ LOG.warn("Total record count is too high! Metric for the number of affected rows may not be updated correctly");
+ }
+ count = count < cap ? count : cap;
+ for (int i = 0; i <= count && totalRows > 0; i++) {
+ int rowCount = totalRows < Integer.MAX_VALUE ? (int) totalRows : Integer.MAX_VALUE;
+ stageMetrics.count(RECORDS_UPDATED_METRIC, rowCount);
+ totalRows -= rowCount;
+ }
+ }
+
+ private long extractRecordCount(Map metadata) {
+ String value = metadata.get(RECORD_COUNT);
+ return value == null ? 0L : Long.parseLong(value);
+ }
+ }
+
/**
* Sink configuration.
*/
@@ -94,6 +231,8 @@ public static class GCSBatchSinkConfig extends GCPReferenceSinkConfig implements
private static final String NAME_FORMAT = "format";
private static final String NAME_SCHEMA = "schema";
private static final String NAME_LOCATION = "location";
+ private static final String NAME_FS_PROPERTIES = "fileSystemProperties";
+ private static final String NAME_FILE_NAME_BASE = "outputFileNameBase";
private static final String SCHEME = "gs://";
@Name(NAME_PATH)
@@ -132,6 +271,18 @@ public static class GCSBatchSinkConfig extends GCPReferenceSinkConfig implements
"This value is ignored if the bucket already exists")
protected String location;
+ @Name(NAME_FS_PROPERTIES)
+ @Macro
+ @Nullable
+ @Description("Advanced feature to specify any additional properties that should be used with the sink.")
+ private String fileSystemProperties;
+
+ @Name(NAME_FILE_NAME_BASE)
+ @Macro
+ @Nullable
+ @Description("Advanced feature to specify file output name prefix.")
+ private String outputFileNameBase;
+
@Override
public void validate() {
// no-op
@@ -171,6 +322,13 @@ public void validate(FailureCollector collector) {
} catch (IllegalArgumentException e) {
collector.addFailure(e.getMessage(), null).withConfigProperty(NAME_SCHEMA).withStacktrace(e.getStackTrace());
}
+
+ try {
+ getFileSystemProperties();
+ } catch (IllegalArgumentException e) {
+ collector.addFailure("File system properties must be a valid json.", null)
+ .withConfigProperty(NAME_FS_PROPERTIES).withStacktrace(e.getStackTrace());
+ }
}
public String getBucket() {
@@ -215,5 +373,22 @@ public String getDelimiter() {
public String getLocation() {
return location;
}
+
+ public Map getFileSystemProperties() {
+ if (fileSystemProperties == null || fileSystemProperties.isEmpty()) {
+ return Collections.emptyMap();
+ }
+ try {
+ return new Gson().fromJson(fileSystemProperties, new TypeToken>() {
+ }.getType());
+ } catch (JsonSyntaxException e) {
+ throw new IllegalArgumentException("Unable to parse filesystem properties: " + e.getMessage(), e);
+ }
+ }
+
+ @Nullable
+ public String getOutputFileNameBase() {
+ return outputFileNameBase;
+ }
}
}
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSMultiBatchSink.java b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSMultiBatchSink.java
index 4421ed8604..966a846617 100644
--- a/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSMultiBatchSink.java
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSMultiBatchSink.java
@@ -1,5 +1,5 @@
/*
- * Copyright © 2015-2016 Cask Data, Inc.
+ * Copyright © 2015-2020 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
@@ -18,6 +18,7 @@
import com.google.auth.Credentials;
import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
import com.google.common.base.Strings;
import io.cdap.cdap.api.annotation.Description;
import io.cdap.cdap.api.annotation.Name;
@@ -100,8 +101,15 @@ public void prepareRun(BatchSinkContext context) throws IOException, Instantiati
Credentials credentials = config.getServiceAccountFilePath() == null ?
null : GCPUtils.loadServiceAccountCredentials(config.getServiceAccountFilePath());
Storage storage = GCPUtils.getStorage(config.getProject(), credentials);
- if (storage.get(config.getBucket()) == null) {
- GCPUtils.createBucket(storage, config.getBucket(), config.getLocation(), cmekKey);
+ try {
+ if (storage.get(config.getBucket()) == null) {
+ GCPUtils.createBucket(storage, config.getBucket(), config.getLocation(), cmekKey);
+ }
+ } catch (StorageException e) {
+ // Add more descriptive error message
+ throw new RuntimeException(
+ String.format("Unable to access or create bucket %s. ", config.getBucket())
+ + "Ensure you entered the correct bucket path and have permissions for it.", e);
}
for (Map.Entry argument : argumentCopy.entrySet()) {
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputFormatProvider.java b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputFormatProvider.java
new file mode 100644
index 0000000000..0bced920c2
--- /dev/null
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputFormatProvider.java
@@ -0,0 +1,250 @@
+package io.cdap.plugin.gcp.gcs.sink;
+
+import com.google.cloud.storage.Blob;
+import com.google.common.annotations.VisibleForTesting;
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.etl.api.validation.FormatContext;
+import io.cdap.cdap.etl.api.validation.ValidatingOutputFormat;
+import io.cdap.plugin.gcp.common.GCPUtils;
+import io.cdap.plugin.gcp.gcs.StorageClient;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.io.NullWritable;
+import org.apache.hadoop.mapreduce.JobContext;
+import org.apache.hadoop.mapreduce.JobStatus;
+import org.apache.hadoop.mapreduce.OutputCommitter;
+import org.apache.hadoop.mapreduce.OutputFormat;
+import org.apache.hadoop.mapreduce.RecordWriter;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
+import org.apache.hadoop.util.ReflectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * OutputFormatProvider for GCSSink
+ */
+public class GCSOutputFormatProvider implements ValidatingOutputFormat {
+
+ private static final Logger LOG = LoggerFactory.getLogger(GCSOutputFormatProvider.class);
+ private static final String DELEGATE_OUTPUTFORMAT_CLASSNAME = "gcssink.delegate.outputformat.classname";
+ private static final String OUTPUT_FOLDER = "gcssink.metric.output.folder";
+ public static final String RECORD_COUNT_FORMAT = "recordcount.%s";
+ private final ValidatingOutputFormat delegate;
+
+ public GCSOutputFormatProvider(ValidatingOutputFormat delegate) {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public void validate(FormatContext context) {
+ delegate.validate(context);
+ }
+
+ @Override
+ public String getOutputFormatClassName() {
+ return GCSOutputFormat.class.getName();
+ }
+
+ @Override
+ public Map getOutputFormatConfiguration() {
+ Map outputFormatConfiguration = new HashMap<>(delegate.getOutputFormatConfiguration());
+ outputFormatConfiguration.put(DELEGATE_OUTPUTFORMAT_CLASSNAME, delegate.getOutputFormatClassName());
+ return outputFormatConfiguration;
+ }
+
+ /**
+ * OutputFormat for GCS Sink
+ */
+ public static class GCSOutputFormat extends OutputFormat {
+ private OutputFormat delegateFormat;
+
+ private OutputFormat getDelegateFormatInstance(Configuration configuration) throws IOException {
+ if (delegateFormat != null) {
+ return delegateFormat;
+ }
+
+ String delegateClassName = configuration.get(DELEGATE_OUTPUTFORMAT_CLASSNAME);
+ try {
+ delegateFormat = (OutputFormat) ReflectionUtils
+ .newInstance(configuration.getClassByName(delegateClassName), configuration);
+ return delegateFormat;
+ } catch (ClassNotFoundException e) {
+ throw new IOException(
+ String.format("Unable to instantiate output format for class %s", delegateClassName),
+ e);
+ }
+ }
+
+ @Override
+ public RecordWriter getRecordWriter(TaskAttemptContext taskAttemptContext) throws
+ IOException, InterruptedException {
+ RecordWriter originalWriter = getDelegateFormatInstance(taskAttemptContext.getConfiguration())
+ .getRecordWriter(taskAttemptContext);
+ return new GCSRecordWriter(originalWriter);
+ }
+
+ @Override
+ public void checkOutputSpecs(JobContext jobContext) throws IOException, InterruptedException {
+ getDelegateFormatInstance(jobContext.getConfiguration()).checkOutputSpecs(jobContext);
+ }
+
+ @Override
+ public OutputCommitter getOutputCommitter(TaskAttemptContext taskAttemptContext) throws IOException,
+ InterruptedException {
+ OutputCommitter delegateCommitter = getDelegateFormatInstance(taskAttemptContext.getConfiguration())
+ .getOutputCommitter(taskAttemptContext);
+ return new GCSOutputCommitter(delegateCommitter);
+ }
+ }
+
+ /**
+ * OutputCommitter for GCS
+ */
+ public static class GCSOutputCommitter extends OutputCommitter {
+
+ private final OutputCommitter delegate;
+
+ public GCSOutputCommitter(OutputCommitter delegate) {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public void setupJob(JobContext jobContext) throws IOException {
+ delegate.setupJob(jobContext);
+ }
+
+ @Override
+ public void cleanupJob(JobContext jobContext) throws IOException {
+ delegate.cleanupJob(jobContext);
+ }
+
+ @Override
+ public void commitJob(JobContext jobContext) throws IOException {
+ delegate.commitJob(jobContext);
+ }
+
+ @Override
+ public void abortJob(JobContext jobContext, JobStatus.State state) throws IOException {
+ delegate.abortJob(jobContext, state);
+ }
+
+ @Override
+ public void setupTask(TaskAttemptContext taskAttemptContext) throws IOException {
+ delegate.setupTask(taskAttemptContext);
+ }
+
+ @Override
+ public boolean needsTaskCommit(TaskAttemptContext taskAttemptContext) throws IOException {
+ return delegate.needsTaskCommit(taskAttemptContext);
+ }
+
+ @Override
+ public void commitTask(TaskAttemptContext taskAttemptContext) throws IOException {
+ /*On commit task, there seems to be some inconsistency across different hadoop implementations regarding the path
+ where output file is stored. For some implementations it appears in the path returned by FileOutputCommitter
+ getCommittedTaskPath and for some it does not.Before commit, the files appear to be consistently present in path
+ returned by FileOutputCommitter getTaskAttemptPath. Hence, find the output file from taskAttemptPath and add
+ metadata before commit happens. After commit, file would have been moved out of the taskAttemptPath. */
+ try {
+ updateMetricMetaData(taskAttemptContext);
+ } catch (Exception exception) {
+ LOG.warn("Unable to record metric for task. Metric emitted for the number of affected rows may be incorrect.",
+ exception);
+ }
+
+ delegate.commitTask(taskAttemptContext);
+ }
+
+ private void updateMetricMetaData(TaskAttemptContext taskAttemptContext) throws IOException {
+ if (!(delegate instanceof FileOutputCommitter)) {
+ return;
+ }
+
+ FileOutputCommitter fileOutputCommitter = (FileOutputCommitter) delegate;
+ Configuration configuration = taskAttemptContext.getConfiguration();
+ //Task is not yet committed, so should be available in attempt path
+ Path taskAttemptPath = fileOutputCommitter.getTaskAttemptPath(taskAttemptContext);
+ if (configuration == null || taskAttemptPath == null) {
+ return;
+ }
+
+ //read the count from configuration
+ String keyInConfig = String.format(RECORD_COUNT_FORMAT, taskAttemptContext.getTaskAttemptID());
+ Map metaData = new HashMap<>();
+ metaData.put(GCSBatchSink.RECORD_COUNT, String.valueOf(configuration.getLong(keyInConfig, 0L)));
+ StorageClient storageClient = getStorageClient(configuration);
+ //update metadata on the output file present in the directory for this task
+ Blob blob = storageClient.pickABlob(taskAttemptPath.toString());
+ if (blob == null) {
+ LOG.info("Could not find a file in path %s to apply count metadata.", taskAttemptPath.toString());
+ return;
+ }
+ storageClient.setMetaData(blob, metaData);
+ }
+
+ @VisibleForTesting
+ StorageClient getStorageClient(Configuration configuration) throws IOException {
+ String project = configuration.get(GCPUtils.FS_GS_PROJECT_ID);
+ String serviceAccountPath = configuration.get(GCPUtils.CLOUD_JSON_KEYFILE);
+ return StorageClient.create(project, serviceAccountPath);
+ }
+
+ @Override
+ public void abortTask(TaskAttemptContext taskAttemptContext) throws IOException {
+ delegate.abortTask(taskAttemptContext);
+ }
+
+ @Override
+ public boolean isCommitJobRepeatable(JobContext jobContext) throws IOException {
+ return delegate.isCommitJobRepeatable(jobContext);
+ }
+
+ @Override
+ public boolean isRecoverySupported(JobContext jobContext) throws IOException {
+ return delegate.isRecoverySupported(jobContext);
+ }
+
+ @Override
+ public boolean isRecoverySupported() {
+ return delegate.isRecoverySupported();
+ }
+
+ @Override
+ public void recoverTask(TaskAttemptContext taskContext) throws IOException {
+ delegate.recoverTask(taskContext);
+ }
+ }
+
+ /**
+ * RecordWriter for GCSSink
+ */
+ public static class GCSRecordWriter extends RecordWriter {
+
+ private final RecordWriter originalWriter;
+ private long recordCount;
+
+ public GCSRecordWriter(RecordWriter originalWriter) {
+ this.originalWriter = originalWriter;
+ }
+
+ @Override
+ public void write(NullWritable nullWritable, StructuredRecord structuredRecord) throws IOException,
+ InterruptedException {
+ originalWriter.write(nullWritable, structuredRecord);
+ recordCount++;
+ }
+
+ @Override
+ public void close(TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
+ originalWriter.close(taskAttemptContext);
+ //Since the file details are not available here, pass the value on in configuration
+ taskAttemptContext.getConfiguration()
+ .setLong(String.format(RECORD_COUNT_FORMAT, taskAttemptContext.getTaskAttemptID()), recordCount);
+ }
+ }
+}
diff --git a/src/main/java/io/cdap/plugin/gcp/gcs/source/GCSSource.java b/src/main/java/io/cdap/plugin/gcp/gcs/source/GCSSource.java
index aec773792b..ee2f432956 100644
--- a/src/main/java/io/cdap/plugin/gcp/gcs/source/GCSSource.java
+++ b/src/main/java/io/cdap/plugin/gcp/gcs/source/GCSSource.java
@@ -1,5 +1,5 @@
/*
- * Copyright © 2015 Cask Data, Inc.
+ * Copyright © 2015-2020 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
@@ -199,7 +199,8 @@ public void validate(FailureCollector collector) {
try {
GCSPath.from(path);
} catch (IllegalArgumentException e) {
- collector.addFailure(e.getMessage(), null).withConfigProperty(NAME_PATH).withStacktrace(e.getStackTrace());
+ collector.addFailure(e.getMessage(), null).withConfigProperty(NAME_PATH)
+ .withStacktrace(e.getStackTrace());
}
}
if (!containsMacro(NAME_FILE_SYSTEM_PROPERTIES)) {
diff --git a/src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkTest.java b/src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkTest.java
index 83db7e88d1..5666b56cb4 100644
--- a/src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkTest.java
+++ b/src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkTest.java
@@ -64,7 +64,8 @@ public void testBigQuerySinkConfig() {
Schema.Field.of("timestamp",
Schema.nullableOf(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS))));
- BigQuerySinkConfig config = new BigQuerySinkConfig("44", "ds", "tb", "bucket", schema.toString());
+ BigQuerySinkConfig config = new BigQuerySinkConfig("44", "ds", "tb", "bucket", schema.toString(),
+ "INTEGER", 0L, 100L, 10L);
MockFailureCollector collector = new MockFailureCollector("bqsink");
config.validate(collector);
Assert.assertEquals(0, collector.getValidationFailures().size());
@@ -75,7 +76,8 @@ public void testBigQuerySinkInvalidConfig() {
Schema invalidSchema = Schema.recordOf("record",
Schema.Field.of("id", Schema.of(Schema.Type.LONG)));
- BigQuerySinkConfig config = new BigQuerySinkConfig("reference!!", "ds", "tb", "buck3t$$", invalidSchema.toString());
+ BigQuerySinkConfig config = new BigQuerySinkConfig("reference!!", "ds", "tb", "buck3t$$", invalidSchema.toString(),
+ "INTEGER", 0L, 100L, 10L);
MockFailureCollector collector = new MockFailureCollector("bqsink");
config.validate(collector);
List failures = collector.getValidationFailures();
@@ -129,7 +131,8 @@ private static BigQuerySink getSinkToTest(Job mockJob) throws NoSuchFieldExcepti
Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
Schema.Field.of("name", Schema.of(Schema.Type.STRING)));
BigQuerySinkConfig config =
- new BigQuerySinkConfig("testmetric", "ds", "tb", "bkt", schema.toString());
+ new BigQuerySinkConfig("testmetric", "ds", "tb", "bkt", schema.toString(),
+ null, null, null, null);
BigQuery bigQueryMock = mock(BigQuery.class);
BigQuerySink sink = new BigQuerySink(config);
setBigQuery(sink, bigQueryMock);
@@ -204,7 +207,8 @@ private BigQuerySink getValidationTestSink(boolean truncateTable) throws NoSuchF
Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
Schema.Field.of("name", Schema.of(Schema.Type.STRING)));
BigQuerySinkConfig config =
- new BigQuerySinkConfig("testmetric", "ds", "tb", "bkt", schema.toString());
+ new BigQuerySinkConfig("testmetric", "ds", "tb", "bkt", schema.toString(),
+ "INTEGER", 0L, 100L, 10L);
FieldSetter.setField(config, AbstractBigQuerySinkConfig.class.getDeclaredField("truncateTable"),
truncateTable);
BigQuery bigQueryMock = mock(BigQuery.class);
diff --git a/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSinkTest.java b/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSinkTest.java
new file mode 100644
index 0000000000..8c051ad551
--- /dev/null
+++ b/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSinkTest.java
@@ -0,0 +1,63 @@
+/*
+ * Copyright © 2020 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.gcp.gcs.sink;
+
+import io.cdap.cdap.etl.mock.validation.MockFailureCollector;
+import io.cdap.plugin.gcp.common.GCPReferenceSinkConfig;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.internal.util.reflection.FieldSetter;
+
+public class GCSBatchSinkTest {
+
+ @Test
+ public void testValidFSProperties() throws NoSuchFieldException {
+ GCSBatchSink.GCSBatchSinkConfig config = getConfig(null);
+ MockFailureCollector collector = new MockFailureCollector("gcssink");
+ config.validate(collector);
+ Assert.assertEquals(0, collector.getValidationFailures().size());
+ }
+
+ private GCSBatchSink.GCSBatchSinkConfig getConfig(String fileSystemProperties) throws NoSuchFieldException {
+ GCSBatchSink.GCSBatchSinkConfig gcsBatchSinkConfig = new GCSBatchSink.GCSBatchSinkConfig();
+ FieldSetter
+ .setField(gcsBatchSinkConfig, GCSBatchSink.GCSBatchSinkConfig.class.getDeclaredField("path"), "gs://test");
+ FieldSetter.setField(gcsBatchSinkConfig, GCSBatchSink.GCSBatchSinkConfig.class.getDeclaredField("format"), "avro");
+ FieldSetter
+ .setField(gcsBatchSinkConfig, GCPReferenceSinkConfig.class.getDeclaredField("referenceName"), "testref");
+ FieldSetter
+ .setField(gcsBatchSinkConfig, GCSBatchSink.GCSBatchSinkConfig.class.getDeclaredField("fileSystemProperties"),
+ fileSystemProperties);
+ return gcsBatchSinkConfig;
+ }
+
+ @Test
+ public void testValidFSProperties1() throws NoSuchFieldException {
+ GCSBatchSink.GCSBatchSinkConfig config = getConfig("{\"key\":\"val\"}");
+ MockFailureCollector collector = new MockFailureCollector("gcssink");
+ config.validate(collector);
+ Assert.assertEquals(0, collector.getValidationFailures().size());
+ }
+
+ @Test
+ public void testInvalidFSProperties() throws NoSuchFieldException {
+ GCSBatchSink.GCSBatchSinkConfig config = getConfig("{\"key\":}");
+ MockFailureCollector collector = new MockFailureCollector("gcssink");
+ config.validate(collector);
+ Assert.assertEquals(1, collector.getValidationFailures().size());
+ }
+}
diff --git a/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputformatProviderTest.java b/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputformatProviderTest.java
new file mode 100644
index 0000000000..63405858a0
--- /dev/null
+++ b/src/test/java/io/cdap/plugin/gcp/gcs/sink/GCSOutputformatProviderTest.java
@@ -0,0 +1,100 @@
+package io.cdap.plugin.gcp.gcs.sink;
+
+import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.plugin.gcp.gcs.StorageClient;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.io.NullWritable;
+import org.apache.hadoop.mapreduce.JobContext;
+import org.apache.hadoop.mapreduce.JobStatus;
+import org.apache.hadoop.mapreduce.RecordWriter;
+import org.apache.hadoop.mapreduce.TaskAttemptContext;
+import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+
+/**
+ * Tests for GCSOutputformatProvider
+ */
+public class GCSOutputformatProviderTest {
+
+ @Test
+ public void testRecordWriter() throws IOException, InterruptedException {
+ RecordWriter mockWriter = Mockito.mock(RecordWriter.class);
+ GCSOutputFormatProvider.GCSRecordWriter recordWriterToTest = new GCSOutputFormatProvider.GCSRecordWriter(
+ mockWriter);
+ NullWritable mockWritable = Mockito.mock(NullWritable.class);
+ StructuredRecord mockRecord = Mockito.mock(StructuredRecord.class);
+ for (int i = 0; i < 5; i++) {
+ recordWriterToTest.write(mockWritable, mockRecord);
+ }
+ TaskAttemptContext mockContext = Mockito.mock(TaskAttemptContext.class);
+ Configuration configuration = new Configuration();
+ Mockito.when(mockContext.getConfiguration()).thenReturn(configuration);
+ recordWriterToTest.close(mockContext);
+ //Verify that the delegate calls are being done as expected
+ Mockito.verify(mockWriter, Mockito.times(5)).write(mockWritable, mockRecord);
+ //verify that count is being recorded as expected
+ Assert.assertTrue(configuration.getLong(
+ String.format(GCSOutputFormatProvider.RECORD_COUNT_FORMAT, mockContext.getTaskAttemptID()), 0) == 5);
+ }
+
+ @Test
+ public void testGCSOutputCommitter() throws IOException, URISyntaxException {
+ FileOutputCommitter fileOutputCommitter = Mockito.mock(FileOutputCommitter.class);
+ GCSOutputFormatProvider.GCSOutputCommitter committer = new GCSOutputFormatProvider.GCSOutputCommitter(
+ fileOutputCommitter);
+ GCSOutputFormatProvider.GCSOutputCommitter committerToTest = Mockito.spy(committer);
+ JobContext mockJobContext = Mockito.mock(JobContext.class);
+ JobStatus.State mockJobState = JobStatus.State.SUCCEEDED;
+ TaskAttemptContext mockContext = Mockito.mock(TaskAttemptContext.class);
+ //test all the delegation
+ committerToTest.abortJob(mockJobContext, mockJobState);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).abortJob(mockJobContext, mockJobState);
+
+ committerToTest.abortTask(mockContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).abortTask(mockContext);
+
+ committerToTest.cleanupJob(mockJobContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).cleanupJob(mockJobContext);
+
+ committerToTest.commitJob(mockJobContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).commitJob(mockJobContext);
+
+ committerToTest.isCommitJobRepeatable(mockJobContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).isCommitJobRepeatable(mockJobContext);
+
+ committerToTest.isRecoverySupported();
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).isRecoverySupported();
+
+ committerToTest.isRecoverySupported(mockJobContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).isRecoverySupported(mockJobContext);
+
+ committerToTest.needsTaskCommit(mockContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).needsTaskCommit(mockContext);
+
+ committerToTest.recoverTask(mockContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).recoverTask(mockContext);
+
+ committerToTest.setupJob(mockJobContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).setupJob(mockJobContext);
+
+ committerToTest.setupTask(mockContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).setupTask(mockContext);
+
+ Configuration configuration = new Configuration();
+ Mockito.when(mockContext.getConfiguration()).thenReturn(configuration);
+ Path committedTaskPathMock = Mockito.mock(Path.class);
+ Mockito.when(fileOutputCommitter.getCommittedTaskPath(mockContext)).thenReturn(committedTaskPathMock);
+ Mockito.when(committedTaskPathMock.toString()).thenReturn("gs://test");
+ StorageClient mockStorage = Mockito.mock(StorageClient.class);
+ Mockito.when(committerToTest.getStorageClient(configuration)).thenReturn(mockStorage);
+
+ committerToTest.commitTask(mockContext);
+ Mockito.verify(fileOutputCommitter, Mockito.times(1)).commitTask(mockContext);
+ }
+}
diff --git a/widgets/BigQueryMultiTable-batchsink.json b/widgets/BigQueryMultiTable-batchsink.json
index 0137b31f7e..00fcff08d5 100644
--- a/widgets/BigQueryMultiTable-batchsink.json
+++ b/widgets/BigQueryMultiTable-batchsink.json
@@ -23,6 +23,14 @@
"default": "auto-detect"
}
},
+ {
+ "widget-type": "textbox",
+ "label": "Dataset Project Id",
+ "name": "datasetProject",
+ "widget-attributes": {
+ "placeholder": "The project in which the dataset is located/should be created. Defaults to the project specified in the Project Id property."
+ }
+ },
{
"widget-type": "textbox",
"label": "Dataset",
diff --git a/widgets/BigQueryTable-batchsink.json b/widgets/BigQueryTable-batchsink.json
index af48f9cbab..cfc52651af 100644
--- a/widgets/BigQueryTable-batchsink.json
+++ b/widgets/BigQueryTable-batchsink.json
@@ -2,7 +2,7 @@
"metadata": {
"spec-version": "1.5"
},
- "display-name" : "BigQuery",
+ "display-name": "BigQuery",
"configuration-groups": [
{
"label": "Basic",
@@ -11,7 +11,7 @@
"widget-type": "textbox",
"label": "Reference Name",
"name": "referenceName",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Name used to identify this sink for lineage"
}
},
@@ -19,15 +19,23 @@
"widget-type": "textbox",
"label": "Project ID",
"name": "project",
- "widget-attributes" : {
+ "widget-attributes": {
"default": "auto-detect"
}
},
+ {
+ "widget-type": "textbox",
+ "label": "Dataset Project Id",
+ "name": "datasetProject",
+ "widget-attributes": {
+ "placeholder": "The project in which the dataset is located/should be created. Defaults to the project specified in the Project Id property."
+ }
+ },
{
"widget-type": "textbox",
"label": "Dataset",
"name": "dataset",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Dataset the table belongs to"
}
},
@@ -35,7 +43,7 @@
"widget-type": "textbox",
"label": "Table",
"name": "table",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Table to write to"
}
},
@@ -43,7 +51,7 @@
"widget-type": "textbox",
"label": "Temporary Bucket Name",
"name": "bucket",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Google Cloud Storage bucket for temporary data"
}
},
@@ -51,20 +59,20 @@
"widget-type": "textbox",
"label": "GCS Upload Request Chunk Size",
"name": "gcsChunkSize",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "GCS upload request chunk size in bytes"
}
}
]
},
{
- "label" : "Credentials",
- "properties" : [
+ "label": "Credentials",
+ "properties": [
{
"widget-type": "textbox",
"label": "Service Account File Path",
"name": "serviceFilePath",
- "widget-attributes" : {
+ "widget-attributes": {
"default": "auto-detect"
}
}
@@ -75,8 +83,8 @@
"properties": [
{
"widget-type": "radio-group",
- "name" : "operation",
- "label" : "Operation",
+ "name": "operation",
+ "label": "Operation",
"widget-attributes": {
"layout": "inline",
"default": "insert",
@@ -100,7 +108,7 @@
"name": "relationTableKey",
"widget-type": "csv",
"label": "Table Key",
- "widget-attributes" : {}
+ "widget-attributes": {}
},
{
"name": "dedupeBy",
@@ -109,14 +117,17 @@
"widget-attributes": {
"delimiter": ",",
"kv-delimiter": " ",
- "dropdownOptions": [ "ASC", "DESC"]
+ "dropdownOptions": [
+ "ASC",
+ "DESC"
+ ]
}
},
{
"widget-type": "textbox",
"label": "Partition Filter",
"name": "partitionFilter",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Filter that can be used for partition elimination"
}
},
@@ -161,7 +172,7 @@
"widget-type": "textbox",
"label": "Location",
"name": "location",
- "widget-attributes" : {
+ "widget-attributes": {
"default": "US"
}
},
@@ -181,11 +192,52 @@
"default": "false"
}
},
+ {
+ "widget-type": "radio-group",
+ "label": "Partitioning type",
+ "name": "partitioningType",
+ "widget-attributes": {
+ "layout": "inline",
+ "default": "TIME",
+ "options": [
+ {
+ "id": "TIME",
+ "label": "Time"
+ },
+ {
+ "id": "INTEGER",
+ "label": "Integer"
+ },
+ {
+ "id": "NONE",
+ "label": "None"
+ }
+ ]
+ }
+ },
+ {
+ "widget-type": "Number",
+ "label": "Range Start (inclusive)",
+ "name": "rangeStart",
+ "widget-attributes": {
+ "default": "0"
+ }
+ },
+ {
+ "widget-type": "Number",
+ "label": "Range End (exclusive)",
+ "name": "rangeEnd"
+ },
+ {
+ "widget-type": "Number",
+ "label": "Range Interval",
+ "name": "rangeInterval"
+ },
{
"widget-type": "textbox",
"label": "Partition Field",
"name": "partitionByField",
- "widget-attributes" : {
+ "widget-attributes": {
"placeholder": "Table field for partitioning"
}
},
@@ -209,7 +261,7 @@
"name": "clusteringOrder",
"widget-type": "csv",
"label": "Clustering Order",
- "widget-attributes" : {}
+ "widget-attributes": {}
}
]
}
@@ -233,6 +285,52 @@
}
}
],
+ "filters": [
+ {
+ "name": "PartitioningIntegerFieldsFilter",
+ "condition": {
+ "expression": "partitioningType == 'INTEGER'"
+ },
+ "show": [
+ {
+ "type": "property",
+ "name": "rangeStart"
+ },
+ {
+ "type": "property",
+ "name": "rangeEnd"
+ },
+ {
+ "type": "property",
+ "name": "rangeInterval"
+ },
+ {
+ "type": "property",
+ "name": "clientAccessToken"
+ }
+ ]
+ },
+ {
+ "name": "PartitionFieldFilter",
+ "condition": {
+ "expression": "createPartitionedTable == true || partitioningType == 'INTEGER' || partitioningType == 'TIME'"
+ },
+ "show": [
+ {
+ "type": "property",
+ "name": "partitionByField"
+ },
+ {
+ "type": "property",
+ "name": "partitionFilterRequired"
+ },
+ {
+ "type": "property",
+ "name": "clusteringOrder"
+ }
+ ]
+ }
+ ],
"jump-config": {
"datasets": [
{
diff --git a/widgets/GCS-batchsink.json b/widgets/GCS-batchsink.json
index 53180affc6..06062993b7 100644
--- a/widgets/GCS-batchsink.json
+++ b/widgets/GCS-batchsink.json
@@ -92,6 +92,21 @@
}
}
]
+ },
+ {
+ "label": "Advanced",
+ "properties": [
+ {
+ "widget-type": "textbox",
+ "label": "Output File Prefix",
+ "name": "outputFileNameBase"
+ },
+ {
+ "widget-type": "json-editor",
+ "label": "File System Properties",
+ "name": "fileSystemProperties"
+ }
+ ]
}
],
"outputs": [
diff --git a/widgets/GCSArgumentSetter-action.json b/widgets/GCSArgumentSetter-action.json
new file mode 100644
index 0000000000..938f171619
--- /dev/null
+++ b/widgets/GCSArgumentSetter-action.json
@@ -0,0 +1,92 @@
+{
+ "metadata": {
+ "spec-version": "1.5"
+ },
+ "display-name": "GCS Argument Setter",
+ "configuration-groups": [
+ {
+ "label": "Basic",
+ "properties": [
+ {
+ "widget-type": "textbox",
+ "label": "Reference Name",
+ "name": "referenceName",
+ "widget-attributes": {
+ "placeholder": "Name used to identify this source for lineage"
+ }
+ },
+ {
+ "widget-type": "textbox",
+ "label": "Project ID",
+ "name": "project",
+ "widget-attributes": {
+ "default": "auto-detect"
+ }
+ },
+ {
+ "widget-type": "textbox",
+ "label": "Path",
+ "name": "path",
+ "widget-attributes": {
+ "placeholder": "gs:///path/to/configuration.json"
+ }
+ },
+ {
+ "label": "Provide Service Account as JSON Content",
+ "name": "serviceAccountType",
+ "widget-type": "toggle",
+ "widget-attributes": {
+ "on": {
+ "value": "JSON",
+ "label": "Yes"
+ },
+ "off": {
+ "value": "filePath",
+ "label": "No"
+ },
+ "default": "filePath"
+ }
+ },
+ {
+ "widget-type": "textbox",
+ "label": "Service Account File Path",
+ "name": "serviceFilePath",
+ "widget-attributes": {
+ "default": "auto-detect"
+ }
+ },
+ {
+ "widget-type": "textbox",
+ "label": "Service Account JSON",
+ "name": "serviceAccountJSON"
+ }
+ ]
+ }
+ ],
+ "filters": [
+ {
+ "name": "ServiceAuthenticationTypeFilePath",
+ "condition": {
+ "expression": "serviceAccountType == 'filePath'"
+ },
+ "show": [
+ {
+ "type": "property",
+ "name": "serviceFilePath"
+ }
+ ]
+ },
+ {
+ "name": "ServiceAuthenticationTypeJSON",
+ "condition": {
+ "expression": "serviceAccountType == 'JSON'"
+ },
+ "show": [
+ {
+ "type": "property",
+ "name": "serviceAccountJSON"
+ }
+ ]
+ }
+ ]
+}
diff --git a/widgets/GCSMultiFiles-batchsink.json b/widgets/GCSMultiFiles-batchsink.json
index e14420c2eb..5ac4a2765f 100644
--- a/widgets/GCSMultiFiles-batchsink.json
+++ b/widgets/GCSMultiFiles-batchsink.json
@@ -103,6 +103,16 @@
"widget-attributes": {
"default": "tablename"
}
+ },
+ {
+ "widget-type": "hidden",
+ "label": "Output File Prefix",
+ "name": "outputFileNameBase"
+ },
+ {
+ "widget-type": "hidden",
+ "label": "File System Properties",
+ "name": "fileSystemProperties"
}
]
},