fix(core): add secret consumer when rendering variables for subflows

kestra-io · Jun 20, 2024 · a862f9b · a862f9b
1 parent ca39844
commit a862f9b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/core/src/main/java/io/kestra/core/runners/RunContextInitializer.java b/core/src/main/java/io/kestra/core/runners/RunContextInitializer.java
@@ -114,7 +114,7 @@ private DefaultRunContext forWorker(final DefaultRunContext runContext,
         }
 
         final RunContextLogger runContextLogger = contextLoggerFactory.create(taskRun, task);
-        enrichedVariables.put("addSecretConsumer", (Consumer<String>) runContextLogger::usedSecret);
+        enrichedVariables.put(RunVariables.SECRET_CONSUMER_VARIABLE_NAME, (Consumer<String>) runContextLogger::usedSecret);
 
         enrichedVariables = variablesModifier.apply(enrichedVariables);
 
@@ -192,7 +192,7 @@ public DefaultRunContext forScheduler(final DefaultRunContext runContext,
         final RunContextLogger runContextLogger = contextLoggerFactory.create(triggerContext, trigger);
 
         final Map<String, Object> variables = new HashMap<>(runContext.getVariables());
-        variables.put("addSecretConsumer", (Consumer<String>) runContextLogger::usedSecret);
+        variables.put(RunVariables.SECRET_CONSUMER_VARIABLE_NAME, (Consumer<String>) runContextLogger::usedSecret);
 
         final StorageContext context = StorageContext.forTrigger(
             triggerContext.getTenantId(),

diff --git a/core/src/main/java/io/kestra/core/runners/RunVariables.java b/core/src/main/java/io/kestra/core/runners/RunVariables.java
@@ -7,7 +7,6 @@
 import io.kestra.core.models.flows.Input;
 import io.kestra.core.models.flows.input.SecretInput;
 import io.kestra.core.models.tasks.Task;
-import io.kestra.core.models.tasks.common.EncryptedString;
 import io.kestra.core.models.triggers.AbstractTrigger;
 import lombok.AllArgsConstructor;
 import lombok.With;
@@ -18,12 +17,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 /**
  * Class for building {@link RunContext} variables.
  */
 public final class RunVariables {
+    public static final String SECRET_CONSUMER_VARIABLE_NAME = "addSecretConsumer";
 
     /**
      * Creates an immutable map representation of the given {@link Task}.
@@ -273,6 +274,9 @@ public Map<String, Object> build(final RunContextLogger logger) {
             // adds any additional variables
             if (variables != null) {
                 builder.putAll(variables);
+                if (logger != null && !variables.containsKey(RunVariables.SECRET_CONSUMER_VARIABLE_NAME)) {
+                    builder.put(RunVariables.SECRET_CONSUMER_VARIABLE_NAME, (Consumer<String>) logger::usedSecret);
+                }
             }
 
             return builder.build();

diff --git a/core/src/main/java/io/kestra/core/runners/pebble/functions/SecretFunction.java b/core/src/main/java/io/kestra/core/runners/pebble/functions/SecretFunction.java
@@ -1,6 +1,7 @@
 package io.kestra.core.runners.pebble.functions;
 
 import io.kestra.core.exceptions.IllegalVariableEvaluationException;
+import io.kestra.core.runners.RunVariables;
 import io.kestra.core.secret.SecretService;
 import io.pebbletemplates.pebble.error.PebbleException;
 import io.pebbletemplates.pebble.extension.Function;
@@ -36,7 +37,7 @@ public Object execute(Map<String, Object> args, PebbleTemplate self, EvaluationC
             String secret = secretService.findSecret(flow.get("tenantId"), flow.get("namespace"), key);
 
             try {
-                Consumer<String> addSecretConsumer = (Consumer<String>) context.getVariable("addSecretConsumer");
+                Consumer<String> addSecretConsumer = (Consumer<String>) context.getVariable(RunVariables.SECRET_CONSUMER_VARIABLE_NAME);
                 addSecretConsumer.accept(secret);
             } catch (Exception e) {
                 log.warn("Unable to get secret consumer", e);