From 97be10e4136f958543e9d3555b6fe11151e7ecac Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 10:47:42 +0200 Subject: [PATCH 01/17] add support for custom auth from env Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 9 ++++++--- flytekit/clis/flyte_cli/main.py | 1 + flytekit/configuration/internal.py | 11 +++++++++++ 3 files changed, 18 insertions(+), 3 deletions(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 972b8655d5..ed97d5f798 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -21,6 +21,7 @@ from flytekit.clis.auth import credentials as _credentials_access from flytekit.configuration import AuthType, PlatformConfig +from flytekit.configuration.internal import Credentials from flytekit.exceptions import user as _user_exceptions from flytekit.exceptions.user import FlyteAuthenticationException from flytekit.loggers import cli_logger @@ -264,7 +265,7 @@ def _refresh_credentials_from_command(self): :return: """ - command = self._cfg.command + command = self._cfg.command or Credentials.COMMAND.read() if not command: raise FlyteAuthenticationException("No command specified in configuration for command authentication") cli_logger.debug("Starting external process to generate id token. Command {}".format(command)) @@ -273,13 +274,15 @@ def _refresh_credentials_from_command(self): except subprocess.CalledProcessError as e: cli_logger.error("Failed to generate token from command {}".format(command)) raise _user_exceptions.FlyteAuthenticationException("Problems refreshing token with command: " + str(e)) - self.set_access_token(output.stdout.strip()) + + authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.COMMAND.read() or None + self.set_access_token(output.stdout.strip(), authorization_header_key) def _refresh_credentials_noop(self): pass def refresh_credentials(self): - cfg_auth = self._cfg.auth_mode + cfg_auth = self._cfg.auth_mode or Credentials.AUTH_MODE.read() if type(cfg_auth) is str: try: cfg_auth = AuthType[cfg_auth.upper()] diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 21aec1c4ad..5920b250b5 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -277,6 +277,7 @@ def _get_client(host: str, insecure: bool) -> _friendly_client.SynchronousFlyteC if parent_ctx.obj["cacert"]: kwargs["root_certificates"] = parent_ctx.obj["cacert"] cfg = parent_ctx.obj["config"] + print(cfg) cfg = replace(cfg, endpoint=host, insecure=insecure) return _friendly_client.SynchronousFlyteClient(cfg, **kwargs) diff --git a/flytekit/configuration/internal.py b/flytekit/configuration/internal.py index fb09015b6f..86afdcf9d4 100644 --- a/flytekit/configuration/internal.py +++ b/flytekit/configuration/internal.py @@ -65,6 +65,17 @@ class Credentials(object): This command is executed to return a token using an external process. """ + AUTHORIZATION_METADATA_KEY = ConfigEntry(LegacyConfigEntry(SECTION, "authorization_metadata_key", list), YamlConfigEntry("admin.authorizationHeader", list)) + + """ + The authorization metadata key used for passing access tokens in gRPC requests. + Traditionally this value is 'authorization' however it is made configurable. + The default value in admin is 'flyte-authorization', see + https://github.com/flyteorg/flyteadmin/blob/87cc9424ca4b684515c92f9d40cbcc4f5578e70c/auth/config/config.go#L61-L62 + For reasoning about calling the header 'flyte-authorization' instead of 'authorization', see + https://github.com/flyteorg/flyteadmin/blob/689bfab25c372f4dffcd398dc896cd786912e0df/auth/handlers.go#L187-L191 + """ + CLIENT_ID = ConfigEntry(LegacyConfigEntry(SECTION, "client_id"), YamlConfigEntry("admin.clientId")) """ This is the public identifier for the app which handles authorization for a Flyte deployment. From 044a13db5c59765685401a131331e267083c1fd5 Mon Sep 17 00:00:00 2001 From: sonjaer Date: Wed, 15 Jun 2022 10:57:05 +0200 Subject: [PATCH 02/17] Update flytekit/clis/flyte_cli/main.py Signed-off-by: Sonja Ericsson --- flytekit/clis/flyte_cli/main.py | 1 - 1 file changed, 1 deletion(-) diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 5920b250b5..21aec1c4ad 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -277,7 +277,6 @@ def _get_client(host: str, insecure: bool) -> _friendly_client.SynchronousFlyteC if parent_ctx.obj["cacert"]: kwargs["root_certificates"] = parent_ctx.obj["cacert"] cfg = parent_ctx.obj["config"] - print(cfg) cfg = replace(cfg, endpoint=host, insecure=insecure) return _friendly_client.SynchronousFlyteClient(cfg, **kwargs) From 6d0610552ad4116f89ff93fe45ce6a8a46eebda5 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 11:15:41 +0200 Subject: [PATCH 03/17] fix Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index ed97d5f798..7d136e0298 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -275,7 +275,9 @@ def _refresh_credentials_from_command(self): cli_logger.error("Failed to generate token from command {}".format(command)) raise _user_exceptions.FlyteAuthenticationException("Problems refreshing token with command: " + str(e)) - authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.COMMAND.read() or None + authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.COMMAND.read() + if not authorization_header_key: + self.set_access_token(output.stdout.strip()) self.set_access_token(output.stdout.strip(), authorization_header_key) def _refresh_credentials_noop(self): From 8573d5b27f767de68faaa2fcfac2521344f4af32 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 11:18:26 +0200 Subject: [PATCH 04/17] fix Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 7d136e0298..bd1423df8b 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -275,7 +275,7 @@ def _refresh_credentials_from_command(self): cli_logger.error("Failed to generate token from command {}".format(command)) raise _user_exceptions.FlyteAuthenticationException("Problems refreshing token with command: " + str(e)) - authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.COMMAND.read() + authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.AUTHORIZATION_METADATA_KEY.read() if not authorization_header_key: self.set_access_token(output.stdout.strip()) self.set_access_token(output.stdout.strip(), authorization_header_key) From 34aeec2a377f46a656ef9713595b48504f3e9fce Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 11:30:12 +0200 Subject: [PATCH 05/17] fix Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 3 +-- flytekit/configuration/internal.py | 11 ----------- 2 files changed, 1 insertion(+), 13 deletions(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index bd1423df8b..78d009094b 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -274,8 +274,7 @@ def _refresh_credentials_from_command(self): except subprocess.CalledProcessError as e: cli_logger.error("Failed to generate token from command {}".format(command)) raise _user_exceptions.FlyteAuthenticationException("Problems refreshing token with command: " + str(e)) - - authorization_header_key = self.public_client_config.authorization_metadata_key or Credentials.AUTHORIZATION_METADATA_KEY.read() + authorization_header_key = self.public_client_config.authorization_metadata_key or None if not authorization_header_key: self.set_access_token(output.stdout.strip()) self.set_access_token(output.stdout.strip(), authorization_header_key) diff --git a/flytekit/configuration/internal.py b/flytekit/configuration/internal.py index 86afdcf9d4..fb09015b6f 100644 --- a/flytekit/configuration/internal.py +++ b/flytekit/configuration/internal.py @@ -65,17 +65,6 @@ class Credentials(object): This command is executed to return a token using an external process. """ - AUTHORIZATION_METADATA_KEY = ConfigEntry(LegacyConfigEntry(SECTION, "authorization_metadata_key", list), YamlConfigEntry("admin.authorizationHeader", list)) - - """ - The authorization metadata key used for passing access tokens in gRPC requests. - Traditionally this value is 'authorization' however it is made configurable. - The default value in admin is 'flyte-authorization', see - https://github.com/flyteorg/flyteadmin/blob/87cc9424ca4b684515c92f9d40cbcc4f5578e70c/auth/config/config.go#L61-L62 - For reasoning about calling the header 'flyte-authorization' instead of 'authorization', see - https://github.com/flyteorg/flyteadmin/blob/689bfab25c372f4dffcd398dc896cd786912e0df/auth/handlers.go#L187-L191 - """ - CLIENT_ID = ConfigEntry(LegacyConfigEntry(SECTION, "client_id"), YamlConfigEntry("admin.clientId")) """ This is the public identifier for the app which handles authorization for a Flyte deployment. From 6ff071cfde8416b6ad04d43644d1c77250c739c4 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 12:33:47 +0200 Subject: [PATCH 06/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 4899051169..55a9b5182d 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -49,18 +49,22 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad assert client.check_access_token("abc") +@mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") -def test_refresh_credentials_from_command(mock_call_to_external_process): +def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): command = ["command", "generating", "token"] token = "token" mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) + mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() + cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=AuthType.EXTERNAL_PROCESS, command=command)) cc._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) - + assert cc._metadata[0][0] == "flyte-authorization" + assert cc._metadata[0][1] == "Bearer token" @mock.patch("flytekit.clients.raw.dataproxy_service") @mock.patch("flytekit.clients.raw.get_basic_authorization_header") From 22f6c96c26b753e72f32a78ba28896172019428d Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 12:49:20 +0200 Subject: [PATCH 07/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 55a9b5182d..e5eb21f2b9 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -66,6 +66,7 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad assert cc._metadata[0][0] == "flyte-authorization" assert cc._metadata[0][1] == "Bearer token" + @mock.patch("flytekit.clients.raw.dataproxy_service") @mock.patch("flytekit.clients.raw.get_basic_authorization_header") @mock.patch("flytekit.clients.raw.get_token") @@ -74,13 +75,13 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad @mock.patch("flytekit.clients.raw.grpc.insecure_channel") @mock.patch("flytekit.clients.raw.grpc.secure_channel") def test_refresh_client_credentials_aka_basic( - mock_secure_channel, - mock_channel, - mock_admin, - mock_admin_auth, - mock_get_token, - mock_get_basic_header, - mock_dataproxy, + mock_secure_channel, + mock_channel, + mock_admin, + mock_admin_auth, + mock_get_token, + mock_get_basic_header, + mock_dataproxy, ): mock_secure_channel.return_value = True mock_channel.return_value = True From 227b35f4eb478886b92d51108422c38a488b8fa1 Mon Sep 17 00:00:00 2001 From: sonjaer Date: Wed, 15 Jun 2022 13:08:10 +0200 Subject: [PATCH 08/17] Update test_raw.py Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index e5eb21f2b9..68672225e9 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -75,13 +75,13 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad @mock.patch("flytekit.clients.raw.grpc.insecure_channel") @mock.patch("flytekit.clients.raw.grpc.secure_channel") def test_refresh_client_credentials_aka_basic( - mock_secure_channel, - mock_channel, - mock_admin, - mock_admin_auth, - mock_get_token, - mock_get_basic_header, - mock_dataproxy, + mock_secure_channel, + mock_channel, + mock_admin, + mock_admin_auth, + mock_get_token, + mock_get_basic_header, + mock_dataproxy, ): mock_secure_channel.return_value = True mock_channel.return_value = True From d7606677e507ebe3b0fe17cdb3112805aa29edb5 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 13:23:51 +0200 Subject: [PATCH 09/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 44 ++++++++++++++++++++++--- 1 file changed, 40 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 68672225e9..8011708b13 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -52,19 +52,46 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad @mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): - command = ["command", "generating", "token"] + _test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth, None) + + +@patch("flytekit.configuration.internal.Credentials.COMMAND.read") +@mock.patch("flytekit.clients.raw.auth_service") +@mock.patch("subprocess.run") +def test_refresh_credentials_from_command_from_environment_variable( + mock_call_to_external_process, mock_admin_auth, mock_env +): + _test_refresh_credentials_from_command( + mock_call_to_external_process, + mock_admin_auth, + mock_env, + True, + command=["command", "generating", "token", "env"], + ) + + +def _test_refresh_credentials_from_command( + mock_call_to_external_process, mock_admin_auth, mock_env, env=False, command=None +): + if command is None: + command = ["command", "generating", "token"] token = "token" mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() - cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=AuthType.EXTERNAL_PROCESS, command=command)) + if not env: + cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) + else: + mock_env.return_value = command + cc = RawSynchronousFlyteClient(PlatformConfig()) + cc._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) assert cc._metadata[0][0] == "flyte-authorization" - assert cc._metadata[0][1] == "Bearer token" + assert cc._metadata[0][1] == f"Bearer {token}" @mock.patch("flytekit.clients.raw.dataproxy_service") @@ -200,6 +227,15 @@ def test_basic_strings(mocked_method): @patch.object(RawSynchronousFlyteClient, "_refresh_credentials_from_command") def test_refresh_command(mocked_method): - cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=AuthType.EXTERNAL_PROCESS)) + cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=AuthType.EXTERNALCOMMAND)) + cc.refresh_credentials() + assert mocked_method.called + + +@patch("flytekit.configuration.internal.Credentials.AUTH_MODE.read") +@patch.object(RawSynchronousFlyteClient, "_refresh_credentials_from_command") +def test_refresh_from_environment_variable(mocked_method, read_env): + cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=None)) + read_env.return_value = "external_process" cc.refresh_credentials() assert mocked_method.called From 237fc03fe10737881ad46b7d018b24d693d6ee89 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 13:27:31 +0200 Subject: [PATCH 10/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 8011708b13..5dbcb50795 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -52,7 +52,13 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad @mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): - _test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth, None) + _test_refresh_credentials_from_command( + mock_call_to_external_process=mock_call_to_external_process, + mock_admin_auth=mock_admin_auth, + mock_env=None, + env=False, + command=["command", "generating", "token"], + ) @patch("flytekit.configuration.internal.Credentials.COMMAND.read") @@ -62,19 +68,15 @@ def test_refresh_credentials_from_command_from_environment_variable( mock_call_to_external_process, mock_admin_auth, mock_env ): _test_refresh_credentials_from_command( - mock_call_to_external_process, - mock_admin_auth, - mock_env, - True, + mock_call_to_external_process=mock_call_to_external_process, + mock_admin_auth=mock_admin_auth, + mock_env=mock_env, + env=True, command=["command", "generating", "token", "env"], ) -def _test_refresh_credentials_from_command( - mock_call_to_external_process, mock_admin_auth, mock_env, env=False, command=None -): - if command is None: - command = ["command", "generating", "token"] +def _test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth, mock_env, env, command): token = "token" mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) From 9a1c917b726f86258eeebfd16aa56ab8182ac16f Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 16:37:41 +0200 Subject: [PATCH 11/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 5dbcb50795..53345e8bcc 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -53,11 +53,7 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad @mock.patch("subprocess.run") def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): _test_refresh_credentials_from_command( - mock_call_to_external_process=mock_call_to_external_process, - mock_admin_auth=mock_admin_auth, - mock_env=None, - env=False, - command=["command", "generating", "token"], + mock_call_to_external_process=mock_call_to_external_process, mock_admin_auth=mock_admin_auth ) @@ -65,18 +61,20 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad @mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") def test_refresh_credentials_from_command_from_environment_variable( - mock_call_to_external_process, mock_admin_auth, mock_env + mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env ): _test_refresh_credentials_from_command( mock_call_to_external_process=mock_call_to_external_process, mock_admin_auth=mock_admin_auth, - mock_env=mock_env, + mock_read_command_from_env=mock_read_command_from_env, env=True, - command=["command", "generating", "token", "env"], ) -def _test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth, mock_env, env, command): +def _test_refresh_credentials_from_command( + mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env=None, env=False +): + command = ["command", "generating", "token"] token = "token" mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) @@ -86,7 +84,7 @@ def _test_refresh_credentials_from_command(mock_call_to_external_process, mock_a if not env: cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) else: - mock_env.return_value = command + mock_read_command_from_env.return_value = command cc = RawSynchronousFlyteClient(PlatformConfig()) cc._refresh_credentials_from_command() @@ -236,8 +234,8 @@ def test_refresh_command(mocked_method): @patch("flytekit.configuration.internal.Credentials.AUTH_MODE.read") @patch.object(RawSynchronousFlyteClient, "_refresh_credentials_from_command") -def test_refresh_from_environment_variable(mocked_method, read_env): +def test_refresh_from_environment_variable(mocked_method, mock_read_env): cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=None)) - read_env.return_value = "external_process" + mock_read_env.return_value = "external_process" cc.refresh_credentials() assert mocked_method.called From b8e4b2e8c7650f37cdfd0765ffc8938cbfbe7c19 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 16:39:53 +0200 Subject: [PATCH 12/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 53345e8bcc..c09585459c 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -75,18 +75,16 @@ def _test_refresh_credentials_from_command( mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env=None, env=False ): command = ["command", "generating", "token"] - token = "token" - - mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) - mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() - if not env: cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) else: mock_read_command_from_env.return_value = command cc = RawSynchronousFlyteClient(PlatformConfig()) + token = "token" + mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) + cc._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) From f018d12aa30965960d0411564f830bf1ce492c61 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Wed, 15 Jun 2022 17:02:12 +0200 Subject: [PATCH 13/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index c09585459c..985de5194f 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -67,22 +67,22 @@ def test_refresh_credentials_from_command_from_environment_variable( mock_call_to_external_process=mock_call_to_external_process, mock_admin_auth=mock_admin_auth, mock_read_command_from_env=mock_read_command_from_env, - env=True, ) def _test_refresh_credentials_from_command( - mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env=None, env=False + mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env=None ): + token = "token" command = ["command", "generating", "token"] + mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() - if not env: + if not mock_read_command_from_env: cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) else: mock_read_command_from_env.return_value = command cc = RawSynchronousFlyteClient(PlatformConfig()) - token = "token" mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) cc._refresh_credentials_from_command() @@ -232,8 +232,8 @@ def test_refresh_command(mocked_method): @patch("flytekit.configuration.internal.Credentials.AUTH_MODE.read") @patch.object(RawSynchronousFlyteClient, "_refresh_credentials_from_command") -def test_refresh_from_environment_variable(mocked_method, mock_read_env): +def test_refresh_from_environment_variable(mocked_method, mock_read_auth_mode_from_env): cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=None)) - mock_read_env.return_value = "external_process" + mock_read_auth_mode_from_env.return_value = "external_process" cc.refresh_credentials() assert mocked_method.called From 3bcefd3dcdb33c87db194c90d7335a08a40e2257 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Thu, 16 Jun 2022 10:33:39 +0200 Subject: [PATCH 14/17] fix Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 4 +-- flytekit/configuration/file.py | 5 +++- tests/flytekit/unit/clients/test_raw.py | 36 ++++--------------------- 3 files changed, 11 insertions(+), 34 deletions(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index 78d009094b..b5c3327ac4 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -265,7 +265,7 @@ def _refresh_credentials_from_command(self): :return: """ - command = self._cfg.command or Credentials.COMMAND.read() + command = self._cfg.command if not command: raise FlyteAuthenticationException("No command specified in configuration for command authentication") cli_logger.debug("Starting external process to generate id token. Command {}".format(command)) @@ -283,7 +283,7 @@ def _refresh_credentials_noop(self): pass def refresh_credentials(self): - cfg_auth = self._cfg.auth_mode or Credentials.AUTH_MODE.read() + cfg_auth = self._cfg.auth_mode if type(cfg_auth) is str: try: cfg_auth = AuthType[cfg_auth.upper()] diff --git a/flytekit/configuration/file.py b/flytekit/configuration/file.py index b329a16112..1f79016f63 100644 --- a/flytekit/configuration/file.py +++ b/flytekit/configuration/file.py @@ -32,13 +32,16 @@ class LegacyConfigEntry(object): option: str type_: typing.Type = str + def get_env_name(self): + return f"FLYTE_{self.section.upper()}_{self.option.upper()}" + def read_from_env(self, transform: typing.Optional[typing.Callable] = None) -> typing.Optional[typing.Any]: """ Reads the config entry from environment variable, the structure of the env var is current ``FLYTE_{SECTION}_{OPTION}`` all upper cased. We will change this in the future. :return: """ - env = f"FLYTE_{self.section.upper()}_{self.option.upper()}" + env = self.get_env_name() v = os.environ.get(env, None) if v is None: return None diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 985de5194f..73cb58a144 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -9,6 +9,7 @@ from flytekit.clients.raw import RawSynchronousFlyteClient, get_basic_authorization_header, get_token from flytekit.configuration import AuthType, PlatformConfig +from flytekit.configuration.internal import Credentials def get_admin_stub_mock() -> mock.MagicMock: @@ -52,39 +53,13 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad @mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): - _test_refresh_credentials_from_command( - mock_call_to_external_process=mock_call_to_external_process, mock_admin_auth=mock_admin_auth - ) - - -@patch("flytekit.configuration.internal.Credentials.COMMAND.read") -@mock.patch("flytekit.clients.raw.auth_service") -@mock.patch("subprocess.run") -def test_refresh_credentials_from_command_from_environment_variable( - mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env -): - _test_refresh_credentials_from_command( - mock_call_to_external_process=mock_call_to_external_process, - mock_admin_auth=mock_admin_auth, - mock_read_command_from_env=mock_read_command_from_env, - ) - - -def _test_refresh_credentials_from_command( - mock_call_to_external_process, mock_admin_auth, mock_read_command_from_env=None -): token = "token" command = ["command", "generating", "token"] mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() - if not mock_read_command_from_env: - cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) - else: - mock_read_command_from_env.return_value = command - cc = RawSynchronousFlyteClient(PlatformConfig()) + cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) - cc._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) @@ -230,10 +205,9 @@ def test_refresh_command(mocked_method): assert mocked_method.called -@patch("flytekit.configuration.internal.Credentials.AUTH_MODE.read") @patch.object(RawSynchronousFlyteClient, "_refresh_credentials_from_command") -def test_refresh_from_environment_variable(mocked_method, mock_read_auth_mode_from_env): - cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=None)) - mock_read_auth_mode_from_env.return_value = "external_process" +def test_refresh_from_environment_variable(mocked_method, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv(Credentials.AUTH_MODE.legacy.get_env_name(), AuthType.EXTERNAL_PROCESS.name, prepend=False) + cc = RawSynchronousFlyteClient(PlatformConfig(auth_mode=None).auto(None)) cc.refresh_credentials() assert mocked_method.called From 25bd533d9ec93a333a7636da1cc0cc6ffc592e2b Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Thu, 16 Jun 2022 10:39:08 +0200 Subject: [PATCH 15/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 73cb58a144..1bc1991a79 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -50,9 +50,10 @@ def test_client_set_token(mock_secure_channel, mock_channel, mock_admin, mock_ad assert client.check_access_token("abc") +@mock.patch("flytekit.clients.raw.RawSynchronousFlyteClient.set_access_token") @mock.patch("flytekit.clients.raw.auth_service") @mock.patch("subprocess.run") -def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth): +def test_refresh_credentials_from_command(mock_call_to_external_process, mock_admin_auth, mock_set_access_token): token = "token" command = ["command", "generating", "token"] @@ -63,8 +64,7 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad cc._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) - assert cc._metadata[0][0] == "flyte-authorization" - assert cc._metadata[0][1] == f"Bearer {token}" + mock_set_access_token.assert_called_with(token, cc.public_client_config.authorization_metadata_key) @mock.patch("flytekit.clients.raw.dataproxy_service") From c1e7014bc4a9f9303c6421b43406eb699d4b73e6 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Thu, 16 Jun 2022 10:40:00 +0200 Subject: [PATCH 16/17] fix Signed-off-by: Sonja Ericsson --- tests/flytekit/unit/clients/test_raw.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/clients/test_raw.py b/tests/flytekit/unit/clients/test_raw.py index 1bc1991a79..6636c1e0c8 100644 --- a/tests/flytekit/unit/clients/test_raw.py +++ b/tests/flytekit/unit/clients/test_raw.py @@ -58,13 +58,13 @@ def test_refresh_credentials_from_command(mock_call_to_external_process, mock_ad command = ["command", "generating", "token"] mock_admin_auth.AuthMetadataServiceStub.return_value = get_admin_stub_mock() - cc = RawSynchronousFlyteClient(PlatformConfig(command=command)) + client = RawSynchronousFlyteClient(PlatformConfig(command=command)) mock_call_to_external_process.return_value = CompletedProcess(command, 0, stdout=token) - cc._refresh_credentials_from_command() + client._refresh_credentials_from_command() mock_call_to_external_process.assert_called_with(command, capture_output=True, text=True, check=True) - mock_set_access_token.assert_called_with(token, cc.public_client_config.authorization_metadata_key) + mock_set_access_token.assert_called_with(token, client.public_client_config.authorization_metadata_key) @mock.patch("flytekit.clients.raw.dataproxy_service") From af9b033938874ffa5a73dadaef204350f3c57d15 Mon Sep 17 00:00:00 2001 From: Sonja Ericsson Date: Thu, 16 Jun 2022 13:51:36 +0200 Subject: [PATCH 17/17] fix Signed-off-by: Sonja Ericsson --- flytekit/clients/raw.py | 1 - 1 file changed, 1 deletion(-) diff --git a/flytekit/clients/raw.py b/flytekit/clients/raw.py index b5c3327ac4..9bdf85a178 100644 --- a/flytekit/clients/raw.py +++ b/flytekit/clients/raw.py @@ -21,7 +21,6 @@ from flytekit.clis.auth import credentials as _credentials_access from flytekit.configuration import AuthType, PlatformConfig -from flytekit.configuration.internal import Credentials from flytekit.exceptions import user as _user_exceptions from flytekit.exceptions.user import FlyteAuthenticationException from flytekit.loggers import cli_logger