Updates Secret to use env_name field as environment variable name on …

…injection if exists

Signed-off-by: Geert Pingen <[email protected]>

flytekit/core/context_manager.py

@@ -349,32 +349,51 @@ 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: Optional[str] = None,
+        key: Optional[str] = None,
+        group_version: Optional[str] = None,
+        env_name: 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)
-        env_var = self.get_secrets_env_var(group, key, group_version)
-        fpath = self.get_secrets_file(group, key, group_version)
+        self.check_env_name_key(env_name)
+        env_var = self.get_secrets_env_var(group, key, group_version, env_name)
+        if env_name is None:
+            fpath = self.get_secrets_file(group, key, group_version)
         v = os.environ.get(env_var)
         if v is not None:
             return v
         if os.path.exists(fpath):
             with open(fpath, "r") 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}"
+            f"Unable to find secret for key {key} in group {group} or name {env_name}"
+            f"in Env Var:{env_var} and FilePath: {fpath}"
         )
 
-    def get_secrets_env_var(self, group: str, key: Optional[str] = None, group_version: Optional[str] = None) -> str:
+    def get_secrets_env_var(
+        self,
+        group: Optional[str] = None,
+        key: Optional[str] = None,
+        group_version: Optional[str] = None,
+        env_name: Optional[str] = None,
+    ) -> str:
         """
         Returns a string that matches the ENV Variable to look for the secrets
         """
+        self.check_env_name_key(env_name)
+        if env_name is not None:
+            return env_name
         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: Optional[str] = None, group_version: Optional[str] = None) -> str:
+    def get_secrets_file(
+        self, group: Optional[str] = None, key: Optional[str] = None, group_version: Optional[str] = None
+    ) -> str:
         """
         Returns a path that matches the file to look for the secrets
         """
@@ -384,10 +403,15 @@ def get_secrets_file(self, group: str, key: Optional[str] = None, group_version:
         return os.path.join(self._base_dir, *l)
 
     @staticmethod
-    def check_group_key(group: str):
+    def check_group_key(group: Optional[str]):
         if group is None or group == "":
             raise ValueError("secrets group is a mandatory field.")
 
+    @staticmethod
+    def check_env_name_key(env_name: Optional[str]):
+        if env_name is not None and len(env_name) <= 0:
+            raise ValueError(f"Invalid env_name {env_name}")
+
 
 @dataclass(frozen=True)
 class CompilationState(object):

flytekit/models/security.py

@@ -16,8 +16,8 @@ class Secret(_common.FlyteIdlEntity):
         group is the Name of the secret. For example in kubernetes secrets is the name of the secret
         key is optional and can be an individual secret identifier within the secret For k8s this is required
         version is the version of the secret. This is an optional field
-        env_name is name of the environment variable if this Secret is injected as an environment variable. This is an optional field
         mount_requirement provides a hint to the system as to how the secret should be injected
+        env_name is name of the environment variable if this Secret is injected as an environment variable. This is an optional field
     """
 
     class MountType(Enum):
@@ -39,8 +39,8 @@ class MountType(Enum):
     group: str
     key: Optional[str] = None
     group_version: Optional[str] = None
-    env_name: Optional[str] = None
     mount_requirement: MountType = MountType.ANY
+    env_name: Optional[str] = None
 
     def __post_init__(self):
         if self.group is None:
@@ -51,8 +51,8 @@ def to_flyte_idl(self) -> _sec.Secret:
             group=self.group,
             group_version=self.group_version,
             key=self.key,
-            env_name=self.env_name,
             mount_requirement=self.mount_requirement.value,
+            env_name=self.env_name,
         )
 
     @classmethod
@@ -61,8 +61,8 @@ def from_flyte_idl(cls, pb2_object: _sec.Secret) -> "Secret":
             group=pb2_object.group,
             group_version=pb2_object.group_version if pb2_object.group_version else None,
             key=pb2_object.key if pb2_object.key else None,
-            env_name=pb2_object.env_name if pb2_object.env_name else None,
             mount_requirement=Secret.MountType(pb2_object.mount_requirement),
+            env_name=pb2_object.env_name if pb2_object.env_name else None,
         )
 
 

plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py

@@ -55,8 +55,8 @@ def _secret_to_dict(secret: Secret) -> typing.Dict[str, typing.Optional[str]]:
             "group": secret.group,
             "key": secret.key,
             "group_version": secret.group_version,
-            "env_name": secret.name,
             "mount_requirement": secret.mount_requirement.value,
+            "env_name": secret.env_name,
         }
 
     def secret_connect_args_to_dicts(self) -> typing.Optional[typing.Dict[str, typing.Dict[str, typing.Optional[str]]]]:

tests/flytekit/common/parameterizers.py

@@ -238,6 +238,7 @@
     None,
     security.Secret(group="x", key="g"),
     security.Secret(group="x", key="y", mount_requirement=security.Secret.MountType.ANY),
+    security.Secret(group="x", key="y", env_name="z", mount_requirement=security.Secret.MountType.ANY),
     security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.FILE),
 ]