diff --git a/README.md b/README.md index f9a4748..8458770 100644 --- a/README.md +++ b/README.md @@ -329,6 +329,26 @@ for ns, ns_info in info.namespaces.items(): ) ``` +### List Namespaces + +All the names of active namespaces can be listed. + +```python +namespaces = index.list_namespaces() +for ns in namespaces: + print(ns) # name of the namespace +``` + +### Delete a Namespace + +A namespace can be deleted entirely. +If no such namespace exists, and exception is raised. +The default namespaces cannot be deleted. + +```python +index.delete_namespace(namespace="ns") +``` + # Contributing ## Preparing the environment diff --git a/tests/__init__.py b/tests/__init__.py index 3196cb3..146d250 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,7 +1,10 @@ import time import asyncio -NAMESPACES = ["", "ns"] +from upstash_vector import Index, AsyncIndex +from upstash_vector.core.index_operations import DEFAULT_NAMESPACE + +NAMESPACES = [DEFAULT_NAMESPACE, "ns"] def assert_eventually(assertion, retry_delay=0.5, timeout=5.0): @@ -38,3 +41,51 @@ async def assert_eventually_async(assertion, retry_delay=0.5, timeout=5.0): raise AssertionError("Couldn't run the assertion") raise last_err + + +def ensure_ns_exists(index: Index, ns: str): + """ + Ensures the given namespace exists in the index by upserting some + random vector into it, and calling reset. + + No need to call this method, if you are upserting some vector/data. + """ + if ns == DEFAULT_NAMESPACE: + return + + index.upsert( + vectors=[("0", [0.1, 0.1])], + namespace=ns, + ) + + def assertion(): + info = index.info() + assert info.namespaces[ns].pending_vector_count == 0 + + assert_eventually(assertion) + + index.reset(namespace=ns) + + +async def ensure_ns_exists_async(index: AsyncIndex, ns: str): + """ + Ensures the given namespace exists in the index by upserting some + random vector into it, and calling reset. + + No need to call this method, if you are upserting some vector/data. + """ + if ns == DEFAULT_NAMESPACE: + return + + await index.upsert( + vectors=[("0", [0.1, 0.1])], + namespace=ns, + ) + + async def assertion(): + info = await index.info() + assert info.namespaces[ns].pending_vector_count == 0 + + await assert_eventually_async(assertion) + + await index.reset(namespace=ns) diff --git a/tests/core/test_info.py b/tests/core/test_info.py index 5039657..7b83345 100644 --- a/tests/core/test_info.py +++ b/tests/core/test_info.py @@ -1,10 +1,19 @@ import pytest -from tests import assert_eventually, assert_eventually_async, NAMESPACES +from tests import ( + assert_eventually, + assert_eventually_async, + NAMESPACES, + ensure_ns_exists, + ensure_ns_exists_async, +) from upstash_vector import Index, AsyncIndex def test_info(index: Index): + for ns in NAMESPACES: + ensure_ns_exists(index, ns) + info = index.info() assert info.vector_count == 0 @@ -32,6 +41,9 @@ def assertion(): @pytest.mark.asyncio async def test_info_async(async_index: AsyncIndex): + for ns in NAMESPACES: + await ensure_ns_exists_async(async_index, ns) + info = await async_index.info() assert info.vector_count == 0 diff --git a/tests/core/test_namespace_operations.py b/tests/core/test_namespace_operations.py new file mode 100644 index 0000000..b2399fc --- /dev/null +++ b/tests/core/test_namespace_operations.py @@ -0,0 +1,63 @@ +import pytest + +from tests import NAMESPACES, ensure_ns_exists, ensure_ns_exists_async +from upstash_vector import Index, AsyncIndex +from upstash_vector.core.index_operations import DEFAULT_NAMESPACE + + +def test_list_namespaces(index: Index): + for ns in NAMESPACES: + ensure_ns_exists(index, ns) + + all_ns = index.list_namespaces() + + assert len(all_ns) == len(NAMESPACES) + assert NAMESPACES[0] in all_ns + assert NAMESPACES[1] in all_ns + + +@pytest.mark.asyncio +async def test_list_namespaces_async(async_index: AsyncIndex): + for ns in NAMESPACES: + await ensure_ns_exists_async(async_index, ns) + + all_ns = await async_index.list_namespaces() + + assert len(all_ns) == len(NAMESPACES) + assert NAMESPACES[0] in all_ns + assert NAMESPACES[1] in all_ns + + +def test_delete_namespaces(index: Index): + for ns in NAMESPACES: + ensure_ns_exists(index, ns) + + for ns in NAMESPACES: + if ns == DEFAULT_NAMESPACE: + continue + + # Should not fail + index.delete_namespace(namespace=ns) + + info = index.info() + + # Only default namespace should exist + assert len(info.namespaces) == 1 + + +@pytest.mark.asyncio +async def test_delete_namespaces_async(async_index: AsyncIndex): + for ns in NAMESPACES: + await ensure_ns_exists_async(async_index, ns) + + for ns in NAMESPACES: + if ns == DEFAULT_NAMESPACE: + continue + + # Should not fail + await async_index.delete_namespace(namespace=ns) + + info = await async_index.info() + + # Only default namespace should exist + assert len(info.namespaces) == 1 diff --git a/upstash_vector/core/index_operations.py b/upstash_vector/core/index_operations.py index 15687ef..214827c 100644 --- a/upstash_vector/core/index_operations.py +++ b/upstash_vector/core/index_operations.py @@ -27,6 +27,8 @@ RANGE_PATH = "/range" FETCH_PATH = "/fetch" INFO_PATH = "/info" +LIST_NAMESPACES = "/list-namespaces" +DELETE_NAMESPACE = "/delete-namespace" def _path_for(namespace: str, path: str) -> str: @@ -320,6 +322,19 @@ def info(self) -> InfoResult: self._execute_request(payload=None, path=INFO_PATH) ) + def list_namespaces(self) -> List[str]: + """ + Returns the list of names of namespaces. + """ + return self._execute_request(payload=None, path=LIST_NAMESPACES) + + def delete_namespace(self, namespace: str) -> None: + """ + Deletes the given namespace if it exists, or raises + exception if no such namespace exists. + """ + self._execute_request(payload=None, path=_path_for(namespace, DELETE_NAMESPACE)) + class AsyncIndexOperations: async def _execute_request_async(self, payload, path): @@ -606,3 +621,19 @@ async def info(self) -> InfoResult: return InfoResult._from_json( await self._execute_request_async(payload=None, path=INFO_PATH) ) + + async def list_namespaces(self) -> List[str]: + """ + Returns the list of names of namespaces. + """ + result = await self._execute_request_async(payload=None, path=LIST_NAMESPACES) + return result + + async def delete_namespace(self, namespace: str) -> None: + """ + Deletes the given namespace if it exists, or raises + exception if no such namespace exists. + """ + await self._execute_request_async( + payload=None, path=_path_for(namespace, DELETE_NAMESPACE) + )