SK-1736: Add validations for invoke connection and Fix inconsistencies (

#146)

* SK-1736: Add get log level function and validations for invoke connection

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)
 

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)

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."

skyflow/utils/_utils.py

@@ -105,6 +105,9 @@ 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)
+
     try:
         return requests.Request(
             method = request.method.value,
@@ -114,7 +117,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)
 
 

skyflow/vault/controller/_connections.py

@@ -1,5 +1,4 @@
 import json
-
 import requests
 from skyflow.error import SkyflowError
 from skyflow.utils import construct_invoke_connection_request, SkyflowMessages, get_metrics, \
@@ -35,6 +34,7 @@ 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)
+            raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value,
+                               SkyflowMessages.ErrorCodes.SERVER_ERROR.value)

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)
 

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"}