find env based on parameter env_var

Signed-off-by: yuteng <[email protected]>

Author: 0yukali0

flytekit/core/context_manager.py

@@ -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:
         """
@@ -413,20 +414,26 @@ def get(
             env_prefixes.append("")
 
         for env_prefix in env_prefixes:
-            env_var = self._get_secrets_env_var(
+            group_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(group_env_var)
+            if v is not None:
+                return v.strip()
+
+        if env_var is not None:
+            v = os.environ.get(env_var, None)
             if v is not None:
                 return v.strip()
 
         fpath = self.get_secrets_file(group, key, group_version)
         if os.path.exists(fpath):
             with open(fpath, encode_mode) as f:
                 return f.read().strip()
+        env_vars = ",".join([k for k in filter(None, (group_env_var, env_var))])
         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:{env_vars} and FilePath: {fpath}"
         )
 
     def get_secrets_env_var(
@@ -443,7 +450,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)}"
 

tests/flytekit/unit/core/test_context_manager.py

@@ -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):