diff --git a/tests/test_vectors.py b/tests/test_vectors.py index 57b4266..e549589 100644 --- a/tests/test_vectors.py +++ b/tests/test_vectors.py @@ -1,6 +1,7 @@ import uuid import turbopuffer as tpuf import tests +import pytest from datetime import datetime @@ -485,3 +486,12 @@ def test_attribute_types(): ]], ) assert len(results) == 1 + +def test_not_found_error(): + ns = tpuf.Namespace(tests.test_prefix + 'not_found') + + with pytest.raises(tpuf.NotFoundError): + ns.query( + top_k=5, + vector=[0.0, 0.0], + ) diff --git a/turbopuffer/__init__.py b/turbopuffer/__init__.py index d5a0bf0..4ff5a3f 100644 --- a/turbopuffer/__init__.py +++ b/turbopuffer/__init__.py @@ -28,4 +28,4 @@ def dump_json_bytes(obj): return json.dumps(obj, cls=NumpyEncoder).encode() from turbopuffer.namespace import Namespace, namespaces from turbopuffer.vectors import VectorColumns, VectorRow, VectorResult from turbopuffer.query import VectorQuery, Filters -from turbopuffer.error import TurbopufferError, AuthenticationError, APIError +from turbopuffer.error import TurbopufferError, AuthenticationError, APIError, NotFoundError diff --git a/turbopuffer/backend.py b/turbopuffer/backend.py index b5b2992..a7c9660 100644 --- a/turbopuffer/backend.py +++ b/turbopuffer/backend.py @@ -5,7 +5,7 @@ import requests import turbopuffer as tpuf import gzip -from turbopuffer.error import AuthenticationError, APIError +from turbopuffer.error import AuthenticationError, raise_api_error from typing import Optional, List @@ -121,7 +121,7 @@ def make_api_request(self, try: content = response.json() except json.JSONDecodeError as err: - raise APIError(response.status_code, traceback.format_exception_only(err), response.text) + raise_api_error(response.status_code, traceback.format_exception_only(err), response.text) if response.ok: performance['total_time'] = time.monotonic() - start @@ -130,9 +130,9 @@ def make_api_request(self, 'performance': performance, }) else: - raise APIError(response.status_code, content.get('status', 'error'), content.get('error', '')) + raise_api_error(response.status_code, content.get('status', 'error'), content.get('error', '')) else: - raise APIError(response.status_code, 'Server returned non-JSON response', response.text) + raise_api_error(response.status_code, 'Server returned non-JSON response', response.text) except requests.HTTPError as http_err: retry_attempt += 1 # print(traceback.format_exc()) @@ -141,6 +141,6 @@ def make_api_request(self, time.sleep(2 ** retry_attempt) # exponential falloff up to 64 seconds for 6 retries. else: print(f'Request failed after {retry_attempt} attempts...') - raise APIError(http_err.response.status_code, + raise_api_error(http_err.response.status_code, f'Request to {http_err.request.url} failed after {retry_attempt} attempts', str(http_err)) diff --git a/turbopuffer/error.py b/turbopuffer/error.py index 51760bd..1bcb38c 100644 --- a/turbopuffer/error.py +++ b/turbopuffer/error.py @@ -11,3 +11,12 @@ def __init__(self, status_code: int, status_name: str, message: str): self.status_code = status_code self.status_name = status_name super().__init__(f'{status_name} (HTTP {status_code}): {message}') + +class NotFoundError(APIError): + pass + +def raise_api_error(status_code, status_name, message): + if status_code == 404: + raise NotFoundError(status_code, status_name, message) + else: + raise APIError(status_code, status_name, message)