fix(vector-stores): scope Pinecone delete_col()/reset() to namespace (#6287)
This commit is contained in:
@@ -12,8 +12,9 @@ const upsert = jest.fn().mockResolvedValue(undefined);
|
|||||||
const query = jest.fn().mockResolvedValue({ matches: [] });
|
const query = jest.fn().mockResolvedValue({ matches: [] });
|
||||||
const fetch = jest.fn().mockResolvedValue({ records: {} });
|
const fetch = jest.fn().mockResolvedValue({ records: {} });
|
||||||
const deleteOne = jest.fn().mockResolvedValue(undefined);
|
const deleteOne = jest.fn().mockResolvedValue(undefined);
|
||||||
|
const deleteAll = jest.fn().mockResolvedValue(undefined);
|
||||||
|
|
||||||
const nsHandle = { upsert, query, fetch, deleteOne };
|
const nsHandle = { upsert, query, fetch, deleteOne, deleteAll };
|
||||||
const namespace = jest.fn().mockReturnValue(nsHandle);
|
const namespace = jest.fn().mockReturnValue(nsHandle);
|
||||||
|
|
||||||
const describeIndexStats = jest
|
const describeIndexStats = jest
|
||||||
@@ -47,6 +48,7 @@ const __mocks__ = {
|
|||||||
query,
|
query,
|
||||||
fetch,
|
fetch,
|
||||||
deleteOne,
|
deleteOne,
|
||||||
|
deleteAll,
|
||||||
namespace,
|
namespace,
|
||||||
describeIndexStats,
|
describeIndexStats,
|
||||||
index,
|
index,
|
||||||
@@ -94,6 +96,7 @@ beforeEach(() => {
|
|||||||
__mocks__.query.mockResolvedValue({ matches: [] });
|
__mocks__.query.mockResolvedValue({ matches: [] });
|
||||||
__mocks__.fetch.mockResolvedValue({ records: {} });
|
__mocks__.fetch.mockResolvedValue({ records: {} });
|
||||||
__mocks__.deleteOne.mockResolvedValue(undefined);
|
__mocks__.deleteOne.mockResolvedValue(undefined);
|
||||||
|
__mocks__.deleteAll.mockResolvedValue(undefined);
|
||||||
__mocks__.describeIndexStats.mockResolvedValue({
|
__mocks__.describeIndexStats.mockResolvedValue({
|
||||||
totalRecordCount: 0,
|
totalRecordCount: 0,
|
||||||
namespaces: {},
|
namespaces: {},
|
||||||
@@ -104,6 +107,7 @@ beforeEach(() => {
|
|||||||
query: __mocks__.query,
|
query: __mocks__.query,
|
||||||
fetch: __mocks__.fetch,
|
fetch: __mocks__.fetch,
|
||||||
deleteOne: __mocks__.deleteOne,
|
deleteOne: __mocks__.deleteOne,
|
||||||
|
deleteAll: __mocks__.deleteAll,
|
||||||
};
|
};
|
||||||
__mocks__.namespace.mockReturnValue(nsHandle);
|
__mocks__.namespace.mockReturnValue(nsHandle);
|
||||||
__mocks__.index.mockReturnValue({
|
__mocks__.index.mockReturnValue({
|
||||||
@@ -424,6 +428,14 @@ describe("deleteCol", () => {
|
|||||||
await db.initialize();
|
await db.initialize();
|
||||||
expect(__mocks__.createIndex).toHaveBeenCalledTimes(2);
|
expect(__mocks__.createIndex).toHaveBeenCalledTimes(2);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("clears only the namespace and never drops the shared index", async () => {
|
||||||
|
const db = await initDb({ namespace: "tenant-a" });
|
||||||
|
await db.deleteCol();
|
||||||
|
expect(__mocks__.namespace).toHaveBeenCalledWith("tenant-a");
|
||||||
|
expect(__mocks__.deleteAll).toHaveBeenCalled();
|
||||||
|
expect(__mocks__.deleteIndex).not.toHaveBeenCalled();
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
describe("list", () => {
|
describe("list", () => {
|
||||||
|
|||||||
@@ -316,6 +316,11 @@ export class PineconeDB implements VectorStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async deleteCol(): Promise<void> {
|
async deleteCol(): Promise<void> {
|
||||||
|
if (this.namespace) {
|
||||||
|
await this.initialize();
|
||||||
|
await this.namespacedIndex().deleteAll();
|
||||||
|
return;
|
||||||
|
}
|
||||||
if (this._initPromise) {
|
if (this._initPromise) {
|
||||||
await this._initPromise.catch(() => {});
|
await this._initPromise.catch(() => {});
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -358,12 +358,16 @@ class PineconeDB(VectorStoreBase):
|
|||||||
return self.client.list_indexes()
|
return self.client.list_indexes()
|
||||||
|
|
||||||
def delete_col(self):
|
def delete_col(self):
|
||||||
"""Delete an index/collection."""
|
"""Delete the index, or clear only the configured namespace if one is set."""
|
||||||
try:
|
try:
|
||||||
self.client.delete_index(self.collection_name)
|
if self.namespace is not None:
|
||||||
logger.info(f"Index {self.collection_name} deleted successfully")
|
self.index.delete(delete_all=True, namespace=self.namespace)
|
||||||
|
logger.info(f"Namespace {self.namespace} in index {self.collection_name} cleared successfully")
|
||||||
|
else:
|
||||||
|
self.client.delete_index(self.collection_name)
|
||||||
|
logger.info(f"Index {self.collection_name} deleted successfully")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error deleting index {self.collection_name}: {e}")
|
logger.error(f"Error deleting index {self.collection_name} (namespace={self.namespace}): {e}")
|
||||||
|
|
||||||
def col_info(self) -> Dict:
|
def col_info(self) -> Dict:
|
||||||
"""
|
"""
|
||||||
@@ -419,8 +423,7 @@ class PineconeDB(VectorStoreBase):
|
|||||||
int: Total number of vectors.
|
int: Total number of vectors.
|
||||||
"""
|
"""
|
||||||
stats = self.index.describe_index_stats()
|
stats = self.index.describe_index_stats()
|
||||||
if self.namespace:
|
if self.namespace is not None:
|
||||||
# Safely get the namespace stats and return vector_count, defaulting to 0 if not found
|
|
||||||
namespace_summary = (stats.namespaces or {}).get(self.namespace)
|
namespace_summary = (stats.namespaces or {}).get(self.namespace)
|
||||||
if namespace_summary:
|
if namespace_summary:
|
||||||
return namespace_summary.vector_count or 0
|
return namespace_summary.vector_count or 0
|
||||||
|
|||||||
@@ -139,9 +139,40 @@ def test_list_cols(pinecone_db):
|
|||||||
pinecone_db.client.list_indexes.assert_called()
|
pinecone_db.client.list_indexes.assert_called()
|
||||||
|
|
||||||
|
|
||||||
def test_delete_col(pinecone_db):
|
@pytest.mark.parametrize(
|
||||||
|
"namespace,drops_whole_index",
|
||||||
|
[(None, True), ("test_namespace", False), ("", False)],
|
||||||
|
)
|
||||||
|
def test_delete_col_scopes_to_namespace(pinecone_db, namespace, drops_whole_index):
|
||||||
|
pinecone_db.namespace = namespace
|
||||||
pinecone_db.delete_col()
|
pinecone_db.delete_col()
|
||||||
pinecone_db.client.delete_index.assert_called_with("test_index")
|
if drops_whole_index:
|
||||||
|
pinecone_db.client.delete_index.assert_called_with("test_index")
|
||||||
|
else:
|
||||||
|
pinecone_db.index.delete.assert_called_with(delete_all=True, namespace=namespace)
|
||||||
|
pinecone_db.client.delete_index.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_delete_col_namespace_delete_error_is_swallowed(pinecone_db):
|
||||||
|
pinecone_db.index.delete.side_effect = Exception("Namespace not found")
|
||||||
|
pinecone_db.delete_col()
|
||||||
|
|
||||||
|
|
||||||
|
def test_reset_with_namespace_does_not_drop_or_recreate_index(pinecone_db, mock_pinecone_client):
|
||||||
|
mock_pinecone_client.list_indexes.return_value.names.return_value = ["test_index"]
|
||||||
|
mock_pinecone_client.create_index.reset_mock()
|
||||||
|
pinecone_db.reset()
|
||||||
|
pinecone_db.index.delete.assert_called_with(delete_all=True, namespace="test_namespace")
|
||||||
|
mock_pinecone_client.delete_index.assert_not_called()
|
||||||
|
mock_pinecone_client.create_index.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_reset_without_namespace_drops_and_recreates_index(pinecone_db, mock_pinecone_client):
|
||||||
|
pinecone_db.namespace = None
|
||||||
|
mock_pinecone_client.list_indexes.return_value.names.return_value = []
|
||||||
|
pinecone_db.reset()
|
||||||
|
mock_pinecone_client.delete_index.assert_called_with("test_index")
|
||||||
|
mock_pinecone_client.create_index.assert_called()
|
||||||
|
|
||||||
|
|
||||||
def test_col_info(pinecone_db):
|
def test_col_info(pinecone_db):
|
||||||
@@ -190,6 +221,17 @@ def test_count_with_none_vector_count(pinecone_db):
|
|||||||
pinecone_db.index.describe_index_stats.assert_called_once()
|
pinecone_db.index.describe_index_stats.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_count_with_empty_string_namespace(pinecone_db):
|
||||||
|
pinecone_db.namespace = ""
|
||||||
|
stats_mock = MagicMock()
|
||||||
|
stats_mock.namespaces = {"": MagicMock(vector_count=3)}
|
||||||
|
stats_mock.total_vector_count = 99
|
||||||
|
pinecone_db.index.describe_index_stats.return_value = stats_mock
|
||||||
|
|
||||||
|
count = pinecone_db.count()
|
||||||
|
assert count == 3
|
||||||
|
|
||||||
|
|
||||||
def test_list_error_returns_list_not_dict(pinecone_db):
|
def test_list_error_returns_list_not_dict(pinecone_db):
|
||||||
"""list() error path must return [[]] so callers can do result[0]."""
|
"""list() error path must return [[]] so callers can do result[0]."""
|
||||||
pinecone_db.index.query.side_effect = Exception("connection error")
|
pinecone_db.index.query.side_effect = Exception("connection error")
|
||||||
|
|||||||
Reference in New Issue
Block a user