diff --git a/arcade/arcade/worker/config/deployment.py b/arcade/arcade/worker/config/deployment.py index 38cebc6f..34691b44 100644 --- a/arcade/arcade/worker/config/deployment.py +++ b/arcade/arcade/worker/config/deployment.py @@ -1,6 +1,7 @@ import base64 import io import os +import re import secrets import tarfile from pathlib import Path @@ -11,7 +12,7 @@ import toml from arcadepy import Arcade, NotFoundError from httpx import Client from packaging.requirements import Requirement -from pydantic import BaseModel, field_validator, model_validator +from pydantic import BaseModel, field_serializer, field_validator, model_validator # Base class for versioned packages @@ -63,26 +64,47 @@ class Pypi(PackageRepository): trusted_host: str = "pypi.org" +class Secret(BaseModel): + value: str + pattern: str | None = None + + class Config(BaseModel): id: str enabled: bool = True timeout: int = 30 retries: int = 3 - secret: str + secret: Secret | None = None - # Validate that the secret is a non-empty string and not 'dev' - @field_validator("secret") + # Validate and parse the secret if required + @field_validator("secret", mode="before") @classmethod - def valid_secret(cls, v: str) -> str: - if v.strip("") == "" or v == "dev": + def valid_secret(cls, v: str | Secret | None) -> Secret: + # If the secret is a string, attempt to parse it as an environment variable or return the secret + if isinstance(v, str): + secret = get_env_secret(v) + # If the secret has been manually set, return it + elif isinstance(v, Secret): + secret = v + else: + raise TypeError("Secret must be a string or a Secret object") + # Check that the secret is not the default dev secret or empty + if secret.value.strip() == "" or secret.value == "dev": raise ValueError("Secret must be a non-empty string and not 'dev'") - return v + return secret + + @field_serializer("secret") + def serialize_secret(self, secret: Secret) -> str: + if secret.pattern: + return f"$env:{secret.pattern}" + else: + return secret.value # Cloud request for deploying a worker class Request(BaseModel): name: str - secret: str + secret: Secret enabled: bool timeout: int retries: int @@ -90,6 +112,10 @@ class Request(BaseModel): custom_repositories: list[PackageRepository] | None = None local_packages: list[LocalPackage] | None = None + @field_serializer("secret") + def serialize_secret(self, secret: Secret) -> str: + return secret.value + def execute(self, cloud_client: Client, engine_client: Arcade) -> Any: # Attempt to deploy worker to the cloud try: @@ -113,7 +139,7 @@ class Request(BaseModel): enabled=self.enabled, http={ "uri": cloud_response.json()["data"]["worker_endpoint"], - "secret": self.secret, + "secret": self.secret.value, "timeout": self.timeout, "retry": self.retries, }, @@ -125,7 +151,7 @@ class Request(BaseModel): enabled=self.enabled, http={ "uri": cloud_response.json()["data"]["worker_endpoint"], - "secret": self.secret, + "secret": self.secret.value, "timeout": self.timeout, "retry": self.retries, }, @@ -148,6 +174,8 @@ class Worker(BaseModel): """Convert Deployment to a Request object.""" self.validate_packages() self.compress_local_packages() + if self.config.secret is None: + raise ValueError("Secret is required") return Request( name=self.config.id, secret=self.config.secret, @@ -274,7 +302,7 @@ def create_demo_deployment(toml_path: Path, toolkit_name: str) -> None: enabled=True, timeout=30, retries=3, - secret=secrets.token_hex(16), + secret=Secret(value=secrets.token_hex(16), pattern=None), ), local_source=LocalPackages(packages=[f"./{toolkit_name}"]), ) @@ -292,3 +320,25 @@ def update_deployment_with_local_packages(toml_path: Path, toolkit_name: str) -> else: deployment.worker[0].local_source.packages.append(f"./{toolkit_name}") deployment.save() + + +def get_env_secret(secret: str) -> Secret: + """Parse a secret from an environment variable.""" + # Check if the secret contains the "${env:}" syntax + pattern = r"\${env:([^}]+)}" + matches = re.findall(pattern, secret) + + # Only allow a single match + if matches and len(matches) == 1: + match = matches[0].strip() + # Attempt to lookup and create the secret + print(f"Looking up secret: {match}") + value = os.getenv(match) + if value: + return Secret(value=value, pattern=match) + else: + raise ValueError(f"Environment variable not found: {match}") + elif matches and len(matches) > 1: + raise ValueError(f"Multiple environment variables found in secret: {secret}") + # If no matches are found, return the secret as is + return Secret(value=secret, pattern=None) diff --git a/arcade/tests/deployment/test_config.py b/arcade/tests/deployment/test_config.py index 923626ff..531bfaac 100644 --- a/arcade/tests/deployment/test_config.py +++ b/arcade/tests/deployment/test_config.py @@ -2,6 +2,7 @@ # ruff: noqa: S105 # ruff: noqa: S106 import json +import os from pathlib import Path import pytest @@ -13,6 +14,7 @@ from arcade.worker.config.deployment import ( Package, PackageRepository, Pypi, + Secret, Worker, ) @@ -41,7 +43,7 @@ def test_deployment_parsing(test_dir): assert deployment.worker[0].config.enabled is True assert deployment.worker[0].config.timeout == 10 assert deployment.worker[0].config.retries == 3 - assert deployment.worker[0].config.secret == "test-secret" + assert deployment.worker[0].config.secret == Secret(value="test-secret", pattern=None) # Test pypi section assert deployment.worker[0].pypi_source.packages == [Package(name="arcade-x")] @@ -123,7 +125,6 @@ def test_deployment_dict(test_dir): ] }""") got = deployment.worker[0].request().model_dump(mode="json") - print(got) # Remove encoding part that contains the content got["local_packages"][0].pop("content") expected["local_packages"][0].pop("content") @@ -161,7 +162,7 @@ def test_unconfigured_local_package(test_dir): def test_duplicate_pypi_packages(): worker = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), pypi_source=Pypi(packages=["arcade-slack", "arcade-slack"]), ) with pytest.raises(ValueError): @@ -171,7 +172,7 @@ def test_duplicate_pypi_packages(): def test_duplicate_custom_repository_packages(): worker = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), custom_source=[ PackageRepository( index="pypi", @@ -188,7 +189,7 @@ def test_duplicate_custom_repository_packages(): def test_duplicate_local_packages(): worker = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), local_source=LocalPackages(packages=["./mock_toolkit", "./mock_toolkit"]), ) with pytest.raises(ValueError): @@ -198,7 +199,7 @@ def test_duplicate_local_packages(): def test_duplicate_all_typed_packages(): worker = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), pypi_source=Pypi(packages=["arcade-slack"]), custom_source=[ PackageRepository( @@ -217,11 +218,19 @@ def test_duplicate_all_typed_packages(): def test_duplicate_worker_names(): worker = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), ) worker2 = Worker( toml_path=Path(__file__), - config=Config(id="test", secret="test-secret"), + config=Config(id="test", secret=Secret(value="test-secret", pattern=None)), ) with pytest.raises(ValueError): Deployment(workers=[worker, worker2]) + + +def test_secret_parsing(test_dir): + os.environ["TEST_WORKER_SECRET"] = "test-secret" + deployment = Deployment.from_toml(test_dir / "test_files" / "env.secret.worker.toml") + assert deployment.worker[0].config.secret == Secret( + value="test-secret", pattern="TEST_WORKER_SECRET" + ) diff --git a/arcade/tests/deployment/test_files/env.secret.worker.toml b/arcade/tests/deployment/test_files/env.secret.worker.toml new file mode 100644 index 00000000..309f61af --- /dev/null +++ b/arcade/tests/deployment/test_files/env.secret.worker.toml @@ -0,0 +1,5 @@ + +[[worker]] +[worker.config] +id = "test" +secret = "${env: TEST_WORKER_SECRET}"