modifying execute API to get column nullability state

wrangler-api/src/main/java/io/cdap/wrangler/api/NullHandlingException.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright © 2016-2019 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.wrangler.api;
+
+/**
+ * A Null Handling specific exception used for communicating issues with Null Handling in a column.
+ */
+public class NullHandlingException extends Exception {
+  public NullHandlingException(Exception e) {
+    super(e);
+  }
+
+  public NullHandlingException(String message) {
+    super(message);
+  }
+
+}

wrangler-core/src/main/java/io/cdap/wrangler/executor/RecipePipelineExecutor.java

@@ -30,6 +30,7 @@
 import io.cdap.wrangler.api.ReportErrorAndProceed;
 import io.cdap.wrangler.api.Row;
 import io.cdap.wrangler.api.TransientVariableScope;
+import io.cdap.wrangler.proto.workspace.v2.Workspace.UserDefinedAction;
 import io.cdap.wrangler.schema.DirectiveOutputSchemaGenerator;
 import io.cdap.wrangler.schema.DirectiveSchemaResolutionContext;
 import io.cdap.wrangler.schema.TransientStoreKeys;
@@ -40,6 +41,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -56,10 +59,13 @@ public final class RecipePipelineExecutor implements RecipePipeline<Row, Structu
   private final RecipeParser recipeParser;
   private final ExecutorContext context;
   private List<Directive> directives;
+  private HashMap<String, UserDefinedAction> nullabilityMap;
 
-  public RecipePipelineExecutor(RecipeParser recipeParser, @Nullable ExecutorContext context) {
+  public RecipePipelineExecutor(RecipeParser recipeParser, @Nullable ExecutorContext context,
+      HashMap<String, UserDefinedAction> nullabilityMap) {
     this.context = context;
     this.recipeParser = recipeParser;
+    this.nullabilityMap = nullabilityMap;
   }
 
   /**

wrangler-core/src/test/java/io/cdap/wrangler/TestingRig.java

@@ -89,7 +89,7 @@ public static List<Row> execute(String[] recipe, List<Row> rows, ExecutorContext
 
     String migrate = new MigrateToV2(recipe).migrate();
     RecipeParser parser = new GrammarBasedParser(Contexts.SYSTEM, migrate, registry);
-    return new RecipePipelineExecutor(parser, context).execute(rows);
+    return new RecipePipelineExecutor(parser, context, null).execute(rows);
   }
 
   /**
@@ -112,7 +112,7 @@ public static Pair<List<Row>, List<Row>> executeWithErrors(String[] recipe, List
 
     String migrate = new MigrateToV2(recipe).migrate();
     RecipeParser parser = new GrammarBasedParser(Contexts.SYSTEM, migrate, registry);
-    RecipePipeline pipeline = new RecipePipelineExecutor(parser, context);
+    RecipePipeline pipeline = new RecipePipelineExecutor(parser, context, null);
     List<Row> results = pipeline.execute(rows);
     List<Row> errors = pipeline.errors();
     return new Pair<>(results, errors);
@@ -126,7 +126,7 @@ public static RecipePipeline execute(String[] recipe)
 
     String migrate = new MigrateToV2(recipe).migrate();
     RecipeParser parser = new GrammarBasedParser(Contexts.SYSTEM, migrate, registry);
-    return new RecipePipelineExecutor(parser, new TestingPipelineContext());
+    return new RecipePipelineExecutor(parser, new TestingPipelineContext(), null);
   }
 
   public static RecipeParser parse(String[] recipe) throws DirectiveParseException, DirectiveLoadException {

wrangler-proto/src/main/java/io/cdap/wrangler/proto/workspace/v2/DirectiveExecutionRequest.java

@@ -17,7 +17,9 @@
 
 package io.cdap.wrangler.proto.workspace.v2;
 
+import io.cdap.wrangler.proto.workspace.v2.Workspace.UserDefinedAction;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 
 /**
@@ -26,10 +28,14 @@
 public class DirectiveExecutionRequest {
   private final List<String> directives;
   private final int limit;
+  private final HashMap<String, UserDefinedAction> nullabilityMap;
 
-  public DirectiveExecutionRequest(List<String> directives, int limit) {
+
+  public DirectiveExecutionRequest(List<String> directives, int limit,
+      HashMap<String, UserDefinedAction> nullabilityMap) {
     this.directives = directives;
     this.limit = limit;
+    this.nullabilityMap = nullabilityMap;
   }
 
   public int getLimit() {
@@ -39,4 +45,8 @@ public int getLimit() {
   public List<String> getDirectives() {
     return directives == null ? Collections.emptyList() : directives;
   }
+
+  public HashMap<String, UserDefinedAction> getNullabilityMap() {
+    return nullabilityMap;
+  }
 }

wrangler-proto/src/main/java/io/cdap/wrangler/proto/workspace/v2/Workspace.java

@@ -20,6 +20,7 @@
 import com.google.gson.JsonObject;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Objects;
 import javax.annotation.Nullable;
@@ -38,16 +39,20 @@ public class Workspace {
   // this is for insights page in UI
   private final JsonObject insights;
 
+  private HashMap<String, UserDefinedAction> nullabilityMap;
+
   private Workspace(String workspaceName, String workspaceId, List<String> directives,
                     long createdTimeMillis, long updatedTimeMillis, @Nullable SampleSpec sampleSpec,
-                    JsonObject insights) {
+                    JsonObject insights, HashMap<String, UserDefinedAction> nullabilityMap) {
     this.workspaceName = workspaceName;
     this.workspaceId = workspaceId;
     this.directives = directives;
     this.createdTimeMillis = createdTimeMillis;
     this.updatedTimeMillis = updatedTimeMillis;
     this.sampleSpec = sampleSpec;
     this.insights = insights;
+    this.nullabilityMap = nullabilityMap == null || nullabilityMap.isEmpty() ?
+        new HashMap<>() : nullabilityMap;
   }
 
   public String getWorkspaceName() {
@@ -79,6 +84,15 @@ public JsonObject getInsights() {
     return insights;
   }
 
+  public HashMap<String, UserDefinedAction> getNullabilityMap() {
+    return nullabilityMap;
+  }
+
+  public void setNullabilityMap(
+      HashMap<String, UserDefinedAction> nullabilityMap) {
+    this.nullabilityMap = nullabilityMap;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -111,7 +125,8 @@ public static Builder builder(Workspace existing) {
              .setCreatedTimeMillis(existing.getCreatedTimeMillis())
              .setUpdatedTimeMillis(existing.getUpdatedTimeMillis())
              .setSampleSpec(existing.getSampleSpec())
-             .setInsights(existing.getInsights());
+             .setInsights(existing.getInsights())
+             .setNullabilityMap(existing.getNullabilityMap());
   }
 
   /**
@@ -125,6 +140,7 @@ public static class Builder {
     private long updatedTimeMillis;
     private SampleSpec sampleSpec;
     private JsonObject insights;
+    private HashMap<String, UserDefinedAction> nullabilityMap;
 
     Builder(String name, String workspaceId) {
       this.workspaceName = name;
@@ -159,9 +175,23 @@ public Builder setInsights(JsonObject insights) {
       return this;
     }
 
+    public Builder setNullabilityMap (HashMap<String, UserDefinedAction> nullabilityMap) {
+      this.nullabilityMap = nullabilityMap;
+      return this;
+    }
+
     public Workspace build() {
       return new Workspace(workspaceName, workspaceId, directives, createdTimeMillis, updatedTimeMillis, sampleSpec,
-                           insights);
+                           insights, nullabilityMap);
     }
   }
+
+  /**
+   * UserDefinedAction enum.
+   */
+  public enum UserDefinedAction {
+    SKIP_ROW,
+    SEND_TO_ERROR_COLLECTOR,
+    ERROR_PIPELINE,
+  }
 }

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/AbstractDirectiveHandler.java

@@ -46,6 +46,9 @@
 import io.cdap.wrangler.proto.workspace.ColumnValidationResult;
 import io.cdap.wrangler.proto.workspace.WorkspaceValidationResult;
 import io.cdap.wrangler.proto.workspace.v2.DirectiveExecutionResponse;
+import io.cdap.wrangler.proto.workspace.v2.SampleSpec;
+import io.cdap.wrangler.proto.workspace.v2.Workspace;
+import io.cdap.wrangler.proto.workspace.v2.Workspace.UserDefinedAction;
 import io.cdap.wrangler.registry.CompositeDirectiveRegistry;
 import io.cdap.wrangler.registry.DirectiveRegistry;
 import io.cdap.wrangler.registry.SystemDirectiveRegistry;
@@ -118,7 +121,8 @@ protected <E extends Exception> List<Row> executeDirectives(
       String namespace,
       List<String> directives,
       List<Row> sample,
-      GrammarWalker.Visitor<E> grammarVisitor) throws DirectiveParseException, E, RecipeException {
+      GrammarWalker.Visitor<E> grammarVisitor,
+      HashMap<String, UserDefinedAction> nullabilityMap) throws DirectiveParseException, E, RecipeException {
 
     if (directives.isEmpty()) {
       return sample;
@@ -139,8 +143,11 @@ protected <E extends Exception> List<Row> executeDirectives(
                                                  new ConfigDirectiveContext(DirectiveConfig.EMPTY));
     try (RecipePipelineExecutor executor = new RecipePipelineExecutor(parser,
                                                                       new ServicePipelineContext(
-                                                                        namespace, ExecutorContext.Environment.SERVICE,
-                                                                        getContext(), TRANSIENT_STORE))) {
+                                                                        namespace,
+                                                                          ExecutorContext.Environment.SERVICE,
+                                                                          getContext(),
+                                                                          TRANSIENT_STORE),
+        nullabilityMap)) {
       List<Row> result = executor.execute(sample);
 
       List<ErrorRecordBase> errors = executor.errors()

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/DirectivesHandler.java

@@ -1095,7 +1095,7 @@ private <E extends Exception> List<Row> executeDirectives(NamespacedId id, List<
       // Extract rows from the workspace.
       List<Row> rows = fromWorkspace(workspace);
       return executeDirectives(id.getNamespace().getName(), directives, sample.apply(rows),
-                               grammarVisitor);
+          grammarVisitor, null);
     });
   }
 }

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/RemoteExecutionTask.java

@@ -107,7 +107,7 @@ public void run(RunnableTaskContext runnableTaskContext) throws Exception {
                                                                           namespace,
                                                                           ExecutorContext.Environment.SERVICE,
                                                                           systemAppContext,
-                                                                          new DefaultTransientStore()))) {
+                                                                          new DefaultTransientStore()), null)) {
         rows = executor.execute(rows);
         List<ErrorRecordBase> errors = executor.errors().stream()
             .filter(ErrorRecordBase::isShownInWrangler)

wrangler-service/src/main/java/io/cdap/wrangler/service/directive/WorkspaceHandler.java

@@ -67,6 +67,7 @@
 import io.cdap.wrangler.proto.workspace.v2.ServiceResponse;
 import io.cdap.wrangler.proto.workspace.v2.StageSpec;
 import io.cdap.wrangler.proto.workspace.v2.Workspace;
+import io.cdap.wrangler.proto.workspace.v2.Workspace.UserDefinedAction;
 import io.cdap.wrangler.proto.workspace.v2.WorkspaceCreationRequest;
 import io.cdap.wrangler.proto.workspace.v2.WorkspaceDetail;
 import io.cdap.wrangler.proto.workspace.v2.WorkspaceId;
@@ -169,7 +170,8 @@ public void createWorkspace(HttpServiceRequest request, HttpServiceResponder res
       long now = System.currentTimeMillis();
       Workspace workspace = Workspace.builder(generateWorkspaceName(wsId, creationRequest.getSampleRequest().getPath()),
                                               wsId.getWorkspaceId())
-                              .setCreatedTimeMillis(now).setUpdatedTimeMillis(now).setSampleSpec(spec).build();
+                              .setCreatedTimeMillis(now).setUpdatedTimeMillis(now)
+          .setSampleSpec(spec).setNullabilityMap(new HashMap<>()).build();
       wsStore.saveWorkspace(wsId, new WorkspaceDetail(workspace, rows));
       responder.sendJson(wsId.getWorkspaceId());
     });
@@ -472,6 +474,12 @@ private DirectiveExecutionResponse execute(NamespaceSummary ns, HttpServiceReque
 
     WorkspaceDetail detail = wsStore.getWorkspaceDetail(workspaceId);
     UserDirectivesCollector userDirectivesCollector = new UserDirectivesCollector();
+    HashMap<String, UserDefinedAction> nullabilityMap = executionRequest.getNullabilityMap() == null ?
+        new HashMap<>() : executionRequest.getNullabilityMap();
+    if (!nullabilityMap.isEmpty()) {
+      //change nullabilityMap in Workspace Object
+      changeNullability(nullabilityMap, workspaceId);
+    }
     List<Row> result = executeDirectives(ns.getName(), directives, detail,
                                          userDirectivesCollector);
     DirectiveExecutionResponse response = generateExecutionResponse(result,
@@ -484,6 +492,18 @@ private DirectiveExecutionResponse execute(NamespaceSummary ns, HttpServiceReque
     return response;
   }
 
+  private void changeNullability(HashMap<String, UserDefinedAction> columnMappings,
+       WorkspaceId workspaceId) throws Exception {
+    try {
+      Workspace workspace = wsStore.getWorkspace(workspaceId);
+      workspace.setNullabilityMap(columnMappings);
+      wsStore.updateWorkspace(workspaceId, workspace);
+    } catch (Exception e) {
+      throw new RuntimeException("Error in setting nullabilityMap of columns ", e);
+    }
+    }
+
+
   /**
    * Get source specs, contains some hacky way on dealing with the csv parser
    */
@@ -580,7 +600,7 @@ private <E extends Exception> List<Row> executeLocally(String namespace, List<St
     // load the udd
     composite.reload(namespace);
     return executeDirectives(namespace, directives, new ArrayList<>(detail.getSample()),
-                             grammarVisitor);
+                             grammarVisitor, detail.getWorkspace().getNullabilityMap());
   }
 
   /**

wrangler-test/src/main/java/io/cdap/wrangler/test/TestingRig.java

@@ -56,7 +56,7 @@ public static RecipePipeline pipeline(Class<? extends Directive> directive, Test
 
     String migrate = new MigrateToV2(recipe.toArray()).migrate();
     RecipeParser parser = new GrammarBasedParser(Contexts.SYSTEM, migrate, registry);
-    return new RecipePipelineExecutor(parser, null);
+    return new RecipePipelineExecutor(parser, null, null);
   }
 
   public static RecipeParser parser(Class<? extends Directive> directive, String[] recipe)

wrangler-transform/src/main/java/io/cdap/wrangler/Wrangler.java

@@ -67,6 +67,7 @@
 import io.cdap.wrangler.parser.NoOpDirectiveContext;
 import io.cdap.wrangler.parser.RecipeCompiler;
 import io.cdap.wrangler.proto.Contexts;
+import io.cdap.wrangler.proto.workspace.v2.Workspace.UserDefinedAction;
 import io.cdap.wrangler.registry.CompositeDirectiveRegistry;
 import io.cdap.wrangler.registry.DirectiveInfo;
 import io.cdap.wrangler.registry.DirectiveRegistry;
@@ -365,7 +366,7 @@ && checkPreconditionNotEmpty(false)) {
 
     try {
       // Create the pipeline executor with context being set.
-      pipeline = new RecipePipelineExecutor(recipe, ctx);
+      pipeline = new RecipePipelineExecutor(recipe, ctx, null);
     } catch (Exception e) {
       throw new Exception(String.format("Stage:%s - %s", getContext().getStageName(), e.getMessage()), e);
     }