Skip to content

[Bug] SecretsManager get function supports env_var pararmeters. #3201

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
+17 −4
Open

[Bug] SecretsManager get function supports env_var pararmeters. #3201

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions flytekit/core/context_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,7 @@ def get(
group: Optional[str] = None,
key: Optional[str] = None,
group_version: Optional[str] = None,
env_var: Optional[str] = None,
encode_mode: str = "r",
) -> str:
"""
Expand All @@ -412,11 +413,18 @@ def get(
if ctx.execution_state is None or ctx.execution_state.is_local_execution():
env_prefixes.append("")

default_env_var = ""
for env_prefix in env_prefixes:
env_var = self._get_secrets_env_var(
default_env_var = self._get_secrets_env_var(
group=group, key=key, group_version=group_version, env_prefix=env_prefix
)
v = os.environ.get(env_var)
v = os.environ.get(default_env_var)
if v is not None:
return v.strip()

custom_env_var = env_var
if custom_env_var is not None:
v = os.environ.get(custom_env_var, None)
if v is not None:
return v.strip()

Expand All @@ -426,7 +434,7 @@ def get(
return f.read().strip()
raise ValueError(
f"Please make sure to add secret_requests=[Secret(group={group}, key={key})] in @task. Unable to find secret for key {key} in group {group} "
f"in Env Var:{env_var} and FilePath: {fpath}"
f"in Env Var:{default_env_var}{', ' + custom_env_var if custom_env_var is not None else ''} and FilePath: {fpath}"
)

def get_secrets_env_var(
Expand All @@ -443,7 +451,7 @@ def _get_secrets_env_var(
key: Optional[str] = None,
group_version: Optional[str] = None,
env_prefix: str = "",
):
) -> str:
l = [k.upper() for k in filter(None, (group, group_version, key))]
return f"{env_prefix}{'_'.join(l)}"

Expand Down
5 changes: 5 additions & 0 deletions tests/flytekit/unit/core/test_context_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,11 @@ def test_secrets_manager_env():
os.environ[sec.get_secrets_env_var(group="group", key="key")] = "value"
assert sec.get(group="group", key="key") == "value"

os.environ["LOCAL_ENV_VAR"] = "value"
assert sec.get(group="group", key="key2", env_var="LOCAL_ENV_VAR") == "value"
assert sec.get(key="key", env_var="LOCAL_ENV_VAR") == "value"
assert sec.get(env_var="LOCAL_ENV_VAR") == "value"


@pytest.mark.parametrize("is_local_execution, prefix", [(True, ""), (False, "_FSEC_")])
def test_secrets_manager_execution(monkeypatch, is_local_execution, prefix):
Expand Down
Loading