Support GCP secrets (#1571)

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

flytekit/core/context_manager.py

@@ -329,13 +329,13 @@ def __getattr__(self, item: str) -> _GroupSecrets:
         """
         return self._GroupSecrets(item, self)
 
-    def get(self, group: str, key: str) -> str:
+    def get(self, group: str, key: Optional[str] = None, group_version: Optional[str] = None) -> str:
         """
         Retrieves a secret using the resolution order -> Env followed by file. If not found raises a ValueError
         """
-        self.check_group_key(group, key)
-        env_var = self.get_secrets_env_var(group, key)
-        fpath = self.get_secrets_file(group, key)
+        self.check_group_key(group)
+        env_var = self.get_secrets_env_var(group, key, group_version)
+        fpath = self.get_secrets_file(group, key, group_version)
         v = os.environ.get(env_var)
         if v is not None:
             return v
@@ -346,26 +346,27 @@ def get(self, group: str, key: str) -> str:
             f"Unable to find secret for key {key} in group {group} " f"in Env Var:{env_var} and FilePath: {fpath}"
         )
 
-    def get_secrets_env_var(self, group: str, key: str) -> str:
+    def get_secrets_env_var(self, group: str, key: Optional[str] = None, group_version: Optional[str] = None) -> str:
         """
         Returns a string that matches the ENV Variable to look for the secrets
         """
-        self.check_group_key(group, key)
-        return f"{self._env_prefix}{group.upper()}_{key.upper()}"
+        self.check_group_key(group)
+        l = [k.upper() for k in filter(None, (group, group_version, key))]
+        return f"{self._env_prefix}{'_'.join(l)}"
 
-    def get_secrets_file(self, group: str, key: str) -> str:
+    def get_secrets_file(self, group: str, key: Optional[str] = None, group_version: Optional[str] = None) -> str:
         """
         Returns a path that matches the file to look for the secrets
         """
-        self.check_group_key(group, key)
-        return os.path.join(self._base_dir, group.lower(), f"{self._file_prefix}{key.lower()}")
+        self.check_group_key(group)
+        l = [k.lower() for k in filter(None, (group, group_version, key))]
+        l[-1] = f"{self._file_prefix}{l[-1]}"
+        return os.path.join(self._base_dir, *l)
 
     @staticmethod
-    def check_group_key(group: str, key: str):
+    def check_group_key(group: str):
         if group is None or group == "":
             raise ValueError("secrets group is a mandatory field.")
-        if key is None or key == "":
-            raise ValueError("secrets key is a mandatory field.")
 
 
 @dataclass(frozen=True)

flytekit/models/security.py

@@ -36,15 +36,13 @@ class MountType(Enum):
         """
 
     group: str
-    key: str
+    key: Optional[str] = None
     group_version: Optional[str] = None
     mount_requirement: MountType = MountType.ANY
 
     def __post_init__(self):
         if self.group is None:
             raise ValueError("Group is a required parameter")
-        if self.key is None:
-            raise ValueError("Key is also a required parameter")
 
     def to_flyte_idl(self) -> _sec.Secret:
         return _sec.Secret(
@@ -59,7 +57,7 @@ def from_flyte_idl(cls, pb2_object: _sec.Secret) -> "Secret":
         return cls(
             group=pb2_object.group,
             group_version=pb2_object.group_version if pb2_object.group_version else None,
-            key=pb2_object.key,
+            key=pb2_object.key if pb2_object.key else None,
             mount_requirement=Secret.MountType(pb2_object.mount_requirement),
         )
 

tests/flytekit/unit/core/test_context_manager.py

@@ -115,18 +115,17 @@ def test_secrets_manager_default():
 
 def test_secrets_manager_get_envvar():
     sec = SecretsManager()
-    with pytest.raises(ValueError):
-        sec.get_secrets_env_var("test", "")
     with pytest.raises(ValueError):
         sec.get_secrets_env_var("", "x")
     cfg = SecretsConfig.auto()
     assert sec.get_secrets_env_var("group", "test") == f"{cfg.env_prefix}GROUP_TEST"
+    assert sec.get_secrets_env_var("group", "test", "v1") == f"{cfg.env_prefix}GROUP_V1_TEST"
+    assert sec.get_secrets_env_var("group", group_version="v1") == f"{cfg.env_prefix}GROUP_V1"
+    assert sec.get_secrets_env_var("group") == f"{cfg.env_prefix}GROUP"
 
 
 def test_secrets_manager_get_file():
     sec = SecretsManager()
-    with pytest.raises(ValueError):
-        sec.get_secrets_file("test", "")
     with pytest.raises(ValueError):
         sec.get_secrets_file("", "x")
     cfg = SecretsConfig.auto()
@@ -135,6 +134,12 @@ def test_secrets_manager_get_file():
         "group",
         f"{cfg.file_prefix}test",
     )
+    assert sec.get_secrets_file("group", "test", "v1") == os.path.join(
+        cfg.default_dir,
+        "group",
+        "v1",
+        f"{cfg.file_prefix}test",
+    )
 
 
 def test_secrets_manager_file(tmpdir: py.path.local):
@@ -145,8 +150,6 @@ def test_secrets_manager_file(tmpdir: py.path.local):
     with open(f, "w+") as w:
         w.write("my-password")
 
-    with pytest.raises(ValueError):
-        sec.get("test", "")
     with pytest.raises(ValueError):
         sec.get("", "x")
     # Group dir not exists

tests/flytekit/unit/models/core/test_security.py

@@ -0,0 +1,13 @@
+from flytekit.models.security import Secret
+
+
+def test_secret():
+    obj = Secret("grp", "key")
+    obj2 = Secret.from_flyte_idl(obj.to_flyte_idl())
+    assert obj2.key == "key"
+    assert obj2.group_version is None
+
+    obj = Secret("grp", group_version="v1")
+    obj2 = Secret.from_flyte_idl(obj.to_flyte_idl())
+    assert obj2.key is None
+    assert obj2.group_version == "v1"