Add integration tests for upsert

Author: jhamon

pinecone/exceptions/__init__.py

@@ -12,6 +12,8 @@
 )
 from .exceptions import PineconeConfigurationError, PineconeProtocolError, ListConversionException
 
+PineconeNotFoundException = NotFoundException
+
 __all__ = [
     "PineconeConfigurationError",
     "PineconeProtocolError",
@@ -22,6 +24,7 @@
     "PineconeApiKeyError",
     "PineconeApiException",
     "NotFoundException",
+    "PineconeNotFoundException",
     "UnauthorizedException",
     "ForbiddenException",
     "ServiceException",

pinecone/grpc/grpc_runner.py

@@ -8,10 +8,19 @@
 from .utils import _generate_request_id
 from .config import GRPCClientConfig
 from pinecone.utils.constants import REQUEST_ID, CLIENT_VERSION
-from pinecone.exceptions.exceptions import PineconeException
-from grpc import CallCredentials, Compression
+from grpc import CallCredentials, Compression, StatusCode
+from grpc.aio import AioRpcError
 from google.protobuf.message import Message
 
+from pinecone.exceptions import (
+    PineconeException,
+    PineconeApiValueError,
+    PineconeApiException,
+    UnauthorizedException,
+    PineconeNotFoundException,
+    ServiceException,
+)
+
 
 class GrpcRunner:
     def __init__(self, index_name: str, config: Config, grpc_config: GRPCClientConfig):
@@ -50,7 +59,7 @@ def wrapped():
                     compression=compression,
                 )
             except _InactiveRpcError as e:
-                raise PineconeException(e._state.debug_error_string) from e
+                self._map_exception(e, e._state.code, e._state.details)
 
         return wrapped()
 
@@ -89,8 +98,8 @@ async def wrapped():
                         wait_for_ready=wait_for_ready,
                         compression=compression,
                     )
-            except _InactiveRpcError as e:
-                raise PineconeException(e._state.debug_error_string) from e
+            except AioRpcError as e:
+                self._map_exception(e, e.code(), e.details())
 
         return await wrapped()
 
@@ -108,3 +117,37 @@ def _prepare_metadata(
 
     def _request_metadata(self) -> Dict[str, str]:
         return {REQUEST_ID: _generate_request_id()}
+
+    def _map_exception(self, e: Exception, code: Optional[StatusCode], details: Optional[str]):
+        # Client / connection issues
+        details = details or ""
+
+        if code in [StatusCode.DEADLINE_EXCEEDED]:
+            raise TimeoutError(details) from e
+
+        # Permissions stuff
+        if code in [StatusCode.PERMISSION_DENIED, StatusCode.UNAUTHENTICATED]:
+            raise UnauthorizedException(status=code, reason=details) from e
+
+        # 400ish stuff
+        if code in [StatusCode.NOT_FOUND]:
+            raise PineconeNotFoundException(status=code, reason=details) from e
+        if code in [StatusCode.INVALID_ARGUMENT, StatusCode.OUT_OF_RANGE]:
+            raise PineconeApiValueError(details) from e
+        if code in [
+            StatusCode.ALREADY_EXISTS,
+            StatusCode.FAILED_PRECONDITION,
+            StatusCode.UNIMPLEMENTED,
+            StatusCode.RESOURCE_EXHAUSTED,
+        ]:
+            raise PineconeApiException(status=code, reason=details) from e
+
+        # 500ish stuff
+        if code in [StatusCode.INTERNAL, StatusCode.UNAVAILABLE]:
+            raise ServiceException(status=code, reason=details) from e
+        if code in [StatusCode.UNKNOWN, StatusCode.DATA_LOSS, StatusCode.ABORTED]:
+            # abandon hope, all ye who enter here
+            raise PineconeException(code, details) from e
+
+        # If you get here, you're in a bad place
+        raise PineconeException(code, details) from e

tests/integration/data_asyncio/__init__.py

@@ -0,0 +1,55 @@
+import pytest
+import os
+from ..helpers import get_environment_var, random_string
+
+
+@pytest.fixture(scope="session")
+def api_key():
+    return get_environment_var("PINECONE_API_KEY")
+
+
+@pytest.fixture(scope="session")
+def host():
+    return get_environment_var("INDEX_HOST")
+
+
+@pytest.fixture(scope="session")
+def dimension():
+    return int(get_environment_var("DIMENSION"))
+
+
+def use_grpc():
+    return os.environ.get("USE_GRPC", "false") == "true"
+
+
+def build_client(api_key):
+    if use_grpc():
+        from pinecone.grpc import PineconeGRPC
+
+        return PineconeGRPC(api_key=api_key)
+    else:
+        from pinecone import Pinecone
+
+        return Pinecone(
+            api_key=api_key, additional_headers={"sdk-test-suite": "pinecone-python-client"}
+        )
+
+
+@pytest.fixture(scope="session")
+async def pc(api_key):
+    return build_client(api_key=api_key)
+
+
+@pytest.fixture(scope="session")
+async def asyncio_idx(pc, host):
+    return pc.AsyncioIndex(host=host)
+
+
+@pytest.fixture(scope="session")
+async def namespace():
+    return random_string(10)
+
+
+@pytest.fixture(scope="session")
+async def list_namespace():
+    return random_string(10)

tests/integration/data_asyncio/conftest.py

@@ -0,0 +1,97 @@
+import pytest
+from pinecone import Vector
+from .conftest import use_grpc
+from ..helpers import random_string
+from .utils import build_asyncio_idx, embedding_values, poll_for_freshness
+
+
+@pytest.mark.parametrize("target_namespace", ["", random_string(20)])
+@pytest.mark.skipif(use_grpc() == False, reason="Currently only GRPC supports asyncio")
+async def test_upsert_to_default_namespace(host, dimension, target_namespace):
+    asyncio_idx = build_asyncio_idx(host)
+
+    def emb():
+        return embedding_values(dimension)
+
+    # Upsert with tuples
+    await asyncio_idx.upsert(
+        vectors=[("1", emb()), ("2", emb()), ("3", emb())], namespace=target_namespace
+    )
+
+    # Upsert with objects
+    await asyncio_idx.upsert(
+        vectors=[
+            Vector(id="4", values=emb()),
+            Vector(id="5", values=emb()),
+            Vector(id="6", values=emb()),
+        ],
+        namespace=target_namespace,
+    )
+
+    # Upsert with dict
+    await asyncio_idx.upsert(
+        vectors=[
+            {"id": "7", "values": emb()},
+            {"id": "8", "values": emb()},
+            {"id": "9", "values": emb()},
+        ],
+        namespace=target_namespace,
+    )
+
+    await poll_for_freshness(asyncio_idx, target_namespace, 9)
+
+    # # Check the vector count reflects some data has been upserted
+    stats = await asyncio_idx.describe_index_stats()
+    assert stats.total_vector_count >= 9
+    # default namespace could have other stuff from other tests
+    if target_namespace != "":
+        assert stats.namespaces[target_namespace].vector_count == 9
+
+
+# @pytest.mark.parametrize("target_namespace", [
+#     "",
+#     random_string(20),
+# ])
+# @pytest.mark.skipif(
+#     os.getenv("METRIC") != "dotproduct", reason="Only metric=dotprodouct indexes support hybrid"
+# )
+# async def test_upsert_to_namespace_with_sparse_embedding_values(pc, host, dimension, target_namespace):
+#     asyncio_idx = pc.AsyncioIndex(host=host)
+
+#     # Upsert with sparse values object
+#     await asyncio_idx.upsert(
+#         vectors=[
+#             Vector(
+#                 id="1",
+#                 values=embedding_values(dimension),
+#                 sparse_values=SparseValues(indices=[0, 1], values=embedding_values()),
+#             )
+#         ],
+#         namespace=target_namespace,
+#     )
+
+#     # Upsert with sparse values dict
+#     await asyncio_idx.upsert(
+#         vectors=[
+#             {
+#                 "id": "2",
+#                 "values": embedding_values(dimension),
+#                 "sparse_values": {"indices": [0, 1], "values": embedding_values()},
+#             },
+#             {
+#                 "id": "3",
+#                 "values": embedding_values(dimension),
+#                 "sparse_values": {"indices": [0, 1], "values": embedding_values()},
+#             },
+#         ],
+#         namespace=target_namespace,
+#     )
+
+#     await poll_for_freshness(asyncio_idx, target_namespace, 9)
+
+#     # Check the vector count reflects some data has been upserted
+#     stats = await asyncio_idx.describe_index_stats()
+#     assert stats.total_vector_count >= 9
+
+#     if (target_namespace != ""):
+#         assert stats.namespaces[target_namespace].vector_count == 9