Project environment variable management commands

tests/commands/project/test_envvars.py

@@ -0,0 +1,69 @@
+from contextlib import contextmanager
+
+import pytest
+import requests_mock
+
+from tests.fixture_data import PROJECT_DATA
+from valohai_cli.commands.project.environment_variables.create import create
+from valohai_cli.commands.project.environment_variables.delete import delete
+from valohai_cli.commands.project.environment_variables.list import list
+
+
+@contextmanager
+def env_var_api_mock():
+    with requests_mock.mock() as m:
+        project_url = f"https://app.valohai.com/api/v0/projects/{PROJECT_DATA['id']}/"
+        m.get(project_url, json=PROJECT_DATA)
+        m.post(f"{project_url}environment_variable/", json={})
+        m.delete(f"{project_url}environment_variable/", json={})
+        yield m
+
+
+def test_list_envvars(runner, logged_in_and_linked):
+    with env_var_api_mock():
+        result = runner.invoke(list)
+    for key, ev in PROJECT_DATA["environment_variables"].items():
+        assert key in result.output
+        assert (ev["value"] if not ev["secret"] else "****") in result.output
+
+
+def test_create_envvar_secret(runner, logged_in_and_linked):
+    with env_var_api_mock() as m:
+        result = runner.invoke(create, ["--secret", "hey=ho=lets=go", "eeuu=oouuuhhh"])
+        assert result.exit_code == 0
+        assert [r.json() for r in m.request_history] == [
+            {"name": "hey", "secret": True, "value": "ho=lets=go"},
+            {"name": "eeuu", "secret": True, "value": "oouuuhhh"},
+        ]
+
+
+def test_create_envvar_nonsecret(runner, logged_in_and_linked):
+    with env_var_api_mock() as m:
+        result = runner.invoke(create, ["BANANA_HAMMOCK_ENABLED=true"])
+        print(result.output)
+        assert result.exit_code == 0
+        assert [r.json() for r in m.request_history] == [
+            {"name": "BANANA_HAMMOCK_ENABLED", "secret": False, "value": "true"},
+        ]
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        "foo",
+        "^38=fbar",
+        "3aa=kdkd",
+    ],
+)
+def test_create_envvar_bad(runner, logged_in_and_linked, args):
+    result = runner.invoke(create, args)
+    assert result.exit_code != 0
+
+
+def test_delete_envvar(runner, logged_in_and_linked):
+    with env_var_api_mock() as m:
+        result = runner.invoke(delete, ["hey"])
+        assert result.exit_code == 0
+        assert [r.json() for r in m.request_history] == [
+            {"name": "hey"},
+        ]

tests/fixture_data.py

@@ -19,6 +19,16 @@
     "urls": {
         "display": "https://app.valohai.com/p/nyan/nyan/",
     },
+    "environment_variables": {
+        "test1": {
+            "value": "val1",
+            "secret": False,
+        },
+        "ssshhh": {
+            "value": None,
+            "secret": True,
+        },
+    },
 }
 
 execution_id = str(uuid.uuid4())

valohai_cli/commands/project/environment_variables/__init__.py

@@ -0,0 +1,10 @@
+import click
+
+from valohai_cli.plugin_cli import PluginCLI
+
+
+@click.command(cls=PluginCLI, commands_module="valohai_cli.commands.project.environment_variables")
+def environment_variables() -> None:
+    """
+    Project environment variable related commands.
+    """

valohai_cli/commands/project/environment_variables/create.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import re
+from typing import Any
+
+import click
+
+from valohai_cli.api import request
+from valohai_cli.ctx import get_project
+from valohai_cli.messages import success, warn
+
+VALID_ENVVAR_NAME_RE = re.compile(r"^[a-z0-9_-]+$", re.I)
+
+
+def parse_env_variable(s: str) -> tuple[str, str]:
+    """
+    Parse an environment variable string of the form NAME=VALUE.
+    """
+    if "=" not in s:
+        raise click.BadParameter("Environment variable must be in the form NAME=VALUE")
+    key, _, value = s.partition("=")
+    key = key.strip()
+    if not VALID_ENVVAR_NAME_RE.match(key):
+        raise click.BadParameter(
+            f"Environment variable name must be alphanumeric-and-dashes-and-underscores, got {key}",
+        )
+    return key, value.strip()
+
+
+@click.command()
+@click.argument(
+    "environment_variables",
+    metavar="NAME=VALUE",
+    nargs=-1,
+    required=True,
+    type=parse_env_variable,
+)
+@click.option("--secret/--public", help="Set environment variable as a secret. Default value is --public")
+def create(
+    *,
+    environment_variables: list[tuple[str, str]],
+    secret: bool,
+) -> Any:
+    """
+    Add an environment variable to a project
+    """
+    project = get_project(require=True)
+
+    if not environment_variables:  # pragma: no cover
+        warn("Nothing to do.")
+        return
+
+    for key, value in environment_variables:
+        payload = {
+            "name": key,
+            "value": value,
+            "secret": secret,
+        }
+
+        request(
+            method="post",
+            url=f"/api/v0/projects/{project.id}/environment_variable/",
+            json=payload,
+        )
+
+    if len(environment_variables) == 1:
+        success(f"Added environment variable {environment_variables[0][0]} to project {project.name}.")
+    else:
+        success(f"Added {len(environment_variables)} environment variables to project {project.name}.")

valohai_cli/commands/project/environment_variables/delete.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from typing import Any
+
+import click
+from click.exceptions import Exit
+
+from valohai_cli.api import request
+from valohai_cli.ctx import get_project
+from valohai_cli.exceptions import APINotFoundError
+from valohai_cli.messages import error, success
+
+
+@click.command()
+@click.argument("names", metavar="NAME", nargs=-1, required=True)
+def delete(*, names: list[str]) -> Any:
+    """
+    Delete one or more environment variables.
+    """
+
+    project = get_project(require=True)
+    fail = False
+    for name in names:
+        payload = {"name": name}
+
+        try:
+            request(method="delete", url=f"/api/v0/projects/{project.id}/environment_variable/", json=payload)
+        except APINotFoundError:  # pragma: no cover
+            error(f"Environment variable ({name}) not found in {project.name}")
+            fail = True
+        else:
+            success(f"Successfully deleted environment variable {name} from {project.name}")
+
+    if fail:
+        raise Exit(1)

valohai_cli/commands/project/environment_variables/list.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import Any
+
+import click
+
+from valohai_cli.api import request
+from valohai_cli.ctx import get_project
+from valohai_cli.messages import info
+from valohai_cli.settings import settings
+from valohai_cli.table import print_json, print_table
+
+
+@click.command()
+def list() -> Any:
+    """
+    List environment variables of a project
+    """
+
+    project = get_project(require=True)
+    data = request(method="get", url=f"/api/v0/projects/{project.id}/").json()["environment_variables"]
+
+    if settings.output_format == "json":
+        return print_json(data)
+
+    if not data:  # pragma: no cover
+        info(f"{project}: No environment variables.")
+        return
+
+    formatted_envvars = [
+        {
+            "name": name,
+            "value": "****" if info["secret"] else info["value"],
+            "secret": info["secret"],
+        }
+        for name, info in sorted(data.items())
+    ]
+
+    print_table(
+        formatted_envvars,
+        columns=["name", "value", "secret"],
+        headers=["Name", "Value", "Secret"],
+    )

valohai_cli/plugin_cli.py

@@ -16,6 +16,7 @@
 
 class PluginCLI(click.MultiCommand):
     aliases = {
+        "add": "create",
         "new": "create",
         "start": "run",
     }
@@ -138,7 +139,7 @@ def get_help(self, ctx: Context) -> str:
             # (see https://github.com/pallets/click/pull/1623).
             import json
 
-            return json.dumps(ctx.to_info_dict())
+            return json.dumps(ctx.to_info_dict(), default=str)
 
         return super().get_help(ctx)