diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index aa6b0e3e4d..3762dec92d 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -349,9 +349,12 @@ def __getattr__(self, item: str) -> _GroupSecrets: """ return self._GroupSecrets(item, self) - def get(self, group: str, key: Optional[str] = None, group_version: Optional[str] = None) -> str: + def get( + self, group: str, key: Optional[str] = None, group_version: Optional[str] = None, encode_mode: str = "r" + ) -> str: """ Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError + param encode_mode, defines the mode to open files, it can either be "r" to read file, or "rb" to read binary file """ self.check_group_key(group) env_var = self.get_secrets_env_var(group, key, group_version) @@ -360,7 +363,7 @@ def get(self, group: str, key: Optional[str] = None, group_version: Optional[str if v is not None: return v if os.path.exists(fpath): - with open(fpath, "r") as f: + with open(fpath, encode_mode) as f: return f.read().strip() raise ValueError( f"Unable to find secret for key {key} in group {group} " f"in Env Var:{env_var} and FilePath: {fpath}" diff --git a/tests/flytekit/unit/core/test_context_manager.py b/tests/flytekit/unit/core/test_context_manager.py index b7932663a5..2ec7eb8e19 100644 --- a/tests/flytekit/unit/core/test_context_manager.py +++ b/tests/flytekit/unit/core/test_context_manager.py @@ -1,3 +1,4 @@ +import base64 import os from datetime import datetime @@ -166,6 +167,15 @@ def test_secrets_manager_file(tmpdir: py.path.local): w.write("my-password") assert sec.get("group", "test") == "my-password" assert sec.group.test == "my-password" + + base64_string = "R2Vla3NGb3JHZWV ==" + base64_bytes = base64_string.encode("ascii") + base64_str = base64.b64encode(base64_bytes) + with open(f, "wb") as w: + w.write(base64_str) + assert sec.get("group", "test") != base64_str + assert sec.get("group", "test", encode_mode="rb") == base64_str + del os.environ["FLYTE_SECRETS_DEFAULT_DIR"]