[CDAP-21089] OAuthHandler should not return 5xx on 4xx errors to skip retries and throw errors for invalid macros based on stage validation request.

cdap-api/src/main/java/io/cdap/cdap/api/macro/MacroParserOptions.java

@@ -33,16 +33,18 @@ public class MacroParserOptions {
   private final boolean skipInvalid;
   private final int maxRecurseDepth;
   private final Set<String> functionWhitelist;
+  private final Set<String> doNotSkipInvalidMacroForFunctions;
 
   private MacroParserOptions(boolean evaluateLookups, boolean evaluateFunctions,
       boolean escapingEnabled, boolean skipInvalid,
-      int maxRecurseDepth, Set<String> functionWhitelist) {
+      int maxRecurseDepth, Set<String> functionWhitelist, Set<String> doNotSkipInvalidMacroForFunctions) {
     this.evaluateLookups = evaluateLookups;
     this.evaluateFunctions = evaluateFunctions;
     this.escapingEnabled = escapingEnabled;
     this.maxRecurseDepth = maxRecurseDepth;
     this.skipInvalid = skipInvalid;
     this.functionWhitelist = functionWhitelist;
+    this.doNotSkipInvalidMacroForFunctions = doNotSkipInvalidMacroForFunctions;
   }
 
   public boolean shouldEvaluateLookups() {
@@ -69,6 +71,10 @@ public Set<String> getFunctionWhitelist() {
     return functionWhitelist;
   }
 
+  public Set<String> getDoNotSkipInvalidMacroForFunctions() {
+    return doNotSkipInvalidMacroForFunctions;
+  }
+
   /**
    * @return Builder to create options
    */
@@ -87,6 +93,7 @@ public static class Builder {
     private boolean skipInvalid;
     private int maxRecurseDepth = 10;
     private final Set<String> functionWhitelist = new HashSet<>();
+    private final Set<String> doNotSkipInvalidMacroForFunctions = new HashSet<>();
 
     public Builder disableLookups() {
       evaluateLookups = false;
@@ -113,6 +120,16 @@ public Builder setMaxRecurseDepth(int maxRecurseDepth) {
       return this;
     }
 
+    public Builder setDoNotSkipInvalidMacroForFunctions(Collection<String> whitelist) {
+      doNotSkipInvalidMacroForFunctions.clear();
+      doNotSkipInvalidMacroForFunctions.addAll(whitelist);
+      return this;
+    }
+
+    public Builder setDoNotSkipInvalidMacroForFunctions(String... whitelist) {
+      return setDoNotSkipInvalidMacroForFunctions(Arrays.asList(whitelist));
+    }
+
     public Builder setFunctionWhitelist(Collection<String> whitelist) {
       functionWhitelist.clear();
       functionWhitelist.addAll(whitelist);
@@ -122,10 +139,10 @@ public Builder setFunctionWhitelist(Collection<String> whitelist) {
     public Builder setFunctionWhitelist(String... whitelist) {
       return setFunctionWhitelist(Arrays.asList(whitelist));
     }
-
+    
     public MacroParserOptions build() {
       return new MacroParserOptions(evaluateLookups, evaluateFunctions, escapingEnabled,
-          skipInvalid, maxRecurseDepth, functionWhitelist);
+          skipInvalid, maxRecurseDepth, functionWhitelist, doNotSkipInvalidMacroForFunctions);
     }
   }
 }

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/runtime/plugin/MacroParser.java

@@ -43,6 +43,7 @@ public class MacroParser {
   private final boolean functionsEnabled;
   private final boolean skipInvalid;
   private final Set<String> functionWhitelist;
+  private final Set<String> doNotSkipInvalidMacroForFunctions;
 
   public MacroParser(MacroEvaluator macroEvaluator) {
     this(macroEvaluator, MacroParserOptions.DEFAULT);
@@ -55,6 +56,7 @@ public MacroParser(MacroEvaluator macroEvaluator, MacroParserOptions options) {
     this.functionsEnabled = options.shouldEvaluateFunctions();
     this.functionWhitelist = options.getFunctionWhitelist();
     this.skipInvalid = options.shouldSkipInvalid();
+    this.doNotSkipInvalidMacroForFunctions = options.getDoNotSkipInvalidMacroForFunctions();
   }
 
   /**
@@ -242,7 +244,7 @@ private MacroMetadata getMacroFunctionMetadata(int startIndex, int endIndex, Str
         }
 
       } catch (InvalidMacroException e) {
-        if (!skipInvalid) {
+        if (!skipInvalid || doNotSkipInvalidMacroForFunctions.contains(macroFunction)) {
           throw e;
         }
       }

cdap-app-fabric/src/test/java/io/cdap/cdap/internal/app/runtime/plugin/MacroParserTest.java

@@ -562,4 +562,28 @@ private static void assertSubstitution(String macro, String expected, Map<String
     MacroParser macroParser = new MacroParser(macroEvaluator);
     Assert.assertEquals(expected, macroParser.parse(macro));
   }
+
+  @Test
+  public void testDoNotSkipInvalidMacrosInMacroFunctions() {
+    MacroEvaluator evaluator = new TestMacroEvaluator(Collections.emptyMap(), Collections.emptyMap(), false);
+
+    //TestMacroEvaluator will throw InvalidMacroException if macro function is other than 't' or 'test'
+    // but skip invalid flag will suppress this exception.
+    MacroParser parser = new MacroParser(evaluator, MacroParserOptions.builder().skipInvalidMacros().build());
+    Assert.assertEquals("${oauthAccessToken(provider,credentials)}",
+                        parser.parse("${oauthAccessToken(provider,credentials)}"));
+
+    // If doNotSkipInvalidMacroForFunctions set has that macro, it will not skip validation and will throw Exception
+    MacroParser parserWithDoNotSkipInvalidSet = new MacroParser(evaluator,
+                                                                MacroParserOptions.builder().skipInvalidMacros()
+                                                                  .setDoNotSkipInvalidMacroForFunctions(
+                                                                    "oauthAccessToken", "oauth").build());
+    boolean exceptionThrown = false;
+    try {
+      parserWithDoNotSkipInvalidSet.parse("${oauthAccessToken(provider,credentials)}");
+    } catch (InvalidMacroException e) {
+      exceptionThrown = true;
+    }
+    Assert.assertTrue(exceptionThrown);
+  }
 }

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/service/OAuthHandler.java

@@ -185,7 +185,7 @@ public void putOAuthCredential(HttpServiceRequest request, HttpServiceResponder
 
       if (response.getResponseCode() != 200) {
         throw new OAuthServiceException(
-            HttpURLConnection.HTTP_INTERNAL_ERROR,
+          response.getResponseCode(),
             "Request for refresh token did not return 200. Response code: "
                 + response.getResponseCode()
                 + " , response message: "
@@ -243,7 +243,7 @@ public void getOAuthCredential(HttpServiceRequest request, HttpServiceResponder
 
       if (response.getResponseCode() != 200) {
         throw new OAuthServiceException(
-            HttpURLConnection.HTTP_INTERNAL_ERROR,
+          response.getResponseCode(),
             "Request for refresh token did not return 200. Response code: "
                 + response.getResponseCode()
                 + " , response message: "

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/service/RemoteValidationTask.java

@@ -43,6 +43,7 @@
 
 import java.util.Collections;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 
 /**
@@ -96,10 +97,12 @@ OAuthAccessTokenMacroEvaluator.FUNCTION_NAME, new OAuthAccessTokenMacroEvaluator
     );
     MacroEvaluator macroEvaluator = new DefaultMacroEvaluator(new BasicArguments(arguments), evaluators,
                                                               DefaultMacroEvaluator.MAP_FUNCTIONS);
+    Set<String> doNotSkipInvalidMacroForFunctions = validationRequest.getDoNotSkipInvalidMacroForFunctions();
     MacroParserOptions macroParserOptions = MacroParserOptions.builder()
-      .skipInvalidMacros()
       .setEscaping(false)
       .setFunctionWhitelist(evaluators.keySet())
+      .skipInvalidMacros()
+      .setDoNotSkipInvalidMacroForFunctions(doNotSkipInvalidMacroForFunctions)
       .build();
     Function<Map<String, String>, Map<String, String>> macroFn =
       macroProperties -> systemAppContext

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/service/ValidationHandler.java

@@ -61,6 +61,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
@@ -169,11 +170,14 @@ OAuthAccessTokenMacroEvaluator.FUNCTION_NAME, new OAuthAccessTokenMacroEvaluator
     );
     MacroEvaluator macroEvaluator = new DefaultMacroEvaluator(new BasicArguments(arguments), evaluators,
                                                               DefaultMacroEvaluator.MAP_FUNCTIONS);
+    Set<String> doNotSkipInvalidMacroForFunctions = validationRequest.getDoNotSkipInvalidMacroForFunctions();
     MacroParserOptions macroParserOptions = MacroParserOptions.builder()
-      .skipInvalidMacros()
       .setEscaping(false)
       .setFunctionWhitelist(evaluators.keySet())
+      .skipInvalidMacros()
+      .setDoNotSkipInvalidMacroForFunctions(doNotSkipInvalidMacroForFunctions)
       .build();
+
     Function<Map<String, String>, Map<String, String>> macroFn =
       macroProperties -> getContext().evaluateMacros(namespace, macroProperties, macroEvaluator, macroParserOptions);
     String validationResponse = GSON.toJson(ValidationUtils.validate(

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/test/java/io/cdap/cdap/datapipeline/DataPipelineServiceTest.java

@@ -484,4 +484,34 @@ private StageValidationResponse sendRequest(StageValidationRequest requestBody)
     Assert.assertEquals(200, response.getResponseCode());
     return GSON.fromJson(response.getResponseBodyAsString(), StageValidationResponse.class);
   }
+
+  @Test
+  public void testValidationForInvalidMacrosForFunctions() throws Exception {
+    String stageName = "invalidMacro";
+    Map<String, String> properties =
+      Collections.singletonMap("metadataOperations", "${oauthAccessToken(provider,credential)}");
+    ETLStage stage = new ETLStage(stageName, new ETLPlugin(MockSource.NAME, BatchSource.PLUGIN_TYPE, properties));
+
+    // In case doNotSkipInvalidMacroForFunctions is set in StageValidationRequest,
+    // it will not skip invalid macro given in the list and will throw the exact exception
+    HttpResponse withDoNotSkipInvalidMacroForFunctionsResponse = sendRequestForInvalidMacros(
+      new StageValidationRequest(stage, Collections.emptyList(), false, "oauthAccessToken,oauth"));
+    Assert.assertEquals(500, withDoNotSkipInvalidMacroForFunctionsResponse.getResponseCode());
+
+    // In case doNotSkipInvalidMacroForFunctions is not set in StageValidationRequest,
+    // it will skip invalid macro and will not throw any exception
+    HttpResponse withoutDoNotSkipInvalidMacroForFunctionsResponse = sendRequestForInvalidMacros(
+      new StageValidationRequest(stage, Collections.emptyList(), false));
+    Assert.assertEquals(200, withoutDoNotSkipInvalidMacroForFunctionsResponse.getResponseCode());
+  }
+
+  private HttpResponse sendRequestForInvalidMacros(StageValidationRequest requestBody) throws IOException {
+    URL validatePipelineURL = serviceURI
+      .resolve(String.format("v1/contexts/%s/validations/stage", NamespaceId.DEFAULT.getNamespace()))
+      .toURL();
+    HttpRequest request = HttpRequest.builder(HttpMethod.POST, validatePipelineURL)
+      .withBody(GSON.toJson(requestBody))
+      .build();
+    return HttpRequests.execute(request, new DefaultHttpRequestConfig(false));
+  }
 }

cdap-app-templates/cdap-etl/cdap-etl-proto/src/main/java/io/cdap/cdap/etl/proto/v2/validation/StageValidationRequest.java

@@ -17,9 +17,15 @@
 
 package io.cdap.cdap.etl.proto.v2.validation;
 
+import com.google.common.base.Strings;
 import io.cdap.cdap.etl.proto.v2.ETLStage;
+
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Request to validate a pipeline stage.
@@ -29,13 +35,21 @@ public class StageValidationRequest {
   private final ETLStage stage;
   private final List<StageSchema> inputSchemas;
   private final Boolean resolveMacrosFromPreferences;
+  private final String doNotSkipInvalidMacroForFunctions;
 
+  public StageValidationRequest(ETLStage stage,
+                                List<StageSchema> inputSchemas,
+                                boolean resolveMacrosFromPreferences) {
+    this(stage, inputSchemas, resolveMacrosFromPreferences, null);
+  }
   public StageValidationRequest(ETLStage stage,
       List<StageSchema> inputSchemas,
-      boolean resolveMacrosFromPreferences) {
+      boolean resolveMacrosFromPreferences,
+      String doNotSkipInvalidMacroForFunctions) {
     this.stage = stage;
     this.inputSchemas = inputSchemas;
     this.resolveMacrosFromPreferences = resolveMacrosFromPreferences;
+    this.doNotSkipInvalidMacroForFunctions = doNotSkipInvalidMacroForFunctions;
   }
 
   public ETLStage getStage() {
@@ -50,6 +64,21 @@ public boolean getResolveMacrosFromPreferences() {
     return resolveMacrosFromPreferences != null ? resolveMacrosFromPreferences : false;
   }
 
+  /**
+   *  This method will return macro function names for which invalid macros should not be skipped.
+   * @return Set of macro function names
+   */
+  public Set<String> getDoNotSkipInvalidMacroForFunctions() {
+    Set<String> doNotSkipInvalidMacroForFunctionsSet = new HashSet<>();
+    if (!Strings.isNullOrEmpty(doNotSkipInvalidMacroForFunctions)) {
+      doNotSkipInvalidMacroForFunctionsSet.addAll(Arrays.stream(doNotSkipInvalidMacroForFunctions.split(","))
+                                                    .map(String::trim)
+                                                    .filter(s -> !s.isEmpty())
+                                                    .collect(Collectors.toSet()));
+    }
+    return doNotSkipInvalidMacroForFunctionsSet;
+  }
+
   /**
    * Validate that the request contains all required information. This should be called whenever
    * this instance is created by deserializing user provided input.