diff --git a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_python.go b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_python.go index 04f982e1ce..9633e36c51 100644 --- a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_python.go +++ b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_python.go @@ -203,8 +203,13 @@ func (builtin *RunPythonCapabilities) Interpret(_ string, arguments *builtin_arg } } + envVars, interpretationErr := extractEnvVarsIfDefined(arguments) + if err != nil { + return nil, interpretationErr + } + // build a service config from image and files artifacts expansion. - builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion) + builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion, envVars) if err != nil { return nil, startosis_errors.WrapWithInterpretationError(err, "An error occurred creating service config using image '%s'", image) } diff --git a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_sh.go b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_sh.go index 8c7b060651..d73e062c35 100644 --- a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_sh.go +++ b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/run_sh.go @@ -58,6 +58,12 @@ func NewRunShService(serviceNetwork service_network.ServiceNetwork, runtimeValue IsOptional: true, ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.List], }, + { + Name: EnvVarsArgName, + IsOptional: true, + ZeroValueProvider: builtin_argument.ZeroValueProvider[*starlark.Dict], + Validator: nil, + }, { Name: WaitArgName, IsOptional: true, @@ -90,6 +96,7 @@ func NewRunShService(serviceNetwork service_network.ServiceNetwork, runtimeValue FilesArgName: true, StoreFilesArgName: true, WaitArgName: true, + EnvVarsArgName: true, }, } } @@ -147,8 +154,13 @@ func (builtin *RunShCapabilities) Interpret(_ string, arguments *builtin_argumen } } + envVars, interpretationErr := extractEnvVarsIfDefined(arguments) + if err != nil { + return nil, interpretationErr + } + // build a service config from image and files artifacts expansion. - builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion) + builtin.serviceConfig, err = getServiceConfig(image, filesArtifactExpansion, envVars) if err != nil { return nil, startosis_errors.WrapWithInterpretationError(err, "An error occurred creating service config using image '%s'", image) } diff --git a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/tasks_shared.go b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/tasks_shared.go index e8addad386..5e546a51f8 100644 --- a/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/tasks_shared.go +++ b/core/server/api_container/server/startosis_engine/kurtosis_instruction/tasks/tasks_shared.go @@ -31,6 +31,7 @@ const ( StoreFilesArgName = "store" WaitArgName = "wait" FilesArgName = "files" + EnvVarsArgName = "env_vars" newlineChar = "\n" @@ -248,7 +249,11 @@ func resultMapToString(resultMap map[string]starlark.Comparable, builtinNameForL return fmt.Sprintf("Command returned with exit code '%v' and the following output: %v", exitCode, outputStr) } -func getServiceConfig(image string, filesArtifactExpansion *service_directory.FilesArtifactsExpansion) (*service.ServiceConfig, error) { +func getServiceConfig( + image string, + filesArtifactExpansion *service_directory.FilesArtifactsExpansion, + envVars *map[string]string, +) (*service.ServiceConfig, error) { serviceConfig, err := service.CreateServiceConfig( image, nil, @@ -262,7 +267,7 @@ func getServiceConfig(image string, filesArtifactExpansion *service_directory.Fi // command is completed runTailCommandToPreventContainerToStopOnCreating, nil, - nil, + *envVars, filesArtifactExpansion, nil, 0, @@ -293,3 +298,21 @@ func removeService(ctx context.Context, serviceNetwork service_network.ServiceNe } return nil } + +func extractEnvVarsIfDefined(arguments *builtin_argument.ArgumentValuesSet) (*map[string]string, *startosis_errors.InterpretationError) { + envVars := map[string]string{} + if arguments.IsSet(EnvVarsArgName) { + envVarsStarlark, err := builtin_argument.ExtractArgumentValue[*starlark.Dict](arguments, EnvVarsArgName) + if err != nil { + return nil, startosis_errors.WrapWithInterpretationError(err, "Unable to extract value for '%s' argument", EnvVarsArgName) + } + if envVarsStarlark != nil && envVarsStarlark.Len() > 0 { + var interpretationErr *startosis_errors.InterpretationError + envVars, interpretationErr = kurtosis_types.SafeCastToMapStringString(envVarsStarlark, EnvVarsArgName) + if interpretationErr != nil { + return nil, interpretationErr + } + } + } + return &envVars, nil +} diff --git a/docs/docs/api-reference/starlark-reference/plan.md b/docs/docs/api-reference/starlark-reference/plan.md index ee8ec2f629..34d18ee71e 100644 --- a/docs/docs/api-reference/starlark-reference/plan.md +++ b/docs/docs/api-reference/starlark-reference/plan.md @@ -477,6 +477,13 @@ The `run_sh` instruction executes a one-time execution task. It runs the bash co # OPTIONAL (Default: badouralix/curl-jq) image = "badouralix/curl-jq", + # Defines environment variables that should be set inside the Docker container running the task. + # OPTIONAL (Default: {}) + env_vars = { + "VAR_1": "VALUE_1", + "VAR_2": "VALUE_2", + }, + # A mapping of path_on_task_where_contents_will_be_mounted -> files_artifact_id_to_mount # For more information about file artifacts, see below. # CAUTION: duplicate paths to files or directories to be mounted is not supported, and it will fail diff --git a/internal_testsuites/golang/testsuite/startosis_run_sh_task_test/run_task_sh_task_test.go b/internal_testsuites/golang/testsuite/startosis_run_sh_task_test/run_task_sh_task_test.go index 323dcfa25d..6b12a8d4ef 100644 --- a/internal_testsuites/golang/testsuite/startosis_run_sh_task_test/run_task_sh_task_test.go +++ b/internal_testsuites/golang/testsuite/startosis_run_sh_task_test/run_task_sh_task_test.go @@ -55,6 +55,12 @@ def run(plan): result2 = plan.run_sh(run="cat /temp/tech.txt | tr -d '\n'", files={"/temp": file_artifacts[0]}) plan.verify(result2.output, "==", "kurtosis") ` + + runShWithEnvVar = ` +def run(plan): + result = plan.run_sh(run="mkdir -p kurtosis && echo $EXAMPLE",image="badouralix/curl-jq",env_vars={"EXAMPLE": "value"}) + plan.verify(result.output, "==", "value\n") +` ) func TestStarlark_RunshTaskSimple(t *testing.T) { @@ -101,3 +107,9 @@ func TestStarlark_RunShWithNewLineRemovalPipe(t *testing.T) { _, err := test_helpers.SetupSimpleEnclaveAndRunScript(t, ctx, runshTest, runShWithNewLineRemoval) require.Nil(t, err) } + +func TestStarlark_RunShWithEnvVars(t *testing.T) { + ctx := context.Background() + _, err := test_helpers.SetupSimpleEnclaveAndRunScript(t, ctx, runshTest, runShWithEnvVar) + require.Nil(t, err) +}