Merge pull request #498 from rejasupotaro/raise-custom-error-with-mes…

…sage

Include Vespa error message in custom error raised

vespa/application.py

@@ -12,12 +12,13 @@
 from pandas import DataFrame
 from requests import Session
 from requests.models import Response
-from requests.exceptions import ConnectionError
+from requests.exceptions import ConnectionError, HTTPError, JSONDecodeError
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
 from tenacity import retry, wait_exponential, stop_after_attempt
 from time import sleep
 
+from vespa.exceptions import VespaError
 from vespa.io import VespaQueryResponse, VespaResponse
 from vespa.package import ApplicationPackage
 
@@ -55,6 +56,29 @@ def parse_feed_df(df: DataFrame, include_id: bool, id_field="id") -> List[Dict[s
     return batch
 
 
+def raise_for_status(response: Response) -> None:
+    """
+    Raises an appropriate error if necessary.
+
+    If the response contains an error message, VespaError is raised along with HTTPError to provide more details.
+
+    :param response: Response object from Vespa API.
+    :raises HTTPError: If status_code is between 400 and 599.
+    :raises VespaError: If the response JSON contains an error message.
+    """
+    try:
+        response.raise_for_status()
+    except HTTPError as http_error:
+        try:
+            response_json = response.json()
+        except JSONDecodeError:
+            raise http_error
+        errors = response_json.get("root", {}).get("errors", [])
+        if not errors:
+            raise http_error
+        raise VespaError(errors) from http_error
+
+
 class Vespa(object):
     def __init__(
         self,
@@ -836,7 +860,7 @@ def feed_data_point(
         )
         vespa_format = {"fields": fields}
         response = self.http_session.post(end_point, json=vespa_format, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return VespaResponse(
             json=response.json(),
             status_code=response.status_code,
@@ -858,7 +882,7 @@ def query(
         :raises HTTPError: if one occurred
         """
         response = self.http_session.post(self.app.search_end_point, json=body, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return VespaQueryResponse(
             json=response.json(), status_code=response.status_code, url=str(response.url)
         )
@@ -882,7 +906,7 @@ def delete_data(
             self.app.end_point, namespace, schema, str(data_id)
         )
         response = self.http_session.delete(end_point, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return VespaResponse(
             json=response.json(),
             status_code=response.status_code,
@@ -909,7 +933,7 @@ def delete_all_docs(
             self.app.end_point, namespace, schema, content_cluster_name
         )
         response = self.http_session.delete(end_point, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return response
 
     def get_data(
@@ -931,7 +955,7 @@ def get_data(
             self.app.end_point, namespace, schema, str(data_id)
         )
         response = self.http_session.get(end_point, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return VespaResponse(
             json=response.json(),
             status_code=response.status_code,
@@ -966,7 +990,7 @@ def update_data(
         )
         vespa_format = {"fields": {k: {"assign": v} for k, v in fields.items()}}
         response = self.http_session.put(end_point, json=vespa_format, cert=self.cert)
-        response.raise_for_status()
+        raise_for_status(response)
         return VespaResponse(
             json=response.json(),
             status_code=response.status_code,

vespa/exceptions.py

@@ -0,0 +1,2 @@
+class VespaError(Exception):
+    """Vespa returned an error response"""

vespa/test_application.py

@@ -1,10 +1,15 @@
 # Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
+import json
 import unittest
+import pytest
+from unittest.mock import PropertyMock, patch
 from pandas import DataFrame
+from requests.models import HTTPError, Response
 
 from vespa.package import ApplicationPackage, Schema, Document
-from vespa.application import Vespa, parse_feed_df
+from vespa.application import Vespa, parse_feed_df, raise_for_status
+from vespa.exceptions import VespaError
 
 
 class TestVespa(unittest.TestCase):
@@ -126,6 +131,84 @@ def test_parse_simplified_feed_batch_with_wrong_columns(self):
             _ = parse_feed_df(df=missing_id_df, include_id=True)
 
 
+
+class TestRaiseForStatus(unittest.TestCase):
+    def test_successful_response(self):
+        response = Response()
+        response.status_code = 200
+        try:
+            raise_for_status(response)
+        except Exception as e:
+            self.fail(f"No exceptions were expected to be raised but {type(e).__name__} occurred")
+
+    def test_successful_response_with_error_content(self):
+        with patch("requests.models.Response.content", new_callable=PropertyMock) as mock_content:
+            response_json = {
+                "root": {
+                    "errors": [
+                        {"code": 1, "summary": "summary", "message": "message"},
+                    ],
+                },
+            }
+            mock_content.return_value = json.dumps(response_json).encode("utf-8")
+            response = Response()
+            response.status_code = 200
+            try:
+                raise_for_status(response)
+            except Exception as e:
+                self.fail(f"No exceptions were expected to be raised but {type(e).__name__} occurred")
+
+    def test_failure_response_for_400(self):
+        response = Response()
+        response.status_code = 400
+        response.reason = "reason"
+        response.url = "http://localhost:8080"
+        with pytest.raises(HTTPError) as e:
+            raise_for_status(response)
+        self.assertEqual(str(e.value), "400 Client Error: reason for url: http://localhost:8080")
+
+    def test_failure_response_for_500(self):
+        response = Response()
+        response.status_code = 500
+        response.reason = "reason"
+        response.url = "http://localhost:8080"
+        with pytest.raises(HTTPError) as e:
+            raise_for_status(response)
+        self.assertEqual(str(e.value), "500 Server Error: reason for url: http://localhost:8080")
+
+    def test_failure_response_without_error_content(self):
+        with patch("requests.models.Response.content", new_callable=PropertyMock) as mock_content:
+            response_json = {
+                "root": {
+                    "errors": [],
+                },
+            }
+            mock_content.return_value = json.dumps(response_json).encode("utf-8")
+            response = Response()
+            response.status_code = 400
+            response.reason = "reason"
+            response.url = "http://localhost:8080"
+            with pytest.raises(HTTPError):
+                raise_for_status(response)
+
+    def test_failure_response_with_error_content(self):
+        with patch("requests.models.Response.content", new_callable=PropertyMock) as mock_content:
+            response_json = {
+                "root": {
+                    "errors": [
+                        {"code": 1, "summary": "summary", "message": "message"},
+                    ],
+                },
+            }
+            mock_content.return_value = json.dumps(response_json).encode("utf-8")
+            response = Response()
+            response.status_code = 400
+            response.reason = "reason"
+            response.url = "http://localhost:8080"
+            with pytest.raises(VespaError):
+                raise_for_status(response)
+
+
 class TestVespaCollectData(unittest.TestCase):
     def setUp(self) -> None:
         self.app = Vespa(url="http://localhost", port=8080)

vespa/test_integration_docker.py

@@ -24,6 +24,7 @@
 )
 from vespa.deployment import VespaDocker
 from vespa.application import VespaSync
+from vespa.exceptions import VespaError
 
 
 CONTAINER_STOP_TIMEOUT = 10
@@ -210,7 +211,7 @@ def redeploy_with_container_stopped(self, application_package):
     def redeploy_with_application_package_changes(self, application_package):
         self.vespa_docker = VespaDocker(port=8089)
         app = self.vespa_docker.deploy(application_package=application_package)
-        with pytest.raises(HTTPError):
+        with pytest.raises(VespaError):
             app.query(
                 body={
                     "yql": "select * from sources * where default contains 'music'",