Skip to content

Commit

Permalink
feat: add env_vars as input to run_sh (#2114)
Browse files Browse the repository at this point in the history
## Description:
This change adds `env_vars` to the run_sh. Feature was requested in
issue #2050. Example of used in practice:

```
def run(plan, args={}):
    result = plan.run_sh(
        run = "mkdir -p kurtosis && echo $EXAMPLE",
        image = "badouralix/curl-jq",
        env_vars = {
            "EXAMPLE": "this_is_a_test"
        },
    )
```
outputs:

```
Container images used in this run:
> badouralix/curl-jq - locally cached

> run_sh run="mkdir -p kurtosis && echo $EXAMPLE" image="badouralix/curl-jq" env_vars={"EXAMPLE": "this_is_a_test"}
Command returned with exit code '0' and the following output:
--------------------
this_is_a_test

--------------------

Starlark code successfully run. No output was returned.
```


## Is this change user facing?
YES

## References (if applicable):
Closes: #2050
  • Loading branch information
adschwartz committed Feb 5, 2024
1 parent 1f73821 commit 5a30ea7
Show file tree
Hide file tree
Showing 6 changed files with 64 additions and 5 deletions.
Expand Up @@ -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)
}
Expand Down
Expand Up @@ -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,
Expand Down Expand Up @@ -90,6 +96,7 @@ func NewRunShService(serviceNetwork service_network.ServiceNetwork, runtimeValue
FilesArgName: true,
StoreFilesArgName: true,
WaitArgName: true,
EnvVarsArgName: true,
},
}
}
Expand Down Expand Up @@ -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)
}
Expand Down
Expand Up @@ -31,6 +31,7 @@ const (
StoreFilesArgName = "store"
WaitArgName = "wait"
FilesArgName = "files"
EnvVarsArgName = "env_vars"

newlineChar = "\n"

Expand Down Expand Up @@ -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,
Expand All @@ -262,7 +267,7 @@ func getServiceConfig(image string, filesArtifactExpansion *service_directory.Fi
// command is completed
runTailCommandToPreventContainerToStopOnCreating,
nil,
nil,
*envVars,
filesArtifactExpansion,
nil,
0,
Expand Down Expand Up @@ -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
}
7 changes: 7 additions & 0 deletions docs/docs/api-reference/starlark-reference/plan.md
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion internal_testsuites/golang/scripts/test.sh
Expand Up @@ -12,7 +12,7 @@ lang_root_dirpath="$(dirname "${script_dirpath}")"
# ==================================================================================================
PARALLELISM=2
DOCKER_TIMEOUT="3m" # This must be Go-parseable timeout
KUBERNETES_TIMEOUT="8m" # K8S takes longer than docker
KUBERNETES_TIMEOUT="510s" # K8S takes longer than docker

TESTSUITE_CLUSTER_BACKEND_DOCKER="docker"
TESTSUITE_CLUSTER_BACKEND_KUBERNETES="kubernetes"
Expand Down
Expand Up @@ -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) {
Expand Down Expand Up @@ -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)
}

0 comments on commit 5a30ea7

Please sign in to comment.