SK-1889: Merge branch 'v2' into SK-1889-partial-error-cases-when-continue-on-error

Author: saileshwar-skyflow

skyflow/utils/_skyflow_messages.py

@@ -277,6 +277,7 @@ class ErrorLogs(Enum):
         UPDATE_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Update request resulted in failure."
         QUERY_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Query request resulted in failure."
         GET_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Get request resulted in failure."
+        INVOKE_CONNECTION_REQUEST_REJECTED = f"{ERROR}: [{error_prefix}] Invoke connection request resulted in failure."
 
     class Interface(Enum):
         INSERT = "INSERT"

skyflow/utils/_utils.py

@@ -330,24 +330,30 @@ def parse_invoke_connection_response(api_response: requests.Response):
 
             invoke_connection_response.response = json_content
             return invoke_connection_response
-        except:
+        except Exception as e:
             raise SkyflowError(SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format(content), status_code)
     except HTTPError:
-        message = SkyflowMessages.Error.API_ERROR.value.format(status_code)
-        if api_response and api_response.content:
-            try:
-                error_response = json.loads(content)
-                if isinstance(error_response.get('error'), dict) and 'message' in error_response['error']:
-                    message = error_response['error']['message']
-            except json.JSONDecodeError:
-                message = SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format(content)
-
-        if 'x-request-id' in api_response.headers:
-            message += ' - request id: ' + api_response.headers['x-request-id']
-
+        message = SkyflowMessages.Error.API_ERROR.value.format(status_code)  
+        try:
+            error_response = json.loads(content)                    
+            request_id = api_response.headers['x-request-id']
+            error_from_client = api_response.headers.get('error-from-client')
+
+            status_code = error_response.get('error', {}).get('http_code', 500)  # Default to 500 if not found
+            http_status = error_response.get('error', {}).get('http_status')
+            grpc_code = error_response.get('error', {}).get('grpc_code')
+            details = error_response.get('error', {}).get('details')
+            message = error_response.get('error', {}).get('message', "An unknown error occurred.")
+            
+            if error_from_client is not None:
+                if details is None: details = []
+                details.append({'error_from_client': error_from_client})
+
+            raise SkyflowError(message, status_code, request_id, grpc_code, http_status, details)
+        except json.JSONDecodeError:
+            message = SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format(content)
         raise SkyflowError(message, status_code)
 
-
 def log_and_reject_error(description, status_code, request_id, http_status=None, grpc_code=None, details=None, logger = None):
     raise SkyflowError(description, status_code, request_id, grpc_code, http_status, details)
 

skyflow/utils/_version.py

@@ -1 +1 @@
-SDK_VERSION = '2.0.0b1.dev0+dcb5ddc'
+SDK_VERSION = '2.0.0b1.dev0+dcb5ddc'

skyflow/vault/controller/_connections.py

@@ -3,7 +3,7 @@
 from skyflow.error import SkyflowError
 from skyflow.utils import construct_invoke_connection_request, SkyflowMessages, get_metrics, \
     parse_invoke_connection_response
-from skyflow.utils.logger import log_info
+from skyflow.utils.logger import log_info, log_error_log
 from skyflow.vault.connection import InvokeConnectionRequest
 
 
@@ -27,7 +27,7 @@ def invoke(self, request: InvokeConnectionRequest):
 
         invoke_connection_request.headers['sky-metadata'] = json.dumps(get_metrics())
 
-        log_info(SkyflowMessages.Info.INVOKE_CONNECTION_TRIGGERED, self.__vault_client.get_logger())
+        log_info(SkyflowMessages.Info.INVOKE_CONNECTION_TRIGGERED.value, self.__vault_client.get_logger())
 
         try:
             response = session.send(invoke_connection_request)
@@ -36,5 +36,7 @@ def invoke(self, request: InvokeConnectionRequest):
             return invoke_connection_response
 
         except Exception as e:
+            log_error_log(SkyflowMessages.ErrorLogs.INVOKE_CONNECTION_REQUEST_REJECTED.value, self.__vault_client.get_logger())
+            if isinstance(e, SkyflowError): raise e
             raise SkyflowError(SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value,
                                SkyflowMessages.ErrorCodes.SERVER_ERROR.value)

tests/utils/test__utils.py

@@ -356,7 +356,8 @@ def test_parse_invoke_connection_response_http_error_with_json_error_message(sel
         with self.assertRaises(SkyflowError) as context:
             parse_invoke_connection_response(mock_response)
 
-        self.assertEqual(context.exception.message, "Not Found - request id: 1234")
+        self.assertEqual(context.exception.message, "Not Found")
+        self.assertEqual(context.exception.request_id, "1234")
 
     @patch("requests.Response")
     def test_parse_invoke_connection_response_http_error_without_json_error_message(self, mock_response):
@@ -369,7 +370,7 @@ def test_parse_invoke_connection_response_http_error_without_json_error_message(
         with self.assertRaises(SkyflowError) as context:
             parse_invoke_connection_response(mock_response)
 
-        self.assertEqual(context.exception.message, SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format("Internal Server Error") + " - request id: 1234")
+        self.assertEqual(context.exception.message, SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format("Internal Server Error"))
 
     @patch("skyflow.utils._utils.log_and_reject_error")
     def test_handle_exception_json_error(self, mock_log_and_reject_error):

tests/vault/controller/test__connection.py

@@ -1,9 +1,10 @@
 import unittest
 from unittest.mock import Mock, patch
-
+import requests
 from skyflow.error import SkyflowError
-from skyflow.utils import SkyflowMessages
+from skyflow.utils import SkyflowMessages, parse_invoke_connection_response
 from skyflow.utils.enums import RequestMethod
+from skyflow.utils._version import SDK_VERSION
 from skyflow.vault.connection import InvokeConnectionRequest
 from skyflow.vault.controller import Connection
 
@@ -21,7 +22,8 @@
 INVALID_HEADERS = "invalid_headers"
 INVALID_BODY = "invalid_body"
 FAILURE_STATUS_CODE = 400
-ERROR_RESPONSE_CONTENT = '{"error": {"message": "error occurred"}}'
+ERROR_RESPONSE_CONTENT = '"message": "Invalid Request"'
+ERROR_FROM_CLIENT_RESPONSE_CONTENT = b'{"error": {"message": "Invalid Request"}}'
 
 class TestConnection(unittest.TestCase):
     def setUp(self):
@@ -99,6 +101,25 @@ def test_invoke_request_error(self, mock_send):
 
         with self.assertRaises(SkyflowError) as context:
             self.connection.invoke(request)
-        self.assertEqual(context.exception.message, SkyflowMessages.Error.INVOKE_CONNECTION_FAILED.value)
+        self.assertEqual(context.exception.message, f'Skyflow Python SDK {SDK_VERSION} Response {ERROR_RESPONSE_CONTENT} is not valid JSON.')
+        self.assertEqual(context.exception.message, SkyflowMessages.Error.RESPONSE_NOT_JSON.value.format(ERROR_RESPONSE_CONTENT))
+        self.assertEqual(context.exception.http_code, 400)
+
+    def test_parse_invoke_connection_response_error_from_client(self):
+        mock_response = Mock(spec=requests.Response)
+        mock_response.status_code = FAILURE_STATUS_CODE
+        mock_response.content = ERROR_FROM_CLIENT_RESPONSE_CONTENT
+        mock_response.headers = {
+            'error-from-client': 'true',
+            'x-request-id': '12345'
+        }
+        mock_response.raise_for_status.side_effect = requests.HTTPError()
+
+        with self.assertRaises(SkyflowError) as context:
+            parse_invoke_connection_response(mock_response)
+
+        exception = context.exception
 
+        self.assertTrue(any(detail.get('error_from_client') == 'true' for detail in exception.details))
+        self.assertEqual(exception.request_id, '12345')