From db7556e20ad6f24d1e3c06f38a6761ca67283dc5 Mon Sep 17 00:00:00 2001 From: saileshwar-skyflow Date: Tue, 3 Dec 2024 18:58:39 +0530 Subject: [PATCH 1/3] SK-1736: Add get log level function and validations for invoke connection --- skyflow/client/skyflow.py | 3 +++ skyflow/utils/_skyflow_messages.py | 1 + skyflow/utils/_utils.py | 8 +++++++- skyflow/vault/controller/_connections.py | 11 ++++++----- tests/client/test_skyflow.py | 5 +++++ 5 files changed, 22 insertions(+), 6 deletions(-) diff --git a/skyflow/client/skyflow.py b/skyflow/client/skyflow.py index 3dd6bf4..be3f7d9 100644 --- a/skyflow/client/skyflow.py +++ b/skyflow/client/skyflow.py @@ -57,6 +57,9 @@ def set_log_level(self, log_level): self.__builder._Builder__set_log_level(log_level) return self + def get_log_level(self): + return self.__builder._Builder__log_level + def update_log_level(self, log_level): self.__builder._Builder__set_log_level(log_level) diff --git a/skyflow/utils/_skyflow_messages.py b/skyflow/utils/_skyflow_messages.py index c757fa4..954c5e1 100644 --- a/skyflow/utils/_skyflow_messages.py +++ b/skyflow/utils/_skyflow_messages.py @@ -132,6 +132,7 @@ class Error(Enum): INVALID_URL = f"{error_prefix} Validation error. Connection url {{}} is invalid. Specify a valid connection url." INVALID_PATH_PARAMS = f"{error_prefix} Validation error. Path parameters aren't valid. Specify valid path parameters." INVALID_QUERY_PARAMS = f"{error_prefix} Validation error. Query parameters aren't valid. Specify valid query parameters." + INVALID_REQUEST_METHOD = f"{error_prefix} Validation error. Invalid request method. Specify the request method as enum RequestMethod" MISSING_PRIVATE_KEY = f"{error_prefix} Initialization failed. Unable to read private key in credentials. Verify your private key." MISSING_CLIENT_ID = f"{error_prefix} Initialization failed. Unable to read client ID in credentials. Verify your client ID." diff --git a/skyflow/utils/_utils.py b/skyflow/utils/_utils.py index 1db30a8..905126a 100644 --- a/skyflow/utils/_utils.py +++ b/skyflow/utils/_utils.py @@ -105,6 +105,12 @@ def construct_invoke_connection_request(request, connection_url, logger) -> Prep validate_invoke_connection_params(logger, request.query_params, request.path_params) + if not hasattr(request.method, 'value'): + raise SkyflowError(SkyflowMessages.Error.INVALID_REQUEST_METHOD.value, invalid_input_error_code) + + if not request.method.value: + raise SkyflowError(SkyflowMessages.Error.INVALID_REQUEST_METHOD.value, invalid_input_error_code) + try: return requests.Request( method = request.method.value, @@ -114,7 +120,7 @@ def construct_invoke_connection_request(request, connection_url, logger) -> Prep params = request.query_params, files = files ).prepare() - except requests.exceptions.InvalidURL: + except Exception: raise SkyflowError(SkyflowMessages.Error.INVALID_URL.value.format(connection_url), invalid_input_error_code) diff --git a/skyflow/vault/controller/_connections.py b/skyflow/vault/controller/_connections.py index 52e2c25..aebc584 100644 --- a/skyflow/vault/controller/_connections.py +++ b/skyflow/vault/controller/_connections.py @@ -1,9 +1,8 @@ import json - import requests -from skyflow.error import SkyflowError +from skyflow.generated.rest.exceptions import UnauthorizedException from skyflow.utils import construct_invoke_connection_request, SkyflowMessages, get_metrics, \ - parse_invoke_connection_response + parse_invoke_connection_response, handle_exception from skyflow.utils.logger import log_info from skyflow.vault.connection import InvokeConnectionRequest @@ -35,6 +34,8 @@ def invoke(self, request: InvokeConnectionRequest): session.close() invoke_connection_response = parse_invoke_connection_response(response) return invoke_connection_response + except Exception as e: - print(e) - raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value, SkyflowMessages.ErrorCodes.SERVER_ERROR.value) \ No newline at end of file + handle_exception(e, self.__vault_client.get_logger()) + except UnauthorizedException as e: + handle_exception(e, self.__vault_client.get_logger()) \ No newline at end of file diff --git a/tests/client/test_skyflow.py b/tests/client/test_skyflow.py index 87ef1a8..621cdee 100644 --- a/tests/client/test_skyflow.py +++ b/tests/client/test_skyflow.py @@ -125,6 +125,11 @@ def test_build_calls_validate_vault_config(self, mock_validate_vault_config): self.builder.build() mock_validate_vault_config.assert_called_once_with(self.builder._Builder__logger, VALID_VAULT_CONFIG) + def test_get_log_level(self): + builder = self.builder.set_log_level(LogLevel.ERROR) + client = self.builder.build() + self.assertEqual(LogLevel.ERROR, client.get_log_level()) + def test_add_connection_config_valid(self): result = self.builder.add_connection_config(VALID_CONNECTION_CONFIG) From a741e6e81e0ce14f8d406f59bbb61a0d58c5716a Mon Sep 17 00:00:00 2001 From: saileshwar-skyflow Date: Wed, 4 Dec 2024 13:25:43 +0530 Subject: [PATCH 2/3] SK-1736: Fix inconsistencies --- skyflow/service_account/_utils.py | 4 +--- skyflow/vault/controller/_connections.py | 7 ++++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/skyflow/service_account/_utils.py b/skyflow/service_account/_utils.py index 86da7d2..05035bb 100644 --- a/skyflow/service_account/_utils.py +++ b/skyflow/service_account/_utils.py @@ -117,7 +117,6 @@ def get_signed_tokens(credentials_obj, options): try: expiry_time = time.time() + options.get("time_to_live", 60) prefix = "signed_token_" - response_array=[] if options and options.get("data_tokens"): for token in options["data_tokens"]: @@ -136,9 +135,8 @@ def get_signed_tokens(credentials_obj, options): private_key = credentials_obj.get("privateKey") signed_jwt = jwt.encode(claims, private_key, algorithm="RS256") response_object = get_signed_data_token_response_object(prefix + signed_jwt, token) - response_array.append(response_object) log_info(SkyflowMessages.Info.GET_SIGNED_DATA_TOKEN_SUCCESS.value) - return response_array + return response_object except Exception: raise SkyflowError(SkyflowMessages.Error.INVALID_CREDENTIALS.value, invalid_input_error_code) diff --git a/skyflow/vault/controller/_connections.py b/skyflow/vault/controller/_connections.py index aebc584..a017c92 100644 --- a/skyflow/vault/controller/_connections.py +++ b/skyflow/vault/controller/_connections.py @@ -1,5 +1,7 @@ import json import requests + +from skyflow.error import SkyflowError from skyflow.generated.rest.exceptions import UnauthorizedException from skyflow.utils import construct_invoke_connection_request, SkyflowMessages, get_metrics, \ parse_invoke_connection_response, handle_exception @@ -36,6 +38,5 @@ def invoke(self, request: InvokeConnectionRequest): return invoke_connection_response except Exception as e: - handle_exception(e, self.__vault_client.get_logger()) - except UnauthorizedException as e: - handle_exception(e, self.__vault_client.get_logger()) \ No newline at end of file + raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value, + SkyflowMessages.ErrorCodes.SERVER_ERROR.value) \ No newline at end of file From 1297d5a9636129be3620aac0501d69a67a9b7ad6 Mon Sep 17 00:00:00 2001 From: saileshwar-skyflow Date: Wed, 4 Dec 2024 13:35:47 +0530 Subject: [PATCH 3/3] SK-1736: Added tests --- skyflow/utils/_utils.py | 3 --- skyflow/vault/controller/_connections.py | 4 +--- tests/utils/test__utils.py | 15 +++++++++++++++ 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/skyflow/utils/_utils.py b/skyflow/utils/_utils.py index 905126a..cc400b2 100644 --- a/skyflow/utils/_utils.py +++ b/skyflow/utils/_utils.py @@ -108,9 +108,6 @@ def construct_invoke_connection_request(request, connection_url, logger) -> Prep if not hasattr(request.method, 'value'): raise SkyflowError(SkyflowMessages.Error.INVALID_REQUEST_METHOD.value, invalid_input_error_code) - if not request.method.value: - raise SkyflowError(SkyflowMessages.Error.INVALID_REQUEST_METHOD.value, invalid_input_error_code) - try: return requests.Request( method = request.method.value, diff --git a/skyflow/vault/controller/_connections.py b/skyflow/vault/controller/_connections.py index a017c92..2fc52f1 100644 --- a/skyflow/vault/controller/_connections.py +++ b/skyflow/vault/controller/_connections.py @@ -1,10 +1,8 @@ import json import requests - from skyflow.error import SkyflowError -from skyflow.generated.rest.exceptions import UnauthorizedException from skyflow.utils import construct_invoke_connection_request, SkyflowMessages, get_metrics, \ - parse_invoke_connection_response, handle_exception + parse_invoke_connection_response from skyflow.utils.logger import log_info from skyflow.vault.connection import InvokeConnectionRequest diff --git a/tests/utils/test__utils.py b/tests/utils/test__utils.py index d5d4c72..c9010c9 100644 --- a/tests/utils/test__utils.py +++ b/tests/utils/test__utils.py @@ -122,6 +122,21 @@ def test_construct_invoke_connection_request_with_invalid_headers(self): self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_REQUEST_HEADERS.value) + def test_construct_invoke_connection_request_with_invalid_request_method(self): + mock_connection_request = Mock() + mock_connection_request.path_params = {"param1": "value1"} + mock_connection_request.headers = {"Content-Type": ContentType.JSON.value} + mock_connection_request.body = {"key": "value"} + mock_connection_request.method = "POST" + mock_connection_request.query_params = {"query": "test"} + + connection_url = "https://example.com/{param1}/endpoint" + + with self.assertRaises(SkyflowError) as context: + result = construct_invoke_connection_request(mock_connection_request, connection_url, logger=None) + + self.assertEqual(context.exception.message, SkyflowMessages.Error.INVALID_REQUEST_METHOD.value) + def test_construct_invoke_connection_request_with_invalid_request_body(self): mock_connection_request = Mock() mock_connection_request.path_params = {"param1": "value1"}