From 27e4a6f6a0bcf66dfe85c3660716dbfe46150d4c Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Wed, 28 Jun 2023 20:57:05 +0200 Subject: [PATCH 1/6] Updates Secret model to include new env_name field Signed-off-by: Geert Pingen --- flytekit/models/security.py | 4 ++++ .../flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py | 1 + 2 files changed, 5 insertions(+) diff --git a/flytekit/models/security.py b/flytekit/models/security.py index 9af90a4b8a..79bf4d8f95 100644 --- a/flytekit/models/security.py +++ b/flytekit/models/security.py @@ -16,6 +16,7 @@ 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 """ @@ -38,6 +39,7 @@ class MountType(Enum): group: str key: Optional[str] = None group_version: Optional[str] = None + env_name: Optional[str] = None mount_requirement: MountType = MountType.ANY def __post_init__(self): @@ -49,6 +51,7 @@ 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, ) @@ -58,6 +61,7 @@ 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), ) diff --git a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py b/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py index 8e8c464bd4..d11b26a662 100644 --- a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py +++ b/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py @@ -55,6 +55,7 @@ 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, } From b541a215dfa6696392d6d81ab07e1d8a39050e72 Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Sat, 1 Jul 2023 11:11:42 +0200 Subject: [PATCH 2/6] Updates Secret to use env_name field as environment variable name on injection if exists Signed-off-by: Geert Pingen --- flytekit/core/context_manager.py | 40 +++++++++++++++---- flytekit/models/security.py | 8 ++-- .../flytekitplugins/sqlalchemy/task.py | 2 +- tests/flytekit/common/parameterizers.py | 1 + 4 files changed, 38 insertions(+), 13 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index 53ff504d55..ea2d0158c2 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -349,13 +349,20 @@ 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 @@ -363,18 +370,30 @@ def get(self, group: str, key: Optional[str] = None, group_version: Optional[str 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): diff --git a/flytekit/models/security.py b/flytekit/models/security.py index 79bf4d8f95..bf9cee1895 100644 --- a/flytekit/models/security.py +++ b/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, ) diff --git a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py b/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py index d11b26a662..6857e9878d 100644 --- a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py +++ b/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]]]]: diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index d5b07fe420..6d66526856 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/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), ] From c9afca2d4e65dcbc5446d76ae6b66b9575627089 Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Fri, 7 Jul 2023 21:20:00 +0200 Subject: [PATCH 3/6] First draft to add oneof Signed-off-by: Geert Pingen --- flytekit/core/context_manager.py | 7 +-- flytekit/models/security.py | 46 ++++++++++++++++--- .../flytekitplugins/sqlalchemy/task.py | 2 +- tests/flytekit/common/parameterizers.py | 2 +- 4 files changed, 46 insertions(+), 11 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index ea2d0158c2..fba98faa2b 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -361,17 +361,18 @@ def get( """ self.check_env_name_key(env_name) env_var = self.get_secrets_env_var(group, key, group_version, env_name) + fpath = None 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): + if fpath is not None and 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} or name {env_name}" - 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} or FilePath: {fpath}" ) def get_secrets_env_var( diff --git a/flytekit/models/security.py b/flytekit/models/security.py index bf9cee1895..88f6ea4931 100644 --- a/flytekit/models/security.py +++ b/flytekit/models/security.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import List, Optional @@ -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 - 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 + mount_requirement provides a hint to the system as to how the secret should be injected. Soon to be deprecated. + mount_target provies a target for secret injection. Can be environment variable or file path This is an optional field. """ class MountType(Enum): @@ -36,11 +36,43 @@ class MountType(Enum): Caution: May not be supported in all environments """ + @dataclass + class MountEnvVar(_common.FlyteIdlEntity): + name: str + + def to_flyte_idl(self) -> _sec.Secret.MountEnvVar: + return _sec.Secret.MountEnvVar( + name=self.name, + ) + + @classmethod + def from_flyte_idl(cls, pb2_object: _sec.Secret.MountEnvVar): + return cls( + name=pb2_object.name, + ) + + @dataclass + class MountFile(_common.FlyteIdlEntity): + path: str + + def to_flyte_idl(self) -> _sec.Secret.MountFile: + return _sec.Secret.MountFile( + path=self.path, + ) + + @classmethod + def from_flyte_idl(cls, pb2_object: _sec.Secret.MountFile): + return cls( + path=pb2_object.path, + ) + group: str + env_var: Optional[MountEnvVar] = None + file: Optional[MountFile] = None + key: Optional[str] = None group_version: Optional[str] = None mount_requirement: MountType = MountType.ANY - env_name: Optional[str] = None def __post_init__(self): if self.group is None: @@ -52,7 +84,8 @@ def to_flyte_idl(self) -> _sec.Secret: group_version=self.group_version, key=self.key, mount_requirement=self.mount_requirement.value, - env_name=self.env_name, + env_var=self.env_var.to_flyte_idl() if self.env_var else None, + file=self.file.to_flyte_idl() if self.file else None, ) @classmethod @@ -62,7 +95,8 @@ def from_flyte_idl(cls, pb2_object: _sec.Secret) -> "Secret": group_version=pb2_object.group_version if pb2_object.group_version else None, key=pb2_object.key if pb2_object.key else None, mount_requirement=Secret.MountType(pb2_object.mount_requirement), - env_name=pb2_object.env_name if pb2_object.env_name else None, + env_var=Secret.MountEnvVar.from_flyte_idl(pb2_object.env_var) if pb2_object.HasField("env_var") else None, + file=Secret.MountFile.from_flyte_idl(pb2_object.file) if pb2_object.HasField("file") else None, ) diff --git a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py b/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py index 6857e9878d..74d226572b 100644 --- a/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py +++ b/plugins/flytekit-sqlalchemy/flytekitplugins/sqlalchemy/task.py @@ -56,7 +56,7 @@ def _secret_to_dict(secret: Secret) -> typing.Dict[str, typing.Optional[str]]: "key": secret.key, "group_version": secret.group_version, "mount_requirement": secret.mount_requirement.value, - "env_name": secret.env_name, + "mount_target": secret.mount_target, } def secret_connect_args_to_dicts(self) -> typing.Optional[typing.Dict[str, typing.Dict[str, typing.Optional[str]]]]: diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index 6d66526856..0889620953 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -238,7 +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", mount_target=security.Secret.MountEnvVar(name="z"), mount_requirement=security.Secret.MountType.ANY), security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.FILE), ] From 4082866d0a3b3aab64bea44b1fdc5a2e4f395d3b Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Sat, 8 Jul 2023 12:10:45 +0200 Subject: [PATCH 4/6] Updates Makefile to use python3 Signed-off-by: Geert Pingen --- Makefile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 2d4d86050b..997651eb03 100644 --- a/Makefile +++ b/Makefile @@ -13,7 +13,7 @@ help: .PHONY: install-piptools install-piptools: # pip 22.1 broke pip-tools: https://github.com/jazzband/pip-tools/issues/1617 - python -m pip install -U pip-tools setuptools wheel "pip>=22.0.3,!=22.1" + python3 -m pip install -U pip-tools setuptools wheel "pip>=22.0.3,!=22.1" .PHONY: update_boilerplate update_boilerplate: @@ -58,7 +58,7 @@ unit_test: # Skip tensorflow tests and run them with the necessary env var set so that a working (albeit slower) # library is used to serialize/deserialize protobufs is used. pytest -m "not sandbox_test" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} && \ - PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python3 pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} doc-requirements.txt: export CUSTOM_COMPILE_COMMAND := make doc-requirements.txt doc-requirements.txt: doc-requirements.in install-piptools From 6a5956863a89a5d763560b5b24a8cc35a754526a Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Sat, 8 Jul 2023 16:54:16 +0200 Subject: [PATCH 5/6] Adds test security contexts and fixes test cases Signed-off-by: Geert Pingen --- Makefile | 3 ++- flytekit/models/security.py | 9 ++++----- tests/flytekit/common/parameterizers.py | 4 +++- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/Makefile b/Makefile index 997651eb03..5470d8957f 100644 --- a/Makefile +++ b/Makefile @@ -57,8 +57,9 @@ unit_test_codecov: unit_test: # Skip tensorflow tests and run them with the necessary env var set so that a working (albeit slower) # library is used to serialize/deserialize protobufs is used. + # Can use pytest --lf to only rerun previously failed tests. pytest -m "not sandbox_test" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} && \ - PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python3 pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} doc-requirements.txt: export CUSTOM_COMPILE_COMMAND := make doc-requirements.txt doc-requirements.txt: doc-requirements.in install-piptools diff --git a/flytekit/models/security.py b/flytekit/models/security.py index 88f6ea4931..ba19b5513e 100644 --- a/flytekit/models/security.py +++ b/flytekit/models/security.py @@ -67,12 +67,11 @@ def from_flyte_idl(cls, pb2_object: _sec.Secret.MountFile): ) group: str - env_var: Optional[MountEnvVar] = None - file: Optional[MountFile] = None - key: Optional[str] = None group_version: Optional[str] = None mount_requirement: MountType = MountType.ANY + env_var: Optional[MountEnvVar] = field(default_factory=lambda: None) + file: Optional[MountFile] = field(default_factory=lambda: None) def __post_init__(self): if self.group is None: @@ -84,8 +83,8 @@ def to_flyte_idl(self) -> _sec.Secret: group_version=self.group_version, key=self.key, mount_requirement=self.mount_requirement.value, - env_var=self.env_var.to_flyte_idl() if self.env_var else None, - file=self.file.to_flyte_idl() if self.file else None, + env_var=_sec.Secret.MountEnvVar(name=self.env_var.name) if self.env_var else None, + file=_sec.Secret.MountFile(path=self.file.path) if self.file else None, ) @classmethod diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py index 0889620953..c60a8afc67 100644 --- a/tests/flytekit/common/parameterizers.py +++ b/tests/flytekit/common/parameterizers.py @@ -238,7 +238,9 @@ 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", mount_target=security.Secret.MountEnvVar(name="z"), mount_requirement=security.Secret.MountType.ANY), + security.Secret(group="x", key="y", env_var=security.Secret.MountEnvVar(name="z")), + security.Secret(group="x", key="y", file=security.Secret.MountFile(path="/z")), + security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.ENV_VAR), security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.FILE), ] From 1529e9bc4ea46c6a3fc92af50046ad395af0de55 Mon Sep 17 00:00:00 2001 From: Geert Pingen Date: Thu, 13 Jul 2023 21:20:59 +0200 Subject: [PATCH 6/6] Adds default value for Optional[str] in key check Signed-off-by: Geert Pingen --- flytekit/core/context_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py index fba98faa2b..9de8e92bdd 100644 --- a/flytekit/core/context_manager.py +++ b/flytekit/core/context_manager.py @@ -404,12 +404,12 @@ def get_secrets_file( return os.path.join(self._base_dir, *l) @staticmethod - def check_group_key(group: Optional[str]): + def check_group_key(group: Optional[str] = None): if group is None or group == "": raise ValueError("secrets group is a mandatory field.") @staticmethod - def check_env_name_key(env_name: Optional[str]): + def check_env_name_key(env_name: Optional[str] = None): if env_name is not None and len(env_name) <= 0: raise ValueError(f"Invalid env_name {env_name}")