Enable SecretsManager.get to load and return bytes (flyteorg#1798)

* fix secretsmanager Signed-off-by: Yue Shang <[email protected]> * fix lint issue Signed-off-by: Yue Shang <[email protected]> * add doc Signed-off-by: Yue Shang <[email protected]> * fix github check Signed-off-by: Yue Shang <[email protected]> --------- Signed-off-by: Yue Shang <[email protected]> Signed-off-by: Future Outlier <[email protected]>
Future-Outlier · Oct 3, 2023 · a8435ea · a8435ea
1 parent 154c43c
commit a8435ea
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git 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
@@ -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"]