modifying execute API to get column nullability state

data-integrations · Nov 28, 2023 · 1dc23e5 · 1dc23e5
1 parent 5790e4b
commit 1dc23e5
Show file tree

Hide file tree

Showing 10 changed files with 93 additions and 12 deletions.
diff --git a/wrangler-api/src/main/java/io/cdap/wrangler/api/NullHandlingException.java b/wrangler-api/src/main/java/io/cdap/wrangler/api/NullHandlingException.java
@@ -0,0 +1,15 @@
+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);
+  }
+
+}
diff --git a/wrangler-core/src/main/java/io/cdap/wrangler/executor/RecipePipelineExecutor.java b/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;
   }
 
   /**

diff --git a/wrangler-core/src/test/java/io/cdap/wrangler/TestingRig.java b/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 {

diff --git a/...er-proto/src/main/java/io/cdap/wrangler/proto/workspace/v2/DirectiveExecutionRequest.java b/...er-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;
+  }
 }
diff --git a/wrangler-proto/src/main/java/io/cdap/wrangler/proto/workspace/v2/Workspace.java b/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,6 +39,8 @@ 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) {
@@ -48,6 +51,7 @@ private Workspace(String workspaceName, String workspaceId, List<String> directi
     this.updatedTimeMillis = updatedTimeMillis;
     this.sampleSpec = sampleSpec;
     this.insights = insights;
+    this.nullabilityMap = new HashMap<>();
   }
 
   public String getWorkspaceName() {
@@ -79,6 +83,15 @@ public JsonObject getInsights() {
     return insights;
   }
 
+  public HashMap<String, UserDefinedAction> getColumnMappings() {
+    return nullabilityMap;
+  }
+
+  public void setColumnMappings(
+      HashMap<String, UserDefinedAction> nullabilityMap) {
+    this.nullabilityMap = nullabilityMap;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -164,4 +177,15 @@ public Workspace build() {
                            insights);
     }
   }
+
+  /**
+   * UserDefinedAction enum.
+   */
+  public enum UserDefinedAction {
+    NO_ACTION,
+    SKIP_ROW,
+    SEND_TO_ERROR_COLLECTOR,
+    ERROR_PIPELINE,
+    NULLABLE
+  }
 }
diff --git a/...er-service/src/main/java/io/cdap/wrangler/service/directive/AbstractDirectiveHandler.java b/...er-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,
+      Workspace workspace) 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),
+        workspace.getColumnMappings())) {
       List<Row> result = executor.execute(sample);
 
       List<ErrorRecordBase> errors = executor.errors()

diff --git a/wrangler-service/src/main/java/io/cdap/wrangler/service/directive/DirectivesHandler.java b/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);
     });
   }
 }
diff --git a/wrangler-service/src/main/java/io/cdap/wrangler/service/directive/RemoteExecutionTask.java b/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)

diff --git a/wrangler-service/src/main/java/io/cdap/wrangler/service/directive/WorkspaceHandler.java b/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;
@@ -472,6 +473,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 +491,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.setColumnMappings(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 +599,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());
   }
 
   /**

diff --git a/wrangler-test/src/main/java/io/cdap/wrangler/test/TestingRig.java b/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)