diff --git a/latch/functions/secrets.py b/latch/functions/secrets.py new file mode 100644 index 00000000..c48c149c --- /dev/null +++ b/latch/functions/secrets.py @@ -0,0 +1,53 @@ +import os + +from latch_cli.config.latch import config +from latch_cli.tinyrequests import post +from latch_cli.utils import current_workspace, retrieve_or_login + + +def get_secret(secret_name: str): + """ + A utility to allow users to reference secrets stored in their workspace on + Latch. + + Important: When running an execution locally, whether on your own computer + or using `latch develop`, the only secrets you will be able to access are + the ones in your personal workspace. To use secrets from a shared workspace, + register your workflow and run it on Latch. + + Examples: + >>> get_secret("test-secret") + "test-value-123" + """ + execution_token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID") + if execution_token is None: + return _get_secret_local(secret_name) + + resp = post( + url=config.api.user.get_secret, + json={ + "execution_token": execution_token, + "name": secret_name, + }, + ) + + if resp.status_code != 200: + raise ValueError(resp.json()["error"]) + + return resp.json()["secret"] + + +def _get_secret_local(secret_name: str): + resp = post( + url=config.api.user.get_secret_local, + json={ + "ws_account_id": current_workspace(), + "name": secret_name, + }, + headers={"Authorization": f"Bearer {retrieve_or_login()}"}, + ) + + if resp.status_code != 200: + raise ValueError(resp.json()["error"]) + + return resp.json()["secret"] diff --git a/latch_cli/config/latch.py b/latch_cli/config/latch.py index 392f2bf2..7866b70d 100644 --- a/latch_cli/config/latch.py +++ b/latch_cli/config/latch.py @@ -55,6 +55,8 @@ class _ExecutionAPI: class _UserAPI: jwt: str = "/sdk/access-jwt" list_workspaces: str = "/sdk/get-ws" + get_secret: str = "/secrets/get" + get_secret_local: str = "/secrets/get-local" @dataclass