From 9c3d0733878a076223b10bdab3d454c0e880df49 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 07:16:42 -0500 Subject: [PATCH 1/6] prep 2.7.0 release branch --- pyproject.toml | 2 +- uv.lock | 54 +++++++++++++++++++++++++------------------------- 2 files changed, 28 insertions(+), 28 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 58c436d..f38c258 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "simplevecdb" -version = "2.6.2" +version = "2.7.0" description = "Dead-simple local vector database powered by usearch HNSW." authors = [{ name = "Dayton Dunbar", email = "coderdayton14@gmail.com" }] license = { text = "MIT" } diff --git a/uv.lock b/uv.lock index 68d7b76..025f975 100755 --- a/uv.lock +++ b/uv.lock @@ -894,7 +894,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ @@ -1416,17 +1416,17 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version < '3.11'" }, - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, - { name = "jedi", marker = "python_full_version < '3.11'" }, - { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, - { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, - { name = "pygments", marker = "python_full_version < '3.11'" }, - { name = "stack-data", marker = "python_full_version < '3.11'" }, - { name = "traitlets", marker = "python_full_version < '3.11'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "exceptiongroup" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/85/31/10ac88f3357fc276dc8a64e8880c82e80e7459326ae1d0a211b40abf6665/ipython-8.37.0.tar.gz", hash = "sha256:ca815841e1a41a1e6b73a0b08f3038af9b2252564d01fc405356d34033012216", size = 5606088, upload-time = "2025-05-31T16:39:09.613Z" } wheels = [ @@ -1444,17 +1444,17 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "colorama", marker = "python_full_version >= '3.11' and sys_platform == 'win32'" }, - { name = "decorator", marker = "python_full_version >= '3.11'" }, - { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.11'" }, - { name = "jedi", marker = "python_full_version >= '3.11'" }, - { name = "matplotlib-inline", marker = "python_full_version >= '3.11'" }, - { name = "pexpect", marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit", marker = "python_full_version >= '3.11'" }, - { name = "pygments", marker = "python_full_version >= '3.11'" }, - { name = "stack-data", marker = "python_full_version >= '3.11'" }, - { name = "traitlets", marker = "python_full_version >= '3.11'" }, - { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, + { name = "typing-extensions", marker = "python_full_version < '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/29/e6/48c74d54039241a456add616464ea28c6ebf782e4110d419411b83dae06f/ipython-9.7.0.tar.gz", hash = "sha256:5f6de88c905a566c6a9d6c400a8fed54a638e1f7543d17aae2551133216b1e4e", size = 4422115, upload-time = "2025-11-05T12:18:54.646Z" } wheels = [ @@ -1466,7 +1466,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments", marker = "python_full_version >= '3.11'" }, + { name = "pygments" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -4537,7 +4537,7 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -4599,7 +4599,7 @@ resolution-markers = [ "python_full_version == '3.11.*'", ] dependencies = [ - { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" } }, ] sdist = { url = "https://files.pythonhosted.org/packages/0a/ca/d8ace4f98322d01abcd52d381134344bf7b431eba7ed8b42bdea5a3c2ac9/scipy-1.16.3.tar.gz", hash = "sha256:01e87659402762f43bd2fee13370553a17ada367d42e7487800bf2916535aecb", size = 30597883, upload-time = "2025-10-28T17:38:54.068Z" } wheels = [ @@ -4705,7 +4705,7 @@ wheels = [ [[package]] name = "simplevecdb" -version = "2.6.2" +version = "2.7.0" source = { editable = "." } dependencies = [ { name = "cryptography" }, From 8de1374115deea850c1aee0143e56dc306ff4b39 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 07:17:50 -0500 Subject: [PATCH 2/6] transactional vector writes, reserved ids, explicit-id collision policy --- CHANGELOG.md | 54 +++++ README.md | 3 +- docs/CHANGELOG.md | 54 +++++ docs/Features.md | 14 +- docs/api/async.md | 20 ++ docs/api/core.md | 1 + docs/examples.md | 45 +++- src/simplevecdb/async_core.py | 57 ++++- src/simplevecdb/constants.py | 5 + src/simplevecdb/core.py | 189 ++++++++++++++-- src/simplevecdb/engine/catalog.py | 155 ++++++++++++-- src/simplevecdb/integrations/langchain.py | 4 + src/simplevecdb/types.py | 9 +- src/simplevecdb/utils.py | 16 ++ tests/integration/test_rag.py | 28 ++- tests/unit/test_core.py | 6 +- tests/unit/test_id_reservation_and_tx.py | 250 ++++++++++++++++++++++ 17 files changed, 858 insertions(+), 52 deletions(-) create mode 100644 tests/unit/test_id_reservation_and_tx.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f5d49e9..52acf64 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,60 @@ All notable changes to SimpleVecDB will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.7.0] - 2026-07-29 + +### Transactions cover the vector index + +`tx()` and `db.transaction()` guarded only half the store: the SAVEPOINT +rolled back the catalog rows while the vectors written to the HNSW index +stayed put, leaving the index keyed to rows that no longer existed. Vector +writes now take part in the transaction, and explicit document ids stop +overwriting existing documents by accident. + +#### Breaking + +- **`add_texts(ids=…)` refuses an id that already exists.** It previously + upserted silently, so a stale or guessed id destroyed the stored document + with no error. Pass `on_conflict="replace"` for the old behaviour. The + check runs before anything is written, so a rejected call leaves the + collection untouched — including across internal batches. Repeating an id + within a single call is always an error. +- The LangChain adapter keeps upserting: `SimpleVecDBVectorStore.add_texts` + defaults to `on_conflict="replace"` so LangChain's own contract holds. + +#### Added + +- **`collection.reserve_ids(n)`** — reserve ids without writing rows. Ids + come out of the auto-increment sequence and can never be handed out again, + so a document's own id (a self-reference, a shared group key) can be baked + into its metadata and the whole group written with one `add_texts` call + instead of an insert followed by a patch-up write. +- **`AsyncVectorCollection.atomic(fn)`** — run a transaction from async code. + `fn` is a synchronous callable receiving the sync collection. This is a + callback rather than `async with` because a transaction holds a + `threading.RLock` for its lifetime: entering and exiting in two separate + executor tasks can release the lock from a thread that never acquired it + (`RuntimeError: cannot release un-acquired lock`), and holding it across + awaits starves the pool. Running the whole body in one executor task keeps + acquire and release paired. +- **`AsyncVectorCollection.reserve_ids`** and `on_conflict` on the async + `add_texts`. + +#### Fixed + +- **Vector writes are transactional.** `add_texts`, `delete_by_ids`, + `pending.flush()` and `ttl.sweep()` buffer their HNSW mutations while a + transaction is open and apply them just before the outermost savepoint + releases, so a rollback undoes rows and vectors together. Applying ahead of + the release means a failing vector write can still roll the catalog back; + usearch has no undo, so a failure partway through the buffer is logged + loudly with a `rebuild_index()` recommendation rather than hidden. + +#### Changed + +- A search *inside* a transaction does not see that transaction's own vector + writes — they land at commit. Catalog reads are unaffected. + ## [2.6.2] - 2026-06-06 ### Correctness and contract fixes diff --git a/README.md b/README.md index b7e85fc..d3a7511 100755 --- a/README.md +++ b/README.md @@ -216,7 +216,8 @@ A few of the things SimpleVecDB does well — see delete/callback sweep, and a threshold-driven rebuild scheduler. - **Atomic counters & transactions (v2.6.1)** — `increment_metadata` for JSON deltas in one statement; SAVEPOINT-backed `db.transaction()` / - `collection.tx()` rolling all catalog writes back on error. + `collection.tx()` rolling catalog *and* vector writes back on error, + with `reserve_ids()` for ids you need before the rows exist. - **Async, encryption, clustering, hierarchies** — full async surface (with executor injection), SQLCipher AES-256, K-means / MiniBatch K-means / HDBSCAN, parent/child relationships. diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index f5d49e9..52acf64 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -5,6 +5,60 @@ All notable changes to SimpleVecDB will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.7.0] - 2026-07-29 + +### Transactions cover the vector index + +`tx()` and `db.transaction()` guarded only half the store: the SAVEPOINT +rolled back the catalog rows while the vectors written to the HNSW index +stayed put, leaving the index keyed to rows that no longer existed. Vector +writes now take part in the transaction, and explicit document ids stop +overwriting existing documents by accident. + +#### Breaking + +- **`add_texts(ids=…)` refuses an id that already exists.** It previously + upserted silently, so a stale or guessed id destroyed the stored document + with no error. Pass `on_conflict="replace"` for the old behaviour. The + check runs before anything is written, so a rejected call leaves the + collection untouched — including across internal batches. Repeating an id + within a single call is always an error. +- The LangChain adapter keeps upserting: `SimpleVecDBVectorStore.add_texts` + defaults to `on_conflict="replace"` so LangChain's own contract holds. + +#### Added + +- **`collection.reserve_ids(n)`** — reserve ids without writing rows. Ids + come out of the auto-increment sequence and can never be handed out again, + so a document's own id (a self-reference, a shared group key) can be baked + into its metadata and the whole group written with one `add_texts` call + instead of an insert followed by a patch-up write. +- **`AsyncVectorCollection.atomic(fn)`** — run a transaction from async code. + `fn` is a synchronous callable receiving the sync collection. This is a + callback rather than `async with` because a transaction holds a + `threading.RLock` for its lifetime: entering and exiting in two separate + executor tasks can release the lock from a thread that never acquired it + (`RuntimeError: cannot release un-acquired lock`), and holding it across + awaits starves the pool. Running the whole body in one executor task keeps + acquire and release paired. +- **`AsyncVectorCollection.reserve_ids`** and `on_conflict` on the async + `add_texts`. + +#### Fixed + +- **Vector writes are transactional.** `add_texts`, `delete_by_ids`, + `pending.flush()` and `ttl.sweep()` buffer their HNSW mutations while a + transaction is open and apply them just before the outermost savepoint + releases, so a rollback undoes rows and vectors together. Applying ahead of + the release means a failing vector write can still roll the catalog back; + usearch has no undo, so a failure partway through the buffer is logged + loudly with a `rebuild_index()` recommendation rather than hidden. + +#### Changed + +- A search *inside* a transaction does not see that transaction's own vector + writes — they land at commit. Catalog reads are unaffected. + ## [2.6.2] - 2026-06-06 ### Correctness and contract fixes diff --git a/docs/Features.md b/docs/Features.md index 1062bfe..d2159bc 100644 --- a/docs/Features.md +++ b/docs/Features.md @@ -80,10 +80,16 @@ release-by-release detail, see the [Changelog](CHANGELOG.md). - **Transactions (v2.6.1)** — `with db.transaction() as tx: …` and `with collection.tx(): …` wrap a SAVEPOINT around catalog writes (metadata, counters, edges, events, TTL, and `update_embedding`'s - pending overlay). A raised exception rolls all SQL writes back. Coarse - vector mutations (`add_texts`, `delete_by_ids`) are NOT rolled back — - use `update_embedding` + `pending.flush()` for vector changes that - must be commit-gated. + pending overlay). A raised exception rolls all SQL writes back. + Vector mutations (`add_texts`, `delete_by_ids`, `pending.flush()`, + `ttl.sweep()`) are buffered and applied on commit, so a rollback + leaves neither store changed; the tradeoff is that a search inside + the transaction cannot see its own vector writes. Async callers use + `await collection.atomic(fn)` instead of `async with`. +- **Reserved ids** — `collection.reserve_ids(n)` hands out ids before + the rows exist, so self-referential or grouped documents go in with a + single `add_texts` call. Reusing an existing id now raises instead of + silently overwriting; pass `on_conflict="replace"` for an upsert. ## Relationships diff --git a/docs/api/async.md b/docs/api/async.md index 6521540..0c05c0c 100755 --- a/docs/api/async.md +++ b/docs/api/async.md @@ -44,6 +44,7 @@ db = AsyncVectorDB("vectors.db", max_workers=8) | Sync Method | Async Method | | --------------------------------- | -------------------------------------------------- | | `add_texts()` | `await collection.add_texts()` | +| `reserve_ids()` | `await collection.reserve_ids()` | | `similarity_search()` | `await collection.similarity_search()` | | `similarity_search_batch()` | `await collection.similarity_search_batch()` | | `keyword_search()` | `await collection.keyword_search()` | @@ -56,6 +57,25 @@ Synchronous properties remain unchanged: - `collection.name` - Collection name +## Transactions + +`collection.tx()` has no `async with` equivalent — the transaction holds a +`threading.RLock` for its lifetime, and entering and exiting in two separate +executor tasks can release that lock from a thread that never acquired it. +Pass a synchronous callback to `atomic()` instead, and the whole transaction +runs in one executor thread: + +```python +def move(coll): + coll.delete_by_ids([1]) + return coll.add_texts(["replacement"], embeddings=[[0.1] * 384]) + +new_ids = await collection.atomic(move) +``` + +The callback receives the underlying sync `VectorCollection` and must not +await. Catalog writes and vector writes commit or roll back together. + ## Concurrent Operations Run multiple searches in parallel with `asyncio.gather` or use batch search for better performance: diff --git a/docs/api/core.md b/docs/api/core.md index 8031891..49947f2 100755 --- a/docs/api/core.md +++ b/docs/api/core.md @@ -22,6 +22,7 @@ A named collection of vectors within a database. members: - add_texts - add_texts_streaming + - reserve_ids - similarity_search - similarity_search_batch - keyword_search diff --git a/docs/examples.md b/docs/examples.md index 5ab44ef..7eeb6b5 100755 --- a/docs/examples.md +++ b/docs/examples.md @@ -265,10 +265,7 @@ collection.events.prune(before_seq=seq - 100_000) `db.transaction()` and `collection.tx()` wrap a SAVEPOINT around catalog writes (metadata, counters, edges, events, TTL, and the -pending overlay). A raised exception rolls all SQL writes back. Coarse -vector mutations (`add_texts`, `delete_by_ids`) are *not* rolled back — -use `update_embedding` + `pending.flush()` for vector changes that -must be commit-gated. +pending overlay). A raised exception rolls all SQL writes back. ```python with db.transaction() as tx: @@ -277,6 +274,46 @@ with db.transaction() as tx: # any exception below rolls both writes back ``` +Vector mutations (`add_texts`, `delete_by_ids`, `pending.flush()`, +`ttl.sweep()`) take part too: a SAVEPOINT cannot roll back the HNSW +index, so the vector writes are buffered and applied when the outermost +transaction commits. The tradeoff is that a search *inside* the +transaction does not see the transaction's own vector writes. + +```python +with collection.tx() as coll: + coll.delete_by_ids([1]) + coll.add_texts(["replacement"], embeddings=[vec]) + # rows and vectors both survive, or neither does +``` + +On the async side there is no `async with`; pass a sync callback to +`atomic()` instead — see the async API reference for why. + +```python +await collection.atomic(lambda coll: coll.add_texts(["x"], embeddings=[vec])) +``` + +### Reserving ids + +`reserve_ids(n)` burns ids out of the auto-increment sequence without +writing rows, so a document's own id can be baked into its metadata and +the whole group written in one call. + +```python +ids = collection.reserve_ids(3) +group = {"episode_group_id": ids[0]} +collection.add_texts( + texts, + metadatas=[{**group, "id": i} for i in ids], + embeddings=vectors, + ids=ids, +) +``` + +Passing an `ids=` value that already exists raises. Pass +`on_conflict="replace"` for the old overwrite-in-place behaviour. + ## Benchmark scripts ```bash diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index e274bb2..0bf934e 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -27,14 +27,16 @@ import asyncio import functools from concurrent.futures import ThreadPoolExecutor -from collections.abc import Sequence +from collections.abc import Callable, Sequence from threading import Lock -from typing import Any +from typing import Any, TypeVar import logging from .core import VectorDB, VectorCollection -from .types import Document, DistanceStrategy, Quantization +from .types import Document, DistanceStrategy, OnConflict, Quantization + +T = TypeVar("T") _logger = logging.getLogger(__name__) @@ -69,6 +71,53 @@ async def _run(self, fn, /, *args, **kwargs): self._executor, functools.partial(fn, *args, **kwargs) ) + async def atomic(self, fn: Callable[[VectorCollection], T]) -> T: + """Run `fn` inside a transaction on this collection. + + `fn` is an ordinary synchronous callable and receives the underlying + `VectorCollection`; everything it does — catalog writes and vector + writes alike — commits or rolls back as one unit. + + async def swap(coll): + coll.delete_by_ids([1]) + coll.add_texts(["replacement"], embeddings=[vec]) + + await collection.atomic(swap) + + This is a callback rather than `async with collection.tx()` on + purpose. The transaction holds a `threading.RLock` for its whole + lifetime, and an `async with` would enter and exit in two separate + executor tasks: the pool is free to run them on different threads, + and releasing an RLock from a thread that did not acquire it raises + `RuntimeError: cannot release un-acquired lock`. Landing on the same + thread by luck would be no better — every other call on this + collection would sit on that lock across each `await` in the body, + which starves a pool this small. Running the whole body in one + executor task keeps acquire and release paired on one thread. + + Because `fn` runs off the event loop, it must not await; use the + sync collection API inside it. + + Args: + fn: Callable invoked with the sync collection. + + Returns: + Whatever `fn` returns. + """ + return await self._run(self._in_tx, fn) + + def _in_tx(self, fn: Callable[[VectorCollection], T]) -> T: + """Body of `atomic`, run wholly inside one executor thread.""" + with self._collection.tx() as coll: + return fn(coll) + + async def reserve_ids(self, count: int) -> list[int]: + """Reserve document ids without writing rows. + + See VectorCollection.reserve_ids for full documentation. + """ + return await self._run(self._collection.reserve_ids, count) + async def add_texts( self, texts: Sequence[str], @@ -78,6 +127,7 @@ async def add_texts( *, parent_ids: Sequence[int | None] | None = None, threads: int = 0, + on_conflict: OnConflict = "error", ) -> list[int]: """Add texts with optional embeddings and metadata. @@ -91,6 +141,7 @@ async def add_texts( ids, parent_ids=parent_ids, threads=threads, + on_conflict=on_conflict, ) async def similarity_search( diff --git a/src/simplevecdb/constants.py b/src/simplevecdb/constants.py index ae2c75a..833788b 100755 --- a/src/simplevecdb/constants.py +++ b/src/simplevecdb/constants.py @@ -123,6 +123,11 @@ # Default cadence for opt-in TTL background sweeper threads. TTL_SWEEP_DEFAULT_INTERVAL_S = 60.0 +# Bound-parameter budget for a single SQLite statement. Well under the +# SQLITE_MAX_VARIABLE_NUMBER floor (999 on builds before 3.32), so an +# `IN (...)` list built from caller-supplied ids is chunked, never rejected. +SQLITE_MAX_BOUND_PARAMS = 900 + # Heuristic thresholds that trigger maintenance.rebuild_if_needed. REBUILD_PENDING_THRESHOLD = 5_000 REBUILD_TOMBSTONE_THRESHOLD = 5_000 diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 2198179..788884a 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -26,6 +26,7 @@ from .types import ( Document, DistanceStrategy, + OnConflict, Quantization, Edge, Event, @@ -34,7 +35,7 @@ ClusterResult, ClusterTagCallback, ) -from .utils import _import_optional +from .utils import _import_optional, find_duplicates from .engine.quantization import QuantizationStrategy from .engine.search import SearchEngine from .engine.catalog import CatalogManager, _TxState @@ -296,6 +297,76 @@ def _resolve_index_path(self) -> str | None: return self._index_path + def reserve_ids(self, count: int) -> list[int]: + """Reserve `count` document ids without writing any rows. + + The ids are burned out of the auto-increment sequence, so no later + insert can be handed one of them. Use this when a document's own id + has to appear in its metadata — a self-reference, or a shared group + key across a batch — so the whole group goes in with a single + `add_texts` call instead of an insert followed by a patch-up write. + + ids = collection.reserve_ids(3) + group = {"group_id": ids[0]} + collection.add_texts(texts, [{**group, "id": i} for i in ids], ids=ids) + + Reserving is not inserting: ids that are never passed back as `ids=` + stay unused, leaving a gap in the sequence. + + Args: + count: How many ids to reserve. Must be positive. + + Returns: + `count` consecutive ids in ascending order. + + Raises: + ValueError: If `count` is not positive. + """ + return self._catalog.reserve_ids(count) + + def _index_add( + self, + keys: "np.ndarray", + vectors: "np.ndarray", + *, + threads: int = 0, + ) -> None: + """Add vectors to the HNSW index, deferring the write inside a tx. + + A SQLite SAVEPOINT cannot roll back usearch, so when a transaction + is open the write is buffered on the shared transaction state and + applied just before the outermost savepoint releases. A rollback + discards the buffer, so the catalog and the index stay in step. + + The buffered arrays are copied: they may alias caller-owned memory + that is free to change before the transaction commits. + """ + if self._tx_state.depth > 0: + keys_buf = np.array(keys, dtype=np.uint64, copy=True) + vecs_buf = np.array(vectors, dtype=np.float32, copy=True) + + def _deferred_add() -> None: + self._index.add(keys_buf, vecs_buf, threads=threads) + + self._tx_state.index_ops.append(_deferred_add) + return + self._index.add(keys, vectors, threads=threads) + + def _index_remove(self, keys: list[int]) -> None: + """Remove keys from the HNSW index, deferring the write inside a tx. + + Counterpart to `_index_add`; see there for why the write is held. + """ + if self._tx_state.depth > 0: + keys_buf = [int(k) for k in keys] + + def _deferred_remove() -> None: + self._index.remove(keys_buf) + + self._tx_state.index_ops.append(_deferred_remove) + return + self._index.remove(keys) + def add_texts( self, texts: Sequence[str], @@ -305,29 +376,36 @@ def add_texts( *, parent_ids: Sequence[int | None] | None = None, threads: int = 0, + on_conflict: OnConflict = "error", ) -> list[int]: """ Add texts with optional embeddings and metadata to the collection. - Automatically infers vector dimension from first batch. Supports upsert - (update on conflict) when providing existing IDs. For COSINE distance, - vectors are L2-normalized automatically by usearch. + Automatically infers vector dimension from first batch. For COSINE + distance, vectors are L2-normalized automatically by usearch. + + Supplying `ids` that already exist is an error by default; pass + `on_conflict="replace"` to get upsert behaviour. Args: texts: Document text content to store. metadatas: Optional metadata dicts (one per text). embeddings: Optional pre-computed embeddings (one per text). If None, attempts to use local embedding model. - ids: Optional document IDs for upsert behavior. + ids: Optional explicit document IDs — see `reserve_ids`. parent_ids: Optional parent document IDs for hierarchical relationships. threads: Number of threads for parallel insertion (0=auto). + on_conflict: What an explicit id that already exists does — + ``"error"`` (default) raises before anything is written, + ``"replace"`` overwrites the existing document. Returns: List of inserted/updated document IDs. Raises: - ValueError: If embedding dimensions don't match, or if no embeddings - provided and local embedder not available. + ValueError: If embedding dimensions don't match, if no embeddings + provided and local embedder not available, or if `ids` collide + with existing documents under ``on_conflict="error"``. """ if not texts: return [] @@ -349,6 +427,28 @@ def add_texts( f"parent_ids length ({len(parent_ids)}) must match texts length ({len(texts)})" ) + # Check explicit ids once, for the whole call. The catalog re-checks + # per batch, but by then an earlier batch may already have committed — + # only an up-front check can promise that a rejected call wrote + # nothing. Runs before embedding resolution so a doomed call does not + # pay for embeddings first. + explicit_ids = [i for i in ids if i is not None] if ids is not None else [] + if explicit_ids: + duplicates = find_duplicates(explicit_ids) + if duplicates: + raise ValueError( + f"ids contains duplicate values within a single call: {duplicates}" + ) + if on_conflict == "error": + colliding = self._catalog.existing_ids(explicit_ids) + if colliding: + raise ValueError( + f"add_texts: {len(colliding)} id(s) already exist: " + f"{colliding[:10]}{'...' if len(colliding) > 10 else ''}. " + 'Pass on_conflict="replace" to overwrite them, or use ' + "reserve_ids() to get ids that cannot collide." + ) + # Resolve embeddings if embeddings is None: try: @@ -396,14 +496,17 @@ def add_texts( batch_ids, embeddings=batch_embeds if self._store_embeddings else None, parent_ids=batch_parent_ids, + on_conflict=on_conflict, ) - # Add to usearch index. The catalog rows above are already - # committed, so if this fails the two stores diverge (rows present, - # vectors missing). Log it so the divergence is visible instead of - # silent; recovery is rebuild_index() (needs store_embeddings=True). + # Add to usearch index. Inside a transaction this is buffered and + # applied with the commit, so a rollback drops both stores + # together. Outside one the catalog rows are already committed, so + # a failure here diverges the two (rows present, vectors missing); + # log it rather than fail silently. Recovery is rebuild_index() + # (needs store_embeddings=True). try: - self._index.add( + self._index_add( np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads ) except Exception: @@ -594,7 +697,7 @@ def _process_streaming_batch( doc_ids = self._catalog.add_documents( texts, metas, None, embeddings=embeds_resolved ) - self._index.add(np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads) + self._index_add(np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads) return doc_ids @@ -775,8 +878,9 @@ def delete_by_ids(self, ids: Iterable[int]) -> None: self._catalog.delete_by_ids(ids_list) # Then remove from usearch (if this fails, catalog is clean and - # rebuild_index() can recover the index from stored data) - self._index.remove(ids_list) + # rebuild_index() can recover the index from stored data). Inside a + # transaction the removal is buffered until commit. + self._index_remove(ids_list) def remove_texts( self, @@ -1813,7 +1917,7 @@ def flush(self, *, max_batch: int | None = None) -> int: .copy() ) # add() takes the write lock and does remove+add per existing key. - idx.add(ids, mat) + self._collection._index_add(ids, mat) cat.delete_pending_vectors([int(i) for i in ids]) # Track flush count for the rebuild scheduler. try: @@ -1897,7 +2001,7 @@ def sweep( ) if deleted: try: - self._collection._index.remove(deleted) + self._collection._index_remove(deleted) except Exception: # Catalog rows are already gone; index drift means # subsequent searches may surface phantom hits until a @@ -2143,16 +2247,24 @@ class _DBTransaction: owns skips its per-call commit. On success the SAVEPOINT is released; on exception it's rolled back and the depth is reset. + Vector writes are held too: collections buffer their HNSW mutations + on the shared transaction state (see `VectorCollection._index_add`) + and the outermost transaction applies them just before releasing, so + a rollback leaves neither store changed. The cost of deferring is + that a search inside the transaction cannot see the transaction's + own vector writes. + Yields a mapping-like object so callers can do `tx["collection_name"]` to operate on individual collections. """ - __slots__ = ("_db", "_savepoint_name", "_entered") + __slots__ = ("_db", "_savepoint_name", "_entered", "_index_mark") def __init__(self, db: "VectorDB") -> None: self._db = db self._savepoint_name: str | None = None self._entered = False + self._index_mark = 0 def __enter__(self) -> "_DBTransaction": self._db._lock.acquire() @@ -2163,19 +2275,60 @@ def __enter__(self) -> "_DBTransaction": self._db._tx_state.depth = depth + 1 self._savepoint_name = name self._entered = True + # Everything buffered from here on belongs to this savepoint and + # is dropped if it rolls back. + self._index_mark = len(self._db._tx_state.index_ops) except Exception: self._db._lock.release() raise return self + def _apply_index_ops(self, savepoint: str) -> None: + """Flush buffered HNSW writes before the outermost savepoint releases. + + Running ahead of the RELEASE is what makes the two stores atomic in + the normal case: if a vector write raises, the catalog side can still + be rolled back. usearch has no undo, so a failure partway through the + buffer can still leave vectors applied for rows that are about to + disappear — that case is reported loudly rather than papered over. + """ + ops = self._db._tx_state.index_ops + if not ops: + return + applied = 0 + try: + for op in ops: + op() + applied += 1 + except Exception: + self._db.conn.execute(f"ROLLBACK TO SAVEPOINT {savepoint}") + self._db.conn.execute(f"RELEASE SAVEPOINT {savepoint}") + _logger.error( + "transaction: %d of %d buffered index operations applied before " + "failure; catalog rolled back but the index may retain vectors " + "for rows that no longer exist — run rebuild_index() to resync.", + applied, + len(ops), + exc_info=True, + ) + raise + finally: + ops.clear() + def __exit__(self, exc_type, exc, tb) -> None: try: name = self._savepoint_name assert name is not None + state = self._db._tx_state try: if exc_type is None: + if state.depth == 1: + self._apply_index_ops(name) self._db.conn.execute(f"RELEASE SAVEPOINT {name}") else: + # Drop vector writes buffered inside this savepoint; an + # outer transaction keeps everything it buffered earlier. + del state.index_ops[self._index_mark :] self._db.conn.execute(f"ROLLBACK TO SAVEPOINT {name}") self._db.conn.execute(f"RELEASE SAVEPOINT {name}") finally: diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 8dcc9e3..0933e21 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -15,9 +15,11 @@ from typing import Any, TYPE_CHECKING, Callable from collections.abc import Iterable, Sequence +from ..constants import SQLITE_MAX_BOUND_PARAMS +from ..types import ON_CONFLICT_POLICIES, OnConflict from ..utils import _batched -from ..utils import validate_filter, retry_on_lock, normalize_filter +from ..utils import validate_filter, retry_on_lock, normalize_filter, find_duplicates if TYPE_CHECKING: import sqlite3 @@ -79,17 +81,24 @@ def _check_finite_edge_field(value: Any, field: str) -> None: class _TxState: - """Shared per-VectorDB transaction depth counter (gap 2). + """Shared per-VectorDB transaction state (gap 2). Used by VectorDB.transaction() to mark all collections/catalogs as operating inside an outer SAVEPOINT. Catalog write helpers consult - this to decide whether to commit on exit. + `depth` to decide whether to commit on exit. + + `index_ops` buffers HNSW mutations deferred by the collections taking + part in the transaction. SQLite SAVEPOINTs cannot roll back the + usearch index, so vector writes are held here and applied only when + the outermost transaction is about to release; a rollback truncates + the buffer instead, leaving the two stores in step. """ - __slots__ = ("depth",) + __slots__ = ("depth", "index_ops") def __init__(self) -> None: self.depth: int = 0 + self.index_ops: list[Callable[[], None]] = [] class _CatalogWritable: @@ -548,20 +557,34 @@ def add_documents( ids: Sequence[int | None] | None = None, embeddings: Sequence[Sequence[float]] | None = None, parent_ids: Sequence[int | None] | None = None, + on_conflict: OnConflict = "error", ) -> list[int]: """ - Insert or update document metadata. + Insert document metadata, optionally replacing rows on id collision. Args: texts: Document text content metadatas: Metadata dicts for each document - ids: Optional document IDs for upsert behavior + ids: Optional explicit document IDs embeddings: Optional embedding vectors to store parent_ids: Optional parent document IDs for hierarchical relationships + on_conflict: What an explicit id that already exists does — + ``"error"`` raises and writes nothing, ``"replace"`` overwrites + the stored row. Returns: List of document IDs (rowids) + + Raises: + ValueError: If `on_conflict` is not a recognised policy, if `ids` + repeats an id within the call, or if `on_conflict="error"` + and an id already exists. """ + if on_conflict not in ON_CONFLICT_POLICIES: + raise ValueError( + f"on_conflict must be one of {sorted(ON_CONFLICT_POLICIES)}, " + f"got {on_conflict!r}" + ) if not texts: return [] @@ -617,20 +640,50 @@ def add_documents( real_ids: list[int] = [-1] * len(ids_list) + if explicit_rows: + # A repeat inside one call is always an error: the two rows would + # race for the same id whatever the policy, and under "replace" + # the loser would be silently discarded. + explicit_ids = [int(r[0]) for r in explicit_rows] + duplicates = find_duplicates(explicit_ids) + if duplicates: + raise ValueError( + f"ids contains duplicate values within a single call: {duplicates}" + ) + with self._writable(): if explicit_rows: - self.conn.executemany( - f""" - INSERT INTO {self._table_name}(id, text, metadata, embedding, parent_id) - VALUES (?, ?, ?, ?, ?) - ON CONFLICT(id) DO UPDATE SET - text=excluded.text, - metadata=excluded.metadata, - embedding=excluded.embedding, - parent_id=excluded.parent_id - """, - explicit_rows, - ) + if on_conflict == "error": + # Checked inside _writable(), so no concurrent writer can + # land a colliding row between the check and the INSERT. + colliding = self.existing_ids(explicit_ids) + if colliding: + raise ValueError( + f"add_documents: {len(colliding)} id(s) already exist: " + f"{colliding[:10]}" + f"{'...' if len(colliding) > 10 else ''}. " + 'Pass on_conflict="replace" to overwrite them.' + ) + self.conn.executemany( + f""" + INSERT INTO {self._table_name}(id, text, metadata, embedding, parent_id) + VALUES (?, ?, ?, ?, ?) + """, + explicit_rows, + ) + else: + self.conn.executemany( + f""" + INSERT INTO {self._table_name}(id, text, metadata, embedding, parent_id) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + text=excluded.text, + metadata=excluded.metadata, + embedding=excluded.embedding, + parent_id=excluded.parent_id + """, + explicit_rows, + ) if auto_rows: # Use a single multi-VALUES INSERT ... RETURNING id so we @@ -675,6 +728,72 @@ def add_documents( _logger.debug("Added %d documents, ids=%s", len(real_ids), real_ids[:5]) return real_ids + def existing_ids(self, ids: Sequence[int]) -> list[int]: + """Return the subset of `ids` already present, in ascending order. + + Chunked because SQLite caps the number of bound parameters per + statement (`SQLITE_MAX_VARIABLE_NUMBER`). + """ + found: list[int] = [] + for start in range(0, len(ids), SQLITE_MAX_BOUND_PARAMS): + chunk = ids[start : start + SQLITE_MAX_BOUND_PARAMS] + placeholders = ",".join("?" for _ in chunk) + rows = self.conn.execute( + f"SELECT id FROM {self._table_name} WHERE id IN ({placeholders})", + tuple(chunk), + ).fetchall() + found.extend(int(r[0]) for r in rows) + return sorted(found) + + def reserve_ids(self, count: int) -> list[int]: + """Reserve `count` document ids without inserting any rows. + + Advances the table's AUTOINCREMENT high-water mark so the returned + ids can never be handed out again by a later auto-id insert. This + lets a caller stamp ids into metadata (self-referential rows, a + shared group key) and then write the whole group in one + `add_texts` call, instead of inserting and patching afterwards. + + Reserved ids are not rows: nothing is stored until they are passed + back as `ids=`. Ids that are never used simply stay unallocated. + + Args: + count: How many ids to reserve. Must be positive. + + Returns: + `count` consecutive ids in ascending order. + + Raises: + ValueError: If `count` is not positive. + """ + if count <= 0: + raise ValueError(f"reserve_ids: count must be positive, got {count}") + + with self._writable(): + seq_row = self.conn.execute( + "SELECT seq FROM sqlite_sequence WHERE name = ?", + (self._table_name,), + ).fetchone() + max_row = self.conn.execute( + f"SELECT COALESCE(MAX(id), 0) FROM {self._table_name}" + ).fetchone() + # MAX(id) guards the case where explicit ids were inserted past + # the sequence mark; the reservation must clear both. + base = max(int(seq_row[0]) if seq_row is not None else 0, int(max_row[0])) + new_seq = base + count + if seq_row is None: + self.conn.execute( + "INSERT INTO sqlite_sequence(name, seq) VALUES (?, ?)", + (self._table_name, new_seq), + ) + else: + self.conn.execute( + "UPDATE sqlite_sequence SET seq = ? WHERE name = ?", + (new_seq, self._table_name), + ) + + return list(range(base + 1, new_seq + 1)) + @retry_on_lock(max_retries=5, base_delay=0.1) def delete_by_ids(self, ids: Iterable[int]) -> list[int]: """ diff --git a/src/simplevecdb/integrations/langchain.py b/src/simplevecdb/integrations/langchain.py index 2ce18e6..a5708bf 100755 --- a/src/simplevecdb/integrations/langchain.py +++ b/src/simplevecdb/integrations/langchain.py @@ -95,11 +95,15 @@ def add_texts( embeddings = None if self.embedding: embeddings = self.embedding.embed_documents(texts_list) + # LangChain callers expect add_texts(ids=…) to overwrite an existing + # document, so this adapter keeps upsert semantics even though the + # native API now refuses a colliding id by default. ids = self._collection.add_texts( texts=texts_list, metadatas=metadatas, embeddings=embeddings, ids=kwargs.get("ids"), + on_conflict=kwargs.get("on_conflict", "replace"), ) return [str(id_) for id_ in ids] diff --git a/src/simplevecdb/types.py b/src/simplevecdb/types.py index 7b3947d..76ce8c2 100755 --- a/src/simplevecdb/types.py +++ b/src/simplevecdb/types.py @@ -3,7 +3,7 @@ import dataclasses from dataclasses import field from enum import Enum -from typing import Callable, TypedDict, TYPE_CHECKING +from typing import Callable, Literal, TypedDict, TYPE_CHECKING if TYPE_CHECKING: import numpy as np @@ -38,6 +38,13 @@ class StreamingProgress(TypedDict): ProgressCallback = Callable[[StreamingProgress], None] +# What an explicitly supplied document id that already exists should do. +# "error" refuses the whole write; "replace" overwrites the stored row. +OnConflict = Literal["error", "replace"] + +ON_CONFLICT_POLICIES: frozenset[str] = frozenset({"error", "replace"}) + + class DistanceStrategy(StrEnum): """Supported distance metrics for usearch backend.""" diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index 38d709a..14f22ec 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -34,6 +34,22 @@ def _batched(iterable: Iterable[Any], n: int) -> Iterable[Sequence[Any]]: yield batch +def find_duplicates(values: Sequence[int]) -> list[int]: + """Return the values appearing more than once, in ascending order. + + Single pass over the input — callers hand this whole batches of + caller-supplied document ids, so a quadratic scan is not an option. + """ + seen: set[int] = set() + repeated: set[int] = set() + for value in values: + if value in seen: + repeated.add(value) + else: + seen.add(value) + return sorted(repeated) + + def _import_optional(name: str) -> Any: """Attempt to import a module while honoring tests that stub sys.modules.""" sentinel = object() diff --git a/tests/integration/test_rag.py b/tests/integration/test_rag.py index 44819cd..3a6f463 100755 --- a/tests/integration/test_rag.py +++ b/tests/integration/test_rag.py @@ -4,6 +4,9 @@ import pytest from unittest.mock import Mock +# Model the live RAG test drives. Override to exercise a different one. +OLLAMA_MODEL = os.environ.get("SIMPLEVECDB_TEST_OLLAMA_MODEL", "qwen3.5:0.8b") + # Stub Ollama if not installed _ollama_available = False try: @@ -13,6 +16,25 @@ except ImportError: OllamaClient = Mock() # type: ignore + +def _ollama_model_ready() -> bool: + """True only when a reachable daemon has OLLAMA_MODEL pulled. + + Importing the client proves nothing: the package installs as a plain + dependency, so without this probe the test fails with a 404 on any + machine that has the library but not the model. + """ + if not _ollama_available: + return False + try: + installed = OllamaClient().list().get("models", []) + except Exception: + return False + return any((m.get("model") or m.get("name")) == OLLAMA_MODEL for m in installed) + + +_ollama_ready = _ollama_model_ready() + from simplevecdb import VectorDB # noqa: E402 @@ -52,7 +74,9 @@ def mock_generate(prompt) -> dict[str, str]: # Real Ollama test — runs only when a local Ollama server has the # `qwen3.5:0.8b` model pulled. Skipped in CI (no Ollama daemon, no model) # and skipped locally when the daemon is unreachable. -@pytest.mark.skipif(not _ollama_available, reason="Ollama not installed") +@pytest.mark.skipif( + not _ollama_ready, reason=f"Ollama daemon or model {OLLAMA_MODEL!r} unavailable" +) @pytest.mark.skipif( bool(os.environ.get("CI")), reason="CI environments do not run a local Ollama server", @@ -80,7 +104,7 @@ def test_rag_with_ollama(populated_db): contexts = populated_db.collection("default").similarity_search(query_emb, k=2) context = "\n".join(d.page_content for d, _ in contexts) response = client.generate( - model="qwen3.5:0.8b", + model=OLLAMA_MODEL, prompt=f"Using context: {context}, answer: {query}", ) assert "purple" in response["response"].lower() diff --git a/tests/unit/test_core.py b/tests/unit/test_core.py index f3bc43e..50363b8 100755 --- a/tests/unit/test_core.py +++ b/tests/unit/test_core.py @@ -75,7 +75,11 @@ def test_upsert(populated_db): collection = populated_db.collection("default") new_emb = [0.5, 0.5, 0.5, 0.5] collection.add_texts( - ["updated apple"], embeddings=[new_emb], ids=[1], metadatas=[{"color": "green"}] + ["updated apple"], + embeddings=[new_emb], + ids=[1], + metadatas=[{"color": "green"}], + on_conflict="replace", ) updated = populated_db.conn.execute( diff --git a/tests/unit/test_id_reservation_and_tx.py b/tests/unit/test_id_reservation_and_tx.py new file mode 100644 index 0000000..e830d63 --- /dev/null +++ b/tests/unit/test_id_reservation_and_tx.py @@ -0,0 +1,250 @@ +"""Id reservation, explicit-id collision policy, and transactional vectors. + +Covers: +- ``VectorCollection.reserve_ids`` — hand out ids before the rows exist, so a + self-referential or grouped batch goes in with one write. +- ``add_texts(on_conflict=...)`` — an explicit id that already exists is an + error by default instead of a silent overwrite. +- ``tx()`` covering the HNSW index, not just SQLite: a rollback must undo the + vectors along with the rows. +- ``AsyncVectorCollection.atomic`` — the callback form that keeps a + transaction's acquire and release on one executor thread. +""" + +from __future__ import annotations + +import asyncio +import contextlib + +import pytest + +from simplevecdb import VectorDB +from simplevecdb.async_core import AsyncVectorDB + + +VEC_A = [1.0, 0.0] +VEC_B = [0.0, 1.0] + + +@pytest.fixture +def collection(): + # The VectorDB has to stay referenced for the whole test: dropping it + # closes the SQLite connection out from under the collection. + db = VectorDB(":memory:") + coll = db.collection("docs") + coll.add_texts(["first"], embeddings=[VEC_A]) + yield coll + db.close() + + +class TestReserveIds: + def test_returns_consecutive_ids_without_writing_rows(self, collection): + reserved = collection.reserve_ids(3) + + assert reserved == [reserved[0], reserved[0] + 1, reserved[0] + 2] + # Reserving is not inserting. + assert collection.count() == 1 + + def test_reserved_ids_are_never_auto_assigned_later(self, collection): + reserved = collection.reserve_ids(5) + + auto_id = collection.add_texts(["later"], embeddings=[VEC_B])[0] + + assert auto_id > reserved[-1] + + def test_reserved_ids_are_usable_for_self_referential_rows(self, collection): + reserved = collection.reserve_ids(2) + group = reserved[0] + + written = collection.add_texts( + ["a", "b"], + metadatas=[{"id": i, "group_id": group} for i in reserved], + embeddings=[VEC_A, VEC_B], + ids=reserved, + ) + + assert written == reserved + docs = collection.get_documents(filter_dict={"group_id": group}) + assert {meta["id"] for _, _, meta in docs} == set(reserved) + + def test_reservation_clears_explicitly_inserted_ids(self, collection): + collection.add_texts(["high"], embeddings=[VEC_A], ids=[9_000]) + + reserved = collection.reserve_ids(2) + + assert reserved[0] > 9_000 + + def test_works_on_an_empty_collection(self): + db = VectorDB(":memory:") + coll = db.collection("empty") + + assert coll.reserve_ids(2) == [1, 2] + + @pytest.mark.parametrize("count", [0, -1]) + def test_rejects_non_positive_count(self, collection, count): + with pytest.raises(ValueError, match="must be positive"): + collection.reserve_ids(count) + + +class TestOnConflict: + def test_existing_id_raises_by_default(self, collection): + with pytest.raises(ValueError, match="already exist"): + collection.add_texts(["clobber"], embeddings=[VEC_B], ids=[1]) + + def test_rejected_call_writes_nothing(self, collection): + with contextlib.suppress(ValueError): + collection.add_texts( + ["clobber", "innocent"], + embeddings=[VEC_B, VEC_B], + ids=[1, 4_242], + ) + + assert collection.count() == 1 + assert collection.get_documents()[0][1] == "first" + + def test_replace_overwrites(self, collection): + collection.add_texts( + ["replaced"], + embeddings=[VEC_B], + ids=[1], + metadatas=[{"tag": "new"}], + on_conflict="replace", + ) + + docs = collection.get_documents() + assert len(docs) == 1 + assert docs[0][1] == "replaced" + assert docs[0][2]["tag"] == "new" + + def test_duplicate_ids_within_one_call_raise(self, collection): + with pytest.raises(ValueError, match="duplicate"): + collection.add_texts(["a", "b"], embeddings=[VEC_A, VEC_B], ids=[500, 500]) + + def test_unknown_policy_raises(self, collection): + with pytest.raises(ValueError, match="on_conflict must be one of"): + collection.add_texts( + ["a"], embeddings=[VEC_A], ids=[500], on_conflict="upsert" + ) + + def test_fresh_explicit_ids_still_insert(self, collection): + written = collection.add_texts(["new"], embeddings=[VEC_B], ids=[77]) + + assert written == [77] + assert collection.count() == 2 + + +class TestTransactionCoversTheIndex: + def test_commit_applies_both_stores(self, collection): + with collection.tx() as coll: + coll.add_texts(["second"], embeddings=[VEC_B]) + + assert collection.count() == 2 + assert collection._index.size == 2 + + def test_rollback_undoes_both_stores(self, collection): + with pytest.raises(RuntimeError): + with collection.tx() as coll: + coll.add_texts(["doomed"], embeddings=[VEC_B]) + raise RuntimeError("boom") + + assert collection.count() == 1 + # The regression this guards: the row rolled back but the vector + # stayed, leaving the index keyed to a row that no longer exists. + assert collection._index.size == 1 + + def test_rollback_undoes_deletes(self, collection): + with pytest.raises(RuntimeError): + with collection.tx() as coll: + coll.delete_by_ids([1]) + raise RuntimeError("boom") + + assert collection.count() == 1 + assert collection._index.size == 1 + + def test_vector_writes_are_deferred_until_commit(self, collection): + with collection.tx() as coll: + coll.add_texts(["second"], embeddings=[VEC_B]) + # Documented caveat: the index does not see the transaction's + # own writes until it commits. + assert coll._index.size == 1 + + assert collection._index.size == 2 + + def test_inner_rollback_keeps_outer_writes(self, collection): + with collection.tx() as outer: + outer.add_texts(["outer"], embeddings=[VEC_B]) + with contextlib.suppress(RuntimeError): + with collection.tx() as inner: + inner.add_texts(["inner"], embeddings=[VEC_A]) + raise RuntimeError("boom") + + assert collection.count() == 2 + assert collection._index.size == 2 + + def test_buffer_does_not_leak_between_transactions(self, collection): + with pytest.raises(RuntimeError): + with collection.tx() as coll: + coll.add_texts(["doomed"], embeddings=[VEC_B]) + raise RuntimeError("boom") + + with collection.tx() as coll: + coll.add_texts(["kept"], embeddings=[VEC_B]) + + assert collection.count() == 2 + assert collection._index.size == 2 + + +@pytest.mark.asyncio +class TestAsyncAtomic: + async def _collection(self): + """Return (db, collection); the caller must keep `db` alive.""" + db = AsyncVectorDB(":memory:") + coll = db.collection("docs") + await coll.add_texts(["first"], embeddings=[VEC_A]) + return db, coll + + async def test_commits_and_returns_the_callback_result(self): + db, coll = await self._collection() + + def body(sync_coll): + sync_coll.add_texts(["second"], embeddings=[VEC_B]) + return "done" + + assert await coll.atomic(body) == "done" + assert await coll.count() == 2 + assert coll._collection._index.size == 2 + + async def test_rolls_back_both_stores(self): + db, coll = await self._collection() + + def body(sync_coll): + sync_coll.add_texts(["doomed"], embeddings=[VEC_B]) + raise RuntimeError("boom") + + with pytest.raises(RuntimeError): + await coll.atomic(body) + + assert await coll.count() == 1 + assert coll._collection._index.size == 1 + + async def test_concurrent_atomics_do_not_break_the_lock(self): + """Two-phase `async with` would release the RLock cross-thread here.""" + db, coll = await self._collection() + + def body(sync_coll): + sync_coll.add_texts(["batch"], embeddings=[VEC_B]) + return sync_coll.count() + + # More tasks than the default 4-worker pool, so they queue. + counts = await asyncio.gather(*(coll.atomic(body) for _ in range(8))) + + assert sorted(counts) == list(range(2, 10)) + assert await coll.count() == 9 + + async def test_reserve_ids_is_exposed(self): + db, coll = await self._collection() + + reserved = await coll.reserve_ids(2) + + assert len(reserved) == 2 + assert await coll.count() == 1 From f0ecca5faa200b4e8c1e9fd7f84782ef5a14f9f2 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 08:23:08 -0500 Subject: [PATCH 3/6] async API parity, crash and concurrency fixes --- CHANGELOG.md | 70 ++++++ docs/CHANGELOG.md | 70 ++++++ docs/Features.md | 5 +- docs/api/async.md | 45 +++- docs/examples.md | 10 +- pyproject.toml | 5 +- src/simplevecdb/async_core.py | 272 ++++++++++++++++++++- src/simplevecdb/core.py | 42 +++- src/simplevecdb/engine/catalog.py | 14 +- src/simplevecdb/engine/quantization.py | 6 +- src/simplevecdb/engine/usearch_index.py | 96 +++++--- src/simplevecdb/utils.py | 10 + tests/unit/test_async_parity.py | 233 ++++++++++++++++++ tests/unit/test_concurrency_hardening.py | 286 +++++++++++++++++++++++ tests/unit/test_id_reservation_and_tx.py | 34 +++ uv.lock | 2 +- 16 files changed, 1154 insertions(+), 46 deletions(-) create mode 100644 tests/unit/test_async_parity.py create mode 100644 tests/unit/test_concurrency_hardening.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 52acf64..4185d60 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,76 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [2.7.0] - 2026-07-29 +### The async API now mirrors the sync one + +The async surface was written wrapper by wrapper and had drifted: eight +collection methods and four database methods had no async counterpart, and +the sub-namespaces were flattened into names that no longer matched +(`collection.ttl.sweep()` became `sweep_ttl()`). + +Async sub-namespaces are now generic proxies over the sync ones, so async +code reads as sync code with `await` in front — and a method added to a sync +namespace is reachable from async immediately, with no wrapper to write. + +```python +await collection.edges.upsert(src, dst, kind="cites", weight=0.9) +await collection.ttl.sweep() +await collection.pending.flush() +``` + +#### Added + +- **`collection.edges`, `.events`, `.ttl`, `.pending`, `.maintenance`, + `.counters`** on `AsyncVectorCollection`, matching the sync namespaces name + for name and signature for signature. +- **`async with collection.tx()`** — the async mirror of `collection.tx()`. + A transaction holds a `threading.RLock` for its lifetime, so entering and + exiting on two different pool workers would release a lock the thread never + acquired; each transaction therefore gets a private single-worker executor + and every step runs on that one thread. Operate through the yielded handle: + awaiting work on the outer collection inside the block sends it to the + shared pool, where it blocks on the lock the transaction holds. + `atomic(fn)` remains available and makes that mistake unrepresentable — its + body is synchronous and cannot await at all. +- **`AsyncVectorDB.transaction(fn)`** — database-wide transactions spanning + collections, plus `save`, `as_langchain`, and `as_llama_index`. +- **`AsyncVectorCollection.add_texts_streaming`**. +- **`events.subscribe` is a real async generator**, polling with + `asyncio.sleep` instead of blocking the event loop. + +A parity test now fails if a public sync method gains no async counterpart, +or if a namespace signature drifts. + +#### Fixed + +- **Cancelling an `async with` transaction no longer wedges the database.** + Teardown ran through an `await`, and suspending while a `GeneratorExit` is + in flight raises "async generator ignored GeneratorExit" — leaving the + savepoint open and the DB lock held for the life of the process. Teardown + is now driven without suspending. + +### Crash and concurrency fixes + +- **Deleting from a memory-mapped index segfaulted the process.** `add()` + reloaded a `view=True` index as writable before mutating; `remove()` did + not, and usearch does not raise on a read-only mapping — it crashes. Any + database whose index file passed the 50 MB mmap threshold and then saw a + delete was exposed. Both paths now go through one guard. +- **Index reads could dereference a closed index.** `search`, `get`, + `remove`, `size`, `contains`, and `keys` checked `_index is None` and then + re-read the attribute, so a concurrent `close()` produced `AttributeError` + or `TypeError`. They now snapshot the reference once. +- **A locked write inside a transaction is no longer retried.** Outside a + transaction the write helper rolls a failed attempt back before retrying; + inside one it deliberately does not, so re-running the body re-executed + statements that had already applied, duplicating auto-id rows. +- **`close()` and `delete_collection()` stop TTL sweepers.** Background + sweeper threads outlived both, waking on their interval to query a closed + database or dropped tables and logging a failure every cycle. +- **`usearch>=2.24.0`** — earlier versions can underflow `Index::size()` to + ~1.8e19 under concurrent add/remove (unum-cloud/usearch#697), which drives + spurious full index rebuilds. + ### Transactions cover the vector index `tx()` and `db.transaction()` guarded only half the store: the SAVEPOINT diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 52acf64..4185d60 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -7,6 +7,76 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [2.7.0] - 2026-07-29 +### The async API now mirrors the sync one + +The async surface was written wrapper by wrapper and had drifted: eight +collection methods and four database methods had no async counterpart, and +the sub-namespaces were flattened into names that no longer matched +(`collection.ttl.sweep()` became `sweep_ttl()`). + +Async sub-namespaces are now generic proxies over the sync ones, so async +code reads as sync code with `await` in front — and a method added to a sync +namespace is reachable from async immediately, with no wrapper to write. + +```python +await collection.edges.upsert(src, dst, kind="cites", weight=0.9) +await collection.ttl.sweep() +await collection.pending.flush() +``` + +#### Added + +- **`collection.edges`, `.events`, `.ttl`, `.pending`, `.maintenance`, + `.counters`** on `AsyncVectorCollection`, matching the sync namespaces name + for name and signature for signature. +- **`async with collection.tx()`** — the async mirror of `collection.tx()`. + A transaction holds a `threading.RLock` for its lifetime, so entering and + exiting on two different pool workers would release a lock the thread never + acquired; each transaction therefore gets a private single-worker executor + and every step runs on that one thread. Operate through the yielded handle: + awaiting work on the outer collection inside the block sends it to the + shared pool, where it blocks on the lock the transaction holds. + `atomic(fn)` remains available and makes that mistake unrepresentable — its + body is synchronous and cannot await at all. +- **`AsyncVectorDB.transaction(fn)`** — database-wide transactions spanning + collections, plus `save`, `as_langchain`, and `as_llama_index`. +- **`AsyncVectorCollection.add_texts_streaming`**. +- **`events.subscribe` is a real async generator**, polling with + `asyncio.sleep` instead of blocking the event loop. + +A parity test now fails if a public sync method gains no async counterpart, +or if a namespace signature drifts. + +#### Fixed + +- **Cancelling an `async with` transaction no longer wedges the database.** + Teardown ran through an `await`, and suspending while a `GeneratorExit` is + in flight raises "async generator ignored GeneratorExit" — leaving the + savepoint open and the DB lock held for the life of the process. Teardown + is now driven without suspending. + +### Crash and concurrency fixes + +- **Deleting from a memory-mapped index segfaulted the process.** `add()` + reloaded a `view=True` index as writable before mutating; `remove()` did + not, and usearch does not raise on a read-only mapping — it crashes. Any + database whose index file passed the 50 MB mmap threshold and then saw a + delete was exposed. Both paths now go through one guard. +- **Index reads could dereference a closed index.** `search`, `get`, + `remove`, `size`, `contains`, and `keys` checked `_index is None` and then + re-read the attribute, so a concurrent `close()` produced `AttributeError` + or `TypeError`. They now snapshot the reference once. +- **A locked write inside a transaction is no longer retried.** Outside a + transaction the write helper rolls a failed attempt back before retrying; + inside one it deliberately does not, so re-running the body re-executed + statements that had already applied, duplicating auto-id rows. +- **`close()` and `delete_collection()` stop TTL sweepers.** Background + sweeper threads outlived both, waking on their interval to query a closed + database or dropped tables and logging a failure every cycle. +- **`usearch>=2.24.0`** — earlier versions can underflow `Index::size()` to + ~1.8e19 under concurrent add/remove (unum-cloud/usearch#697), which drives + spurious full index rebuilds. + ### Transactions cover the vector index `tx()` and `db.transaction()` guarded only half the store: the SAVEPOINT diff --git a/docs/Features.md b/docs/Features.md index d2159bc..d14f358 100644 --- a/docs/Features.md +++ b/docs/Features.md @@ -84,8 +84,9 @@ release-by-release detail, see the [Changelog](CHANGELOG.md). Vector mutations (`add_texts`, `delete_by_ids`, `pending.flush()`, `ttl.sweep()`) are buffered and applied on commit, so a rollback leaves neither store changed; the tradeoff is that a search inside - the transaction cannot see its own vector writes. Async callers use - `await collection.atomic(fn)` instead of `async with`. + the transaction cannot see its own vector writes. Async callers get the + same shape — `async with collection.tx() as coll:` — plus + `await collection.atomic(fn)` for a callback that cannot await. - **Reserved ids** — `collection.reserve_ids(n)` hands out ids before the rows exist, so self-referential or grouped documents go in with a single `add_texts` call. Reusing an existing id now raises instead of diff --git a/docs/api/async.md b/docs/api/async.md index 0c05c0c..0488644 100755 --- a/docs/api/async.md +++ b/docs/api/async.md @@ -57,13 +57,44 @@ Synchronous properties remain unchanged: - `collection.name` - Collection name +## Sub-namespaces + +Every sync sub-namespace is mirrored on the async collection, name for name: + +```python +await collection.edges.upsert(src, dst, kind="cites", weight=0.9) +await collection.counters.increment(doc_id, {"hits": 1}) +await collection.ttl.sweep() +await collection.pending.flush() + +async for event in collection.events.subscribe(since=0): + ... +``` + ## Transactions -`collection.tx()` has no `async with` equivalent — the transaction holds a -`threading.RLock` for its lifetime, and entering and exiting in two separate -executor tasks can release that lock from a thread that never acquired it. -Pass a synchronous callback to `atomic()` instead, and the whole transaction -runs in one executor thread: +`async with collection.tx()` mirrors the sync context manager. Catalog writes +and vector writes commit or roll back together: + +```python +async with collection.tx() as coll: + await coll.delete_by_ids([1]) + await coll.add_texts(["replacement"], embeddings=[[0.1] * 384]) +``` + +A transaction holds a `threading.RLock` for its lifetime, so its enter and +exit must happen on one thread — the shared pool cannot promise that, and +releasing an RLock from a thread that never acquired it raises. Each +transaction therefore gets a private single-worker executor. + +**Operate through the yielded handle.** It is bound to that pinned thread; +the outer collection is not. Awaiting work on the outer handle inside the +block sends it to the shared pool, where it blocks on the DB lock this +transaction holds — and that lock is not released until the block exits, +which cannot happen while it is awaiting. + +`atomic()` makes that mistake unrepresentable: the callback is synchronous +and cannot await at all. ```python def move(coll): @@ -73,8 +104,8 @@ def move(coll): new_ids = await collection.atomic(move) ``` -The callback receives the underlying sync `VectorCollection` and must not -await. Catalog writes and vector writes commit or roll back together. +`AsyncVectorDB.transaction(fn)` is the database-wide equivalent, spanning +collections via `tx["name"]`. ## Concurrent Operations diff --git a/docs/examples.md b/docs/examples.md index 7eeb6b5..cb2fa9e 100755 --- a/docs/examples.md +++ b/docs/examples.md @@ -287,13 +287,17 @@ with collection.tx() as coll: # rows and vectors both survive, or neither does ``` -On the async side there is no `async with`; pass a sync callback to -`atomic()` instead — see the async API reference for why. +The async side mirrors this. Operate through the yielded handle — it is +bound to the transaction's pinned thread, and the outer collection is not: ```python -await collection.atomic(lambda coll: coll.add_texts(["x"], embeddings=[vec])) +async with collection.tx() as coll: + await coll.add_texts(["x"], embeddings=[vec]) ``` +`await collection.atomic(fn)` runs a synchronous callback instead, which +cannot await and so cannot make that mistake. See the async API reference. + ### Reserving ids `reserve_ids(n)` burns ids out of the auto-increment sequence without diff --git a/pyproject.toml b/pyproject.toml index f38c258..0941d24 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,10 @@ classifiers = [ dependencies = [ "numpy>=1.24", - "usearch>=2.16.3", + # 2.24.0 is the first release with the Index::size() underflow fix + # (unum-cloud/usearch#697). Below it, concurrent add/remove can make + # size() report ~1.8e19, which drives spurious full index rebuilds. + "usearch>=2.24.0", "scikit-learn>=1.3.0", # Clustering (K-means, MiniBatch K-means, HDBSCAN) "hdbscan>=0.8.33", # Density-based clustering "sqlcipher3-binary>=0.5.0", # Encryption support diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index 0bf934e..e4eefb7 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -27,17 +27,95 @@ import asyncio import functools from concurrent.futures import ThreadPoolExecutor -from collections.abc import Callable, Sequence +from collections.abc import AsyncIterator, Callable, Iterable, Sequence +from contextlib import asynccontextmanager from threading import Lock from typing import Any, TypeVar import logging -from .core import VectorDB, VectorCollection +from .core import VectorDB, VectorCollection, _DBTransaction from .types import Document, DistanceStrategy, OnConflict, Quantization T = TypeVar("T") + +class _AsyncNamespace: + """Awaitable mirror of a sync sub-namespace (`collection.edges`, …). + + Every public callable on the wrapped namespace is re-exposed as a + coroutine that runs the sync call in the executor. Async code therefore + reads exactly like sync code with `await` in front — + `await coll.edges.upsert(...)` against `coll.edges.upsert(...)` — and a + method added to a sync namespace is reachable from async immediately, + with no wrapper to write and no way for the two surfaces to drift. + + Non-callable attributes pass through unchanged. + """ + + __slots__ = ("_namespace", "_run") + + def __init__(self, namespace: Any, run: Callable[..., Any]) -> None: + self._namespace = namespace + self._run = run + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + attr = getattr(self._namespace, name) + if not callable(attr): + return attr + + @functools.wraps(attr) + async def _in_executor(*args: Any, **kwargs: Any) -> Any: + return await self._run(attr, *args, **kwargs) + + return _in_executor + + def __dir__(self) -> list[str]: + return sorted(set(dir(self._namespace)) | set(object.__dir__(self))) + + def __repr__(self) -> str: + return f"Async{type(self._namespace).__name__.lstrip('_')}" + + +class _AsyncEventsNamespace(_AsyncNamespace): + """Events namespace with a real async `subscribe`. + + The sync `subscribe` is a blocking generator that sleeps between polls; + driving it from async would stall the event loop, so the loop is + reimplemented here over the async `read` with `asyncio.sleep`. + """ + + async def subscribe( + self, + *, + since: int = 0, + kind: str | None = None, + poll_interval: float | None = None, + batch: int = 500, + ) -> "AsyncIterator[Any]": + """Async generator yielding events as they appear. Caller controls exit.""" + from . import constants + + interval = ( + constants.EVENTS_POLL_INTERVAL_S if poll_interval is None else poll_interval + ) + last = int(since) + while True: + events = await self._run( + self._namespace.read, since=last, kind=kind, limit=batch + ) + if events: + for event in events: + yield event + last = events[-1].seq + if len(events) == batch: + # Drained a full batch; go again without sleeping. + continue + await asyncio.sleep(interval) + + _logger = logging.getLogger(__name__) @@ -53,9 +131,15 @@ def __init__( self, sync_collection: VectorCollection, executor: ThreadPoolExecutor, + *, + tx_pinned: bool = False, ): self._collection = sync_collection self._executor = executor + # True when `executor` is a transaction's private single-worker pool, + # so a nested tx() reuses that thread instead of pinning a new one. + self._tx_pinned = tx_pinned + self._namespaces: dict[str, _AsyncNamespace] = {} @property def name(self) -> str: @@ -71,6 +155,144 @@ async def _run(self, fn, /, *args, **kwargs): self._executor, functools.partial(fn, *args, **kwargs) ) + def _namespace(self, name: str) -> _AsyncNamespace: + """Awaitable mirror of the sync namespace `name`, created once.""" + proxy = self._namespaces.get(name) + if proxy is None: + cls = _AsyncEventsNamespace if name == "events" else _AsyncNamespace + proxy = cls(getattr(self._collection, name), self._run) + self._namespaces[name] = proxy + return proxy + + @property + def edges(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.edges`.""" + return self._namespace("edges") + + @property + def events(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.events`.""" + return self._namespace("events") + + @property + def ttl(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.ttl`.""" + return self._namespace("ttl") + + @property + def pending(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.pending`.""" + return self._namespace("pending") + + @property + def maintenance(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.maintenance`.""" + return self._namespace("maintenance") + + @property + def counters(self) -> _AsyncNamespace: + """Awaitable mirror of `VectorCollection.counters`.""" + return self._namespace("counters") + + async def add_texts_streaming( + self, + items: Iterable[tuple[str, dict | None, Sequence[float] | None]], + *, + batch_size: int | None = None, + threads: int = 0, + on_progress: Any = None, + ) -> list[int]: + """Stream documents in batches, returning every inserted id. + + The sync method is a generator yielding per-batch progress; the whole + drain runs in one executor task here, so `on_progress` fires from + that thread rather than the event loop. + + See VectorCollection.add_texts_streaming for full documentation. + """ + + def _drain(coll: VectorCollection) -> list[int]: + ids: list[int] = [] + for progress in coll.add_texts_streaming( + items, + batch_size=batch_size, + threads=threads, + on_progress=on_progress, + ): + ids.extend(progress["batch_ids"]) + return ids + + return await self._run(_drain, self._collection) + + @asynccontextmanager + async def tx(self) -> AsyncIterator["AsyncVectorCollection"]: + """Async mirror of `VectorCollection.tx()`. + + async with collection.tx() as coll: + await coll.delete_by_ids([1]) + await coll.add_texts(["replacement"], embeddings=[vec]) + + A transaction holds a `threading.RLock` for its lifetime, so its + enter and exit must happen on one thread. The shared pool cannot + promise that — two executor tasks may land on different workers, and + releasing an RLock from a thread that never acquired it raises + `RuntimeError: cannot release un-acquired lock`. So the transaction + gets a private single-worker executor: one thread, every step on it. + + **Operate through the yielded handle.** It is bound to the pinned + thread; the outer collection is not. Awaiting work on the outer + handle inside the block sends it to the shared pool, where it blocks + on the DB lock this transaction holds — and that lock is only + released when the block exits, which cannot happen while it is + awaiting. Use `atomic()` if you want that mistake to be + unrepresentable: its body is synchronous and cannot await at all. + """ + loop = asyncio.get_running_loop() + # A nested tx must run on the thread that already holds the lock: the + # RLock is reentrant per thread, so pinning a second thread here + # would block forever waiting on the outer transaction. + reuse = self._tx_pinned + pinned = ( + self._executor + if reuse + else ThreadPoolExecutor( + max_workers=1, thread_name_prefix=f"simplevecdb-tx-{self.name}" + ) + ) + + def _on_pinned(fn: Callable[..., Any], *args: Any) -> Any: + return loop.run_in_executor(pinned, functools.partial(fn, *args)) + + try: + manager = self._collection.tx() + scoped_sync = await _on_pinned(manager.__enter__) + scoped = AsyncVectorCollection(scoped_sync, pinned, tx_pinned=True) + try: + yield scoped + except (GeneratorExit, asyncio.CancelledError) as exc: + # Teardown paths where awaiting is unsafe. Suspending while a + # GeneratorExit is in flight raises "async generator ignored + # GeneratorExit", and a cancelled task's next await can be + # cancelled again — either way the savepoint would stay open + # and the DB lock would never be released, wedging every other + # writer. Drive the exit on the pinned thread without + # suspending. + pinned.submit( + manager.__exit__, type(exc), exc, exc.__traceback__ + ).result() + raise + except BaseException as exc: + await _on_pinned(manager.__exit__, type(exc), exc, exc.__traceback__) + raise + await _on_pinned(manager.__exit__, None, None, None) + finally: + if not reuse: + # Non-blocking: the transaction is over, so waiting here would + # only stall the event loop on work the body queued and never + # awaited — work that must not run now that the savepoint has + # closed. cancel_futures drops exactly that. + pinned.shutdown(wait=False, cancel_futures=True) + async def atomic(self, fn: Callable[[VectorCollection], T]) -> T: """Run `fn` inside a transaction on this collection. @@ -813,6 +1035,52 @@ async def _run(self, fn, /, *args, **kwargs): self._executor, functools.partial(fn, *args, **kwargs) ) + async def save(self) -> None: + """Persist every collection's index to disk. + + See VectorDB.save for full documentation. + """ + await self._run(self._db.save) + + async def transaction(self, fn: Callable[["_DBTransaction"], T]) -> T: + """Run `fn` inside a database-wide transaction. + + The DB-level counterpart to `AsyncVectorCollection.atomic`: `fn` is a + synchronous callable receiving the transaction handle, so it can span + collections via `tx["name"]`. Everything it does — catalog and vector + writes across every collection — commits or rolls back together. + + def move(tx): + tx["archive"].add_texts(texts, embeddings=vecs) + tx["inbox"].delete_by_ids(ids) + + await db.transaction(move) + + A callback rather than `async with` for the same reason as `atomic`: + the transaction holds a `threading.RLock` across its lifetime and + must acquire and release it on one thread. `fn` must not await. + """ + + def _in_tx() -> T: + with self._db.transaction() as tx: + return fn(tx) + + return await self._run(_in_tx) + + def as_langchain(self, embeddings: Any = None, collection_name: str = "default"): + """Return a LangChain-compatible vector store over the sync database. + + Synchronous by design: LangChain drives its own async surface. + """ + return self._db.as_langchain(embeddings, collection_name) + + def as_llama_index(self, collection_name: str = "default"): + """Return a LlamaIndex-compatible vector store over the sync database. + + Synchronous by design: LlamaIndex drives its own async surface. + """ + return self._db.as_llama_index(collection_name) + def list_collections(self) -> list[str]: """Return names of all persisted collections in the database.""" return self._db.list_collections() diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 788884a..febba66 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -340,8 +340,13 @@ def _index_add( The buffered arrays are copied: they may alias caller-owned memory that is free to change before the transaction commits. + + Only the thread that owns the transaction defers. Another thread's + write is not part of it — its rows are already committed — so + buffering would hand its vectors to a transaction that can roll + back and drop them. """ - if self._tx_state.depth > 0: + if self._tx_state.owned_by_current_thread(): keys_buf = np.array(keys, dtype=np.uint64, copy=True) vecs_buf = np.array(vectors, dtype=np.float32, copy=True) @@ -357,7 +362,7 @@ def _index_remove(self, keys: list[int]) -> None: Counterpart to `_index_add`; see there for why the write is held. """ - if self._tx_state.depth > 0: + if self._tx_state.owned_by_current_thread(): keys_buf = [int(k) for k in keys] def _deferred_remove() -> None: @@ -2272,6 +2277,7 @@ def __enter__(self) -> "_DBTransaction": depth = self._db._tx_state.depth name = f"simplevecdb_tx_{depth + 1}" self._db.conn.execute(f"SAVEPOINT {name}") + self._db._tx_state.owner = threading.get_ident() self._db._tx_state.depth = depth + 1 self._savepoint_name = name self._entered = True @@ -2333,6 +2339,8 @@ def __exit__(self, exc_type, exc, tb) -> None: self._db.conn.execute(f"RELEASE SAVEPOINT {name}") finally: self._db._tx_state.depth = max(0, self._db._tx_state.depth - 1) + if self._db._tx_state.depth == 0: + self._db._tx_state.owner = None # Outermost commit: if depth fell to 0, finalize the # implicit Python sqlite3 transaction so changes flush. if self._db._tx_state.depth == 0 and exc_type is None: @@ -2769,6 +2777,20 @@ def delete_collection(self, name: str) -> None: # Close any cached collection's open index before removing the file for cached_key, cached_col in list(self._collections.items()): if cached_key[0] == name: + # Stop the TTL sweeper first: left running it keeps + # querying tables this method is about to drop, logging a + # failure every interval for the life of the process. + ttl_ns = cached_col.__dict__.get("_ttl_ns") + if ttl_ns is not None: + try: + ttl_ns.stop_background() + except Exception: + _logger.warning( + "Failed to stop TTL sweeper for collection %r " + "during delete", + name, + exc_info=True, + ) try: cached_col._index.close() except Exception: @@ -3045,6 +3067,22 @@ def close(self) -> None: if getattr(self, "_closed", False): return self._closed = True + # Stop opt-in TTL sweepers first. They are daemon threads holding a + # reference to this connection: left running, they wake on their + # interval and fail against a closed database, logging a warning + # every cycle until the process exits — and a sweep in flight can + # race the conn.close() below. + for col in self._collections.values(): + ttl_ns = col.__dict__.get("_ttl_ns") + if ttl_ns is not None: + try: + ttl_ns.stop_background() + except Exception: + _logger.warning( + "Failed to stop TTL sweeper for collection %s during close", + col.name, + exc_info=True, + ) try: self.save() except Exception: diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 0933e21..d6dc73e 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -92,13 +92,25 @@ class _TxState: usearch index, so vector writes are held here and applied only when the outermost transaction is about to release; a rollback truncates the buffer instead, leaving the two stores in step. + + `owner` is the thread id holding the transaction. A writer that is not + the owner is by definition not inside the transaction — its rows are + already committed — so it must apply its vectors immediately rather + than buffer them into a transaction that may roll back and discard + them. Vector writes happen outside the DB lock, so identity is what + decides this, not `depth` alone. """ - __slots__ = ("depth", "index_ops") + __slots__ = ("depth", "index_ops", "owner") def __init__(self) -> None: self.depth: int = 0 self.index_ops: list[Callable[[], None]] = [] + self.owner: int | None = None + + def owned_by_current_thread(self) -> bool: + """True when the calling thread is inside this transaction.""" + return self.depth > 0 and self.owner == threading.get_ident() class _CatalogWritable: diff --git a/src/simplevecdb/engine/quantization.py b/src/simplevecdb/engine/quantization.py index 5c265a3..f453994 100755 --- a/src/simplevecdb/engine/quantization.py +++ b/src/simplevecdb/engine/quantization.py @@ -109,7 +109,11 @@ def deserialize(self, blob: bytes, dim: int | None) -> np.ndarray: ValueError: If quantization mode unsupported or dim missing for BIT """ if self.quantization == Quantization.FLOAT: - return np.frombuffer(blob, dtype=np.float32) + # np.frombuffer aliases the blob and is read-only. Every other + # branch below returns a fresh writable array (astype/where copy), + # so copy here too rather than hand back a vector whose + # mutability depends on the quantization mode. + return np.frombuffer(blob, dtype=np.float32).copy() elif self.quantization == Quantization.INT8: return np.frombuffer(blob, dtype=np.int8).astype(np.float32) / 127.0 diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index 9a1df7f..ef2f659 100755 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -200,9 +200,13 @@ def ndim(self) -> int | None: @property def size(self) -> int: """Number of vectors in the index.""" - if self._index is None: + # Snapshot the reference: close() clears `_index` without the write + # lock, so re-reading the attribute after the None check can hand a + # concurrent reader `len(None)`. + index = self._index + if index is None: return 0 - return len(self._index) + return len(index) @property def is_memory_mapped(self) -> bool: @@ -240,13 +244,7 @@ def add( # Lazy index creation on first add with self._write_lock: - # If currently in view mode, need to reload as writable - if self._is_view and self._index is not None: - _logger.debug("Upgrading from view to writable mode for add operation") - from usearch.index import Index - - self._index = Index.restore(str(self._path), view=False) - self._is_view = False + self._ensure_writable("add") if self._index is None: self._ndim = vectors.shape[1] @@ -308,7 +306,11 @@ def search( """ from .. import constants - if self._index is None or self.size == 0: + # Snapshot once: close() nulls `_index`, and search deliberately takes + # no lock, so re-reading the attribute later can hand this method None + # halfway through. + index = self._index + if index is None or len(index) == 0: # Return empty results for empty index empty_keys = np.array([], dtype=np.uint64) empty_dists = np.array([], dtype=np.float32) @@ -332,20 +334,21 @@ def search( query = query / np.maximum(norms, 1e-12) # Adaptive search: brute-force for small indexes, HNSW for large + size = len(index) if exact is None: - use_exact = self.size < constants.USEARCH_BRUTEFORCE_THRESHOLD + use_exact = size < constants.USEARCH_BRUTEFORCE_THRESHOLD else: use_exact = exact if use_exact: _logger.debug( "Using brute-force search (index size %d < threshold %d)", - self.size, + size, constants.USEARCH_BRUTEFORCE_THRESHOLD, ) # usearch search is thread-safe for reads - matches = self._index.search(query, k, exact=use_exact, threads=threads) + matches = index.search(query, k, exact=use_exact, threads=threads) # Handle single query vs batch keys = np.asarray(matches.keys, dtype=np.uint64) @@ -353,6 +356,25 @@ def search( return keys, distances + def _ensure_writable(self, operation: str) -> None: + """Reload a memory-mapped index as writable before mutating it. + + A `view=True` index is mapped read-only; mutating one does not raise, + it segfaults the process inside usearch. Every mutating entry point + must pass through here first. + + Caller must hold `_write_lock`. + """ + if not self._is_view or self._index is None: + return + _logger.debug( + "Upgrading from view to writable mode for %s operation", operation + ) + from usearch.index import Index + + self._index = Index.restore(str(self._path), view=False) + self._is_view = False + def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: """ Remove vectors by their keys. @@ -376,11 +398,22 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: return 0 with self._write_lock: + # Mapped read-only indexes must be reloaded writable first, or the + # remove below segfaults rather than raising. + self._ensure_writable("remove") + + # Re-read after taking the lock: close() also takes it, so the + # None check above may be stale by now. _ensure_writable can + # rebind `_index` too, so this must come after it. + index = self._index + if index is None: + return 0 + # Filter to only keys that exist in the index - existing_mask = np.array([int(k) in self._index for k in keys], dtype=bool) + existing_mask = np.array([int(k) in index for k in keys], dtype=bool) existing_keys = keys[existing_mask] if len(existing_keys) > 0: - self._index.remove(existing_keys) + index.remove(existing_keys) removed = int(existing_mask.sum()) self._dirty = True _logger.debug("Removed %d vectors from index", removed) @@ -388,9 +421,10 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: def contains(self, key: int) -> bool: """Check if a key exists in the index.""" - if self._index is None: + index = self._index # see `size` for why this is snapshotted + if index is None: return False - return key in self._index + return key in index def save(self) -> None: """Save index to disk atomically if modified. @@ -457,7 +491,10 @@ def save(self) -> None: def close(self) -> None: """Save and close the index.""" self.save() - self._index = None + # Under the write lock so a mutation in flight finishes against a + # live index rather than losing it mid-operation. + with self._write_lock: + self._index = None def __len__(self) -> int: return self.size @@ -467,18 +504,24 @@ def __contains__(self, key: int) -> bool: def keys(self) -> list[int]: """Return all keys in the index.""" - if self._index is None: + index = self._index # see `size` for why this is snapshotted + if index is None: return [] - return [int(k) for k in self._index.keys] + return [int(k) for k in index.keys] - def _vectors_from_index(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: + def _vectors_from_index( + self, keys: NDArray[np.uint64], index: Any | None = None + ) -> NDArray[np.float32]: """Fetch stored vectors for keys, unpacking BIT-quantized bytes to ±1 floats. For BIT quantization usearch stores packed bytes (ndim/8 per vector); a plain float cast would yield the wrong shape and meaningless values, so the bits are unpacked back to the float dimension. + + `index` lets a caller pass the snapshot it already validated, rather + than re-reading `_index` and racing close(). """ - raw = self._index[keys] + raw = (self._index if index is None else index)[keys] if self._quantization == Quantization.BIT: return _unpack_bits(np.asarray(raw, dtype=np.uint8), self._ndim or 1) return np.asarray(raw, dtype=np.float32) @@ -493,14 +536,15 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: Returns: Array of vectors, shape (len(keys), ndim). Missing keys return zeros. """ - if self._index is None or len(keys) == 0: + index = self._index # snapshot; see `search` for why + if index is None or len(keys) == 0: return np.array([], dtype=np.float32).reshape(0, self._ndim or 1) keys = np.asarray(keys, dtype=np.uint64) ndim = self._ndim or 1 # Filter to existing keys for batch retrieval - existing_mask = np.array([int(k) in self._index for k in keys], dtype=bool) + existing_mask = np.array([int(k) in index for k in keys], dtype=bool) if not existing_mask.any(): _logger.warning( @@ -511,7 +555,7 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: if existing_mask.all(): # Fast path: all keys exist, batch retrieve - return self._vectors_from_index(keys) + return self._vectors_from_index(keys, index) # Mixed: some keys missing _logger.warning( @@ -520,7 +564,7 @@ def get(self, keys: NDArray[np.uint64]) -> NDArray[np.float32]: ) result = np.zeros((len(keys), ndim), dtype=np.float32) existing_keys = keys[existing_mask] - result[existing_mask] = self._vectors_from_index(existing_keys) + result[existing_mask] = self._vectors_from_index(existing_keys, index) return result def __del__(self) -> None: diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index 14f22ec..3b80536 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -113,6 +113,16 @@ def retry_on_lock( def decorator(func: F) -> F: @wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: + # Never retry inside the caller's transaction. There, the write + # helper suppresses the per-call commit, so a statement that + # failed halfway through leaves its earlier statements applied + # and un-rolled-back; re-running the body would insert the + # auto-id rows a second time. The transaction owns atomicity, so + # let the error reach it and be rolled back as a unit. + tx_state = getattr(args[0], "_tx_state", None) if args else None + if tx_state is not None and tx_state.owned_by_current_thread(): + return func(*args, **kwargs) + last_exception: sqlite3.OperationalError | None = None total_wait = 0.0 diff --git a/tests/unit/test_async_parity.py b/tests/unit/test_async_parity.py new file mode 100644 index 0000000..99e2828 --- /dev/null +++ b/tests/unit/test_async_parity.py @@ -0,0 +1,233 @@ +"""The async surface must mirror the sync one, permanently. + +The async API used to be hand-written wrapper by wrapper, so it drifted: +eight collection methods and four database methods had no async counterpart, +and the sub-namespaces were flattened into inconsistent names +(`ttl.sweep()` became `sweep_ttl()`). + +Async namespaces are now generic proxies over the sync ones, so a method +added to a sync namespace is reachable from async with no wrapper to write. +These tests hold the line for the parts that are still declared by hand: add +a public sync method without an async counterpart and this file fails. +""" + +from __future__ import annotations + +import asyncio +import inspect + +import pytest + +from simplevecdb.async_core import AsyncVectorCollection, AsyncVectorDB +from simplevecdb.core import VectorCollection, VectorDB + +# Nothing is sync-only any more. `tx()` used to be: a transaction holds a +# threading.RLock for its lifetime, and on the shared pool an `async with` +# could enter and exit on different workers, releasing a lock the thread +# never acquired. Giving each transaction a private single-worker executor +# removes that, so the async surface now mirrors sync completely. +INTENTIONALLY_SYNC_ONLY: set[str] = set() + +NAMESPACES = ("edges", "events", "ttl", "pending", "maintenance", "counters") + +# `events.subscribe` is an async generator polling with `asyncio.sleep` +# instead of a blocking one polling with `time.sleep`; its return annotation +# and `poll_interval` default legitimately differ from the sync method. +SIGNATURE_EXEMPT = {"subscribe"} + + +def public_api(cls: type) -> set[str]: + """Public methods and properties declared on a class.""" + return { + name + for name, value in inspect.getmembers(cls) + if not name.startswith("_") + and (inspect.isfunction(value) or isinstance(value, property)) + } + + +class TestSurfaceParity: + def test_every_collection_method_has_an_async_counterpart(self): + missing = public_api(VectorCollection) - public_api(AsyncVectorCollection) + + assert missing <= INTENTIONALLY_SYNC_ONLY, ( + f"AsyncVectorCollection is missing {sorted(missing - INTENTIONALLY_SYNC_ONLY)}. " + "Add a wrapper, or add the name to INTENTIONALLY_SYNC_ONLY with a reason." + ) + + def test_every_database_method_has_an_async_counterpart(self): + missing = public_api(VectorDB) - public_api(AsyncVectorDB) + + assert not missing, ( + f"AsyncVectorDB is missing {sorted(missing)}. " + "Add a wrapper, or document the omission." + ) + + def test_both_transaction_forms_are_available(self): + """`tx()` mirrors sync; `atomic()` is the deadlock-proof alternative.""" + assert "tx" in public_api(AsyncVectorCollection) + assert "atomic" in public_api(AsyncVectorCollection) + assert "transaction" in public_api(AsyncVectorDB) + + +@pytest.mark.asyncio +class TestNamespaceParity: + @staticmethod + def _collection(): + db = AsyncVectorDB(":memory:") + return db, db.collection("docs") + + @pytest.mark.parametrize("namespace", NAMESPACES) + async def test_namespace_exposes_every_sync_method(self, namespace): + db, coll = self._collection() + sync_ns = getattr(coll._collection, namespace) + async_ns = getattr(coll, namespace) + + for name in dir(sync_ns): + if name.startswith("_"): + continue + assert hasattr(async_ns, name), ( + f"coll.{namespace}.{name} exists on the sync namespace " + f"but not the async one" + ) + await db.close() + + @pytest.mark.parametrize("namespace", NAMESPACES) + async def test_namespace_methods_are_awaitable(self, namespace): + db, coll = self._collection() + async_ns = getattr(coll, namespace) + sync_ns = getattr(coll._collection, namespace) + + for name in dir(sync_ns): + if name.startswith("_") or not callable(getattr(sync_ns, name)): + continue + attr = getattr(async_ns, name) + assert inspect.iscoroutinefunction(attr) or inspect.isasyncgenfunction( + attr + ), f"coll.{namespace}.{name} should be awaitable, got {attr!r}" + await db.close() + + @pytest.mark.parametrize("namespace", NAMESPACES) + async def test_namespace_methods_keep_their_signature(self, namespace): + """functools.wraps must carry the sync signature onto the proxy.""" + db, coll = self._collection() + async_ns = getattr(coll, namespace) + sync_ns = getattr(coll._collection, namespace) + + for name in dir(sync_ns): + sync_attr = getattr(sync_ns, name) + if ( + name.startswith("_") + or not callable(sync_attr) + or name in SIGNATURE_EXEMPT + ): + continue + assert inspect.signature(getattr(async_ns, name)) == inspect.signature( + sync_attr + ), f"coll.{namespace}.{name} signature drifted from sync" + await db.close() + + async def test_async_with_tx_commits(self): + db, coll = self._collection() + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + async with coll.tx() as scoped: + await scoped.add_texts(["committed"], embeddings=[[0.0, 1.0]]) + + assert await coll.count() == 2 + assert coll._collection._index.size == 2 + await db.close() + + async def test_async_with_tx_rolls_back_both_stores(self): + db, coll = self._collection() + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + with pytest.raises(RuntimeError): + async with coll.tx() as scoped: + await scoped.add_texts(["doomed"], embeddings=[[0.0, 1.0]]) + raise RuntimeError("boom") + + assert await coll.count() == 1 + assert coll._collection._index.size == 1 + await db.close() + + async def test_nested_async_with_tx_reuses_the_pinned_thread(self): + """A nested tx must not pin a second thread, or it self-deadlocks.""" + db, coll = self._collection() + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + async def nested(): + async with coll.tx() as outer: + await outer.add_texts(["outer"], embeddings=[[0.1, 0.9]]) + async with outer.tx() as inner: + await inner.add_texts(["inner"], embeddings=[[0.9, 0.1]]) + + await asyncio.wait_for(nested(), timeout=15) + + assert await coll.count() == 3 + assert coll._collection._index.size == 3 + await db.close() + + async def test_cancelling_a_transaction_releases_the_lock(self): + """Teardown must not await, or the savepoint and DB lock leak. + + `async with` is driven by an async generator: cancelling the task + throws GeneratorExit in at the yield, and suspending there raises + "async generator ignored GeneratorExit" — leaving the transaction + open and every other writer wedged behind its lock. + """ + db, coll = self._collection() + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + entered = asyncio.Event() + + async def body(): + async with coll.tx() as scoped: + await scoped.add_texts(["cancelled"], embeddings=[[0.0, 1.0]]) + entered.set() + await asyncio.sleep(30) + + task = asyncio.create_task(body()) + await asyncio.wait_for(entered.wait(), timeout=10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # Would hang if the lock leaked, and the cancelled write must be gone. + async def after(): + async with coll.tx() as scoped: + await scoped.add_texts(["after"], embeddings=[[0.3, 0.7]]) + + await asyncio.wait_for(after(), timeout=15) + + assert await coll.count() == 2 + assert coll._collection._index.size == 2 + await db.close() + + async def test_namespaces_work_through_a_scoped_transaction(self): + db, coll = self._collection() + + async with coll.tx() as scoped: + ids = await scoped.add_texts(["a"], embeddings=[[1.0, 0.0]]) + await scoped.counters.increment(ids[0], {"hits": 1}) + + docs = await coll.get_documents() + assert docs[0][2]["hits"] == 1 + await db.close() + + async def test_events_subscribe_is_a_real_async_generator(self): + """The sync generator blocks between polls; async must not.""" + db, coll = self._collection() + + assert inspect.isasyncgenfunction(coll.events.subscribe) + + await coll.add_texts(["a"], embeddings=[[1.0, 0.0]]) + await coll.events.append("marker") + + seen = [] + async for event in coll.events.subscribe(since=0, poll_interval=0.01): + seen.append(event.kind) + break + + assert seen + await db.close() diff --git a/tests/unit/test_concurrency_hardening.py b/tests/unit/test_concurrency_hardening.py new file mode 100644 index 0000000..dba5f2a --- /dev/null +++ b/tests/unit/test_concurrency_hardening.py @@ -0,0 +1,286 @@ +"""Regressions for crashes and races found auditing the storage layer. + +Each test here corresponds to a defect that was reachable from ordinary API +use, not a hypothetical: + +- Removing from a memory-mapped index segfaulted the process. +- Retrying a locked write inside a transaction duplicated rows. +""" + +from __future__ import annotations + +import asyncio +import sqlite3 +import subprocess +import sys +import threading +import time + +import numpy as np +import pytest + +from simplevecdb import VectorDB +from simplevecdb.async_core import AsyncVectorDB +from simplevecdb.engine.catalog import _TxState +from simplevecdb.utils import DatabaseLockedError, retry_on_lock + + +class TestMemoryMappedRemove: + """`remove()` on a `view=True` index mutated a read-only mapping. + + usearch does not raise on that — it segfaults, taking the process with + it. Any database whose index file passes the mmap threshold and then + sees a delete was exposed. + """ + + def test_delete_against_mapped_index_does_not_crash(self, tmp_path): + # Run out-of-process: the failure mode is SIGSEGV, which no + # in-process assertion can catch. + script = f""" +import numpy as np +from simplevecdb import constants +constants.USEARCH_MMAP_THRESHOLD = 1 # force the mmap path on a small index +from simplevecdb import VectorDB + +path = {str(tmp_path / "v.db")!r} +db = VectorDB(path) +coll = db.collection("docs") +coll.add_texts( + [f"t{{i}}" for i in range(20)], + embeddings=np.random.rand(20, 8).astype(np.float32).tolist(), +) +coll.save() +db.close() + +db2 = VectorDB(path) +coll2 = db2.collection("docs") +assert coll2._index._is_view, "expected the index to load memory-mapped" +coll2.delete_by_ids([1, 2, 3]) +print(coll2.count(), coll2._index.size) +""" + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True + ) + + assert result.returncode == 0, ( + f"expected clean exit, got {result.returncode} " + f"(-11 is SIGSEGV). stderr: {result.stderr[-500:]}" + ) + assert result.stdout.split() == ["17", "17"] + + def test_mapped_index_is_upgraded_in_place(self, tmp_path, monkeypatch): + from simplevecdb import constants + + monkeypatch.setattr(constants, "USEARCH_MMAP_THRESHOLD", 1) + + db = VectorDB(str(tmp_path / "v.db")) + coll = db.collection("docs") + coll.add_texts( + [f"t{i}" for i in range(20)], + embeddings=np.random.rand(20, 8).astype(np.float32).tolist(), + ) + coll.save() + db.close() + + db2 = VectorDB(str(tmp_path / "v.db")) + coll2 = db2.collection("docs") + assert coll2._index._is_view + + coll2.delete_by_ids([1]) + + # The remove must have taken the index out of view mode, not + # mutated the mapping underneath it. + assert not coll2._index._is_view + assert coll2._index.size == 19 + db2.close() + + +class TestSearchIsolation: + """Searches must not surface another thread's uncommitted transaction. + + All collections share one SQLite connection and a connection has no + isolation between its own operations, so raw `db.conn` reads *can* see + uncommitted rows. Vector search cannot, for two reasons worth pinning + down: a transaction's vector writes are buffered until commit, so an + uncommitted row has no index entry to be found by; and the catalog reads + a search makes each take the DB lock already. + + These tests exist so that neither property can be removed silently — drop + the write buffering and the first one fails. + """ + + def test_search_does_not_see_a_rolled_back_write(self): + db = VectorDB(":memory:") + coll = db.collection("docs") + coll.add_texts(["committed"], embeddings=[[1.0, 0.0]]) + + started = threading.Event() + seen: dict[str, int] = {} + + def reader(): + assert started.wait(5) + seen["hits"] = len(coll.similarity_search([1.0, 0.0], k=10)) + + thread = threading.Thread(target=reader) + thread.start() + try: + with coll.tx() as tx: + tx.add_texts(["doomed"], embeddings=[[0.0, 1.0]]) + started.set() + # Give the reader time to reach the search and block on it. + thread.join(0.3) + raise RuntimeError("boom") + except RuntimeError: + pass + thread.join(5) + + assert seen["hits"] == 1 + assert coll.count() == 1 + db.close() + + def test_search_inside_its_own_transaction_does_not_deadlock(self): + """The DB lock is reentrant, so the owning thread may still search.""" + db = VectorDB(":memory:") + coll = db.collection("docs") + coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + with coll.tx() as tx: + tx.add_texts(["inside"], embeddings=[[0.0, 1.0]]) + hits = tx.similarity_search([1.0, 0.0], k=5) + + # Only the committed row is visible: the transaction's own vector + # write is buffered until commit. + assert len(hits) == 1 + assert coll.count() == 2 + db.close() + + +@pytest.mark.asyncio +class TestAsyncSearchIsolation: + """The same guarantees must hold through the executor. + + Async searches run in pool threads while a transaction runs in another, + so this covers the same two properties as the sync case across the + executor boundary. `atomic()` never awaits inside its body, so the lock + holder always makes progress on its own worker and cannot be starved by + queued searches. + """ + + async def test_async_search_does_not_see_a_rolled_back_write(self): + db = AsyncVectorDB(":memory:") + coll = db.collection("docs") + await coll.add_texts(["committed"], embeddings=[[1.0, 0.0]]) + + started = asyncio.Event() + loop = asyncio.get_running_loop() + + def tx_body(sync_coll): + sync_coll.add_texts(["doomed"], embeddings=[[0.0, 1.0]]) + loop.call_soon_threadsafe(started.set) + time.sleep(0.3) # hold the transaction open across the search + raise RuntimeError("boom") + + async def reader(): + await started.wait() + return len(await coll.similarity_search([1.0, 0.0], k=10)) + + tx = asyncio.create_task(coll.atomic(tx_body)) + read = asyncio.create_task(reader()) + + with pytest.raises(RuntimeError): + await tx + assert await read == 1 + assert await coll.count() == 1 + await db.close() + + async def test_saturating_the_pool_with_searches_and_transactions(self): + db = AsyncVectorDB(":memory:") + coll = db.collection("docs") + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + def write(sync_coll): + sync_coll.add_texts(["x"], embeddings=[[0.3, 0.7]]) + + async def task(i): + if i % 2: + return await coll.similarity_search([1.0, 0.0], k=3) + return await coll.atomic(write) + + # Far more tasks than the default 4 workers. + results = await asyncio.wait_for( + asyncio.gather(*(task(i) for i in range(24))), timeout=30 + ) + + assert len(results) == 24 + assert await coll.count() == 13 # seed + 12 writers + assert coll._collection._index.size == 13 + await db.close() + + async def test_search_inside_atomic_does_not_deadlock(self): + db = AsyncVectorDB(":memory:") + coll = db.collection("docs") + await coll.add_texts(["seed"], embeddings=[[1.0, 0.0]]) + + def body(sync_coll): + sync_coll.add_texts(["inner"], embeddings=[[0.0, 1.0]]) + return len(sync_coll.similarity_search([1.0, 0.0], k=5)) + + hits = await asyncio.wait_for(coll.atomic(body), timeout=10) + + assert hits == 1 # the transaction's own vector write is deferred + assert await coll.count() == 2 + await db.close() + + +class TestRetryInsideTransaction: + """`@retry_on_lock` must not re-run a body inside a caller's transaction. + + Outside one, the write helper enters the connection context and a failed + attempt is rolled back before the retry. Inside one it does not, so the + retry would re-execute statements that already applied — duplicating the + auto-id INSERT. + """ + + @staticmethod + def _locking_op(): + calls: list[int] = [] + + class Fake: + _tx_state = _TxState() + + @retry_on_lock(max_retries=3, base_delay=0.001) + def op(self): + calls.append(1) + raise sqlite3.OperationalError("database is locked") + + return Fake(), calls + + def test_retries_when_no_transaction_is_open(self): + fake, calls = self._locking_op() + + with pytest.raises(DatabaseLockedError): + fake.op() + + assert len(calls) == 4 # initial attempt + 3 retries + + def test_does_not_retry_inside_the_owning_thread_transaction(self): + fake, calls = self._locking_op() + fake._tx_state.owner = threading.get_ident() + fake._tx_state.depth = 1 + + # The raw error must surface so the transaction rolls back as a unit. + with pytest.raises(sqlite3.OperationalError): + fake.op() + + assert len(calls) == 1 + + def test_still_retries_for_a_thread_outside_the_transaction(self): + fake, calls = self._locking_op() + # Another thread owns the transaction; this caller is not in it. + fake._tx_state.owner = threading.get_ident() + 1 + fake._tx_state.depth = 1 + + with pytest.raises(DatabaseLockedError): + fake.op() + + assert len(calls) == 4 diff --git a/tests/unit/test_id_reservation_and_tx.py b/tests/unit/test_id_reservation_and_tx.py index e830d63..0525850 100644 --- a/tests/unit/test_id_reservation_and_tx.py +++ b/tests/unit/test_id_reservation_and_tx.py @@ -15,6 +15,7 @@ import asyncio import contextlib +import threading import pytest @@ -181,6 +182,39 @@ def test_inner_rollback_keeps_outer_writes(self, collection): assert collection.count() == 2 assert collection._index.size == 2 + def test_another_thread_is_not_captured_by_an_open_transaction(self, collection): + """A writer outside the transaction must not have its vectors buffered. + + Its rows are already committed, so buffering them into someone else's + transaction means a rollback silently drops the vectors and leaves the + index short — the exact divergence the buffering exists to prevent. + """ + import numpy as np + + started, release = threading.Event(), threading.Event() + + def owner(): + with contextlib.suppress(RuntimeError): + with collection.tx(): + started.set() + release.wait(5) + raise RuntimeError("boom") + + thread = threading.Thread(target=owner) + thread.start() + assert started.wait(5) + + # Stands in for a non-transactional writer past its catalog commit. + collection._index_add( + np.array([99], dtype=np.uint64), + np.array([VEC_B], dtype=np.float32), + ) + + release.set() + thread.join(5) + + assert collection._index.size == 2 + def test_buffer_does_not_leak_between_transactions(self, collection): with pytest.raises(RuntimeError): with collection.tx() as coll: diff --git a/uv.lock b/uv.lock index 025f975..65f1991 100755 --- a/uv.lock +++ b/uv.lock @@ -4773,7 +4773,7 @@ requires-dist = [ { name = "scikit-learn", specifier = ">=1.3.0" }, { name = "sentence-transformers", marker = "extra == 'server'", specifier = ">=5.0" }, { name = "sqlcipher3-binary", specifier = ">=0.5.0" }, - { name = "usearch", specifier = ">=2.16.3" }, + { name = "usearch", specifier = ">=2.24.0" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = ">=0.30" }, ] provides-extras = ["integrations", "server", "examples"] From 36906830eeab9b81e245d93eb6b82dcb586bfe6b Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 08:46:39 -0500 Subject: [PATCH 4/6] one sqlite connection per thread for real transaction isolation --- CHANGELOG.md | 25 +++ docs/CHANGELOG.md | 25 +++ src/simplevecdb/async_core.py | 17 ++ src/simplevecdb/core.py | 92 +++++---- src/simplevecdb/engine/catalog.py | 57 ++++-- src/simplevecdb/engine/connection.py | 229 +++++++++++++++++++++++ tests/unit/test_concurrency_hardening.py | 94 +++++++++- 7 files changed, 478 insertions(+), 61 deletions(-) create mode 100644 src/simplevecdb/engine/connection.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4185d60..bcec794 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,31 @@ or if a namespace signature drifts. savepoint open and the DB lock held for the life of the process. Teardown is now driven without suspending. +### One connection per thread + +A single shared `sqlite3.Connection` served every collection and every +thread. SQLite provides no isolation between operations on one connection, +so a read on one thread could observe another thread's uncommitted rows and +act on data that was about to roll back. + +File-backed databases now open one connection per thread. In WAL mode that +gives snapshot isolation, and readers no longer wait behind a writer. + +In-memory databases keep a single shared connection, and keep the stale-read +behaviour with it. Pooling one requires a shared-cache URI — a plain +`":memory:"` gives every connection its own separate database — and shared +cache takes table-level write locks, so a concurrent reader fails outright +with `SQLITE_LOCKED` ("database table is locked"), which `busy_timeout` does +not wait out. A hard error is worse than a stale read for a database that +cannot outlive the process. + +Connection PRAGMAs (`foreign_keys`, `busy_timeout`, `synchronous`) are now +applied per connection, as they are connection-scoped rather than stored in +the database file. + +`db.conn` and `collection.conn` resolve to the calling thread's connection. +Both remain assignable for injecting a connection. + ### Crash and concurrency fixes - **Deleting from a memory-mapped index segfaulted the process.** `add()` diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 4185d60..bcec794 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -55,6 +55,31 @@ or if a namespace signature drifts. savepoint open and the DB lock held for the life of the process. Teardown is now driven without suspending. +### One connection per thread + +A single shared `sqlite3.Connection` served every collection and every +thread. SQLite provides no isolation between operations on one connection, +so a read on one thread could observe another thread's uncommitted rows and +act on data that was about to roll back. + +File-backed databases now open one connection per thread. In WAL mode that +gives snapshot isolation, and readers no longer wait behind a writer. + +In-memory databases keep a single shared connection, and keep the stale-read +behaviour with it. Pooling one requires a shared-cache URI — a plain +`":memory:"` gives every connection its own separate database — and shared +cache takes table-level write locks, so a concurrent reader fails outright +with `SQLITE_LOCKED` ("database table is locked"), which `busy_timeout` does +not wait out. A hard error is worse than a stale read for a database that +cannot outlive the process. + +Connection PRAGMAs (`foreign_keys`, `busy_timeout`, `synchronous`) are now +applied per connection, as they are connection-scoped rather than stored in +the database file. + +`db.conn` and `collection.conn` resolve to the calling thread's connection. +Both remain assignable for injecting a connection. + ### Crash and concurrency fixes - **Deleting from a memory-mapped index segfaulted the process.** `add()` diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index e4eefb7..2236b7d 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -146,6 +146,15 @@ def name(self) -> str: """Collection name.""" return self._collection.name + @property + def conn(self) -> Any: + """This thread's SQLite connection on the underlying collection. + + A handle, not an operation, so it is not awaitable. Note that reads + issued directly on it bypass the collection API entirely. + """ + return self._collection.conn + def __repr__(self) -> str: return f"AsyncVectorCollection(name={self._collection.name!r})" @@ -1081,6 +1090,14 @@ def as_llama_index(self, collection_name: str = "default"): """ return self._db.as_llama_index(collection_name) + @property + def conn(self) -> Any: + """This thread's SQLite connection on the underlying database. + + A handle, not an operation, so it is not awaitable. + """ + return self._db.conn + def list_collections(self) -> list[str]: """Return names of all persisted collections in the database.""" return self._db.list_collections() diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index febba66..996da4f 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -39,11 +39,11 @@ from .engine.quantization import QuantizationStrategy from .engine.search import SearchEngine from .engine.catalog import CatalogManager, _TxState +from .engine.connection import ConnectionSource, as_source, open_source from .engine.usearch_index import UsearchIndex from .engine.clustering import ClusterEngine, ClusterAlgorithm from . import constants from .encryption import ( - create_encrypted_connection, encrypt_index_file, decrypt_index_file, get_encrypted_index_path, @@ -177,7 +177,7 @@ class VectorCollection: def __init__( self, - conn: sqlite3.Connection, + conn: "sqlite3.Connection | ConnectionSource", db_path: str, name: str, distance_strategy: DistanceStrategy, @@ -187,7 +187,7 @@ def __init__( lock: threading.RLock | None = None, tx_state: _TxState | None = None, ): - self.conn = conn + self._source = as_source(conn) self._db_path = db_path self.name = name self.distance_strategy = distance_strategy @@ -231,7 +231,7 @@ def __init__( # The optional _tx_state is also shared with the parent VectorDB # so db.transaction() can suspend per-call commits everywhere. self._catalog = CatalogManager( - conn=self.conn, + conn=self._source, table_name=self._table_name, fts_table_name=self._fts_table_name, lock=self._lock, @@ -1770,6 +1770,16 @@ def maintenance(self) -> "_MaintenanceNamespace": self.__dict__["_maint_ns"] = ns return ns + @property + def conn(self) -> sqlite3.Connection: + """This thread's SQLite connection.""" + return self._source.conn + + @conn.setter + def conn(self, value: Any) -> None: + """Replace the connection source (tests and connection injection).""" + self._source = as_source(value) + @property def dim(self) -> int | None: """Vector dimension (None if no vectors added yet).""" @@ -2277,7 +2287,6 @@ def __enter__(self) -> "_DBTransaction": depth = self._db._tx_state.depth name = f"simplevecdb_tx_{depth + 1}" self._db.conn.execute(f"SAVEPOINT {name}") - self._db._tx_state.owner = threading.get_ident() self._db._tx_state.depth = depth + 1 self._savepoint_name = name self._entered = True @@ -2339,8 +2348,6 @@ def __exit__(self, exc_type, exc, tb) -> None: self._db.conn.execute(f"RELEASE SAVEPOINT {name}") finally: self._db._tx_state.depth = max(0, self._db._tx_state.depth - 1) - if self._db._tx_state.depth == 0: - self._db._tx_state.owner = None # Outermost commit: if depth fell to 0, finalize the # implicit Python sqlite3 transaction so changes flush. if self._db._tx_state.depth == 0 and exc_type is None: @@ -2378,12 +2385,17 @@ def __enter__(self) -> "VectorCollection": # type: ignore[override] class _CollectionTxShim: """Minimal proxy emulating the VectorDB attributes _DBTransaction reads.""" - __slots__ = ("_lock", "_tx_state", "conn") + __slots__ = ("_lock", "_tx_state", "_collection") def __init__(self, collection: "VectorCollection") -> None: self._lock = collection._lock self._tx_state = collection._tx_state - self.conn = collection.conn + self._collection = collection + + @property + def conn(self) -> sqlite3.Connection: + """Resolved per access: the connection is per-thread.""" + return self._collection.conn class _EdgesNamespace: @@ -2608,45 +2620,43 @@ def __init__( # so all catalogs in this DB suspend per-call commits. self._tx_state: _TxState = _TxState() - # Create connection (encrypted or plain) - if encryption_key is not None: - if self.path == ":memory:": - raise ValueError( - "In-memory databases cannot be encrypted. " - "Use a file path for encrypted databases." - ) - self.conn = create_encrypted_connection( - self.path, - encryption_key, - check_same_thread=False, - timeout=30.0, + # One connection per thread, all against this database. A single + # shared connection cannot isolate transactions — SQLite provides no + # isolation between operations on one connection, so a reader would + # see another thread's uncommitted rows. + if encryption_key is not None and self.path == ":memory:": + raise ValueError( + "In-memory databases cannot be encrypted. " + "Use a file path for encrypted databases." ) - self.conn.execute("PRAGMA journal_mode=WAL") - self.conn.execute("PRAGMA synchronous=NORMAL") - # Native lock-wait window so SQLite blocks the caller in C - # rather than surfacing 'database is locked' immediately - # under multi-writer load (gap 10). - self.conn.execute(f"PRAGMA busy_timeout={constants.SQLITE_BUSY_TIMEOUT_MS}") - self.conn.execute("PRAGMA foreign_keys=ON") - self._encrypted = True + self._source: Any = open_source(self.path, encryption_key=encryption_key) + self._encrypted = encryption_key is not None + if self._encrypted: _logger.info("Opened encrypted database: %s", self.path) - else: - self.conn = sqlite3.connect( - self.path, check_same_thread=False, timeout=30.0 - ) - self.conn.execute("PRAGMA journal_mode=WAL") - self.conn.execute("PRAGMA synchronous=NORMAL") - self.conn.execute(f"PRAGMA busy_timeout={constants.SQLITE_BUSY_TIMEOUT_MS}") - self.conn.execute("PRAGMA foreign_keys=ON") - self._encrypted = False # Verify connection is healthy try: self.conn.execute("SELECT 1") except sqlite3.DatabaseError as e: - self.conn.close() + self._source.close_all() raise RuntimeError(f"Database health check failed: {e}") from e + @property + def conn(self) -> sqlite3.Connection: + """This thread's SQLite connection, opened on first use. + + Each thread gets its own, so one thread's open transaction is + invisible to another's reads — a single shared connection cannot + provide that, because SQLite has no isolation between operations on + one connection. + """ + return self._source.conn + + @conn.setter + def conn(self, value: Any) -> None: + """Replace the connection source (tests and connection injection).""" + self._source = as_source(value) + def transaction(self) -> "_DBTransaction": """Atomic write context spanning all collections (gap 2). @@ -3003,7 +3013,7 @@ def collection( with self._lock: if cache_key not in self._collections: self._collections[cache_key] = VectorCollection( - conn=self.conn, + conn=self._source, db_path=self.path, name=name, distance_strategy=distance_strategy or self.distance_strategy, @@ -3088,7 +3098,7 @@ def close(self) -> None: except Exception: _logger.warning("Failed to save indexes during close", exc_info=True) finally: - self.conn.close() + self._source.close_all() # Clean up ephemeral usearch index files created for in-memory DBs. for col in self._collections.values(): ephemeral = getattr(col, "_ephemeral_index_path", None) diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index d6dc73e..a4c3fd6 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -17,6 +17,7 @@ from ..constants import SQLITE_MAX_BOUND_PARAMS from ..types import ON_CONFLICT_POLICIES, OnConflict +from .connection import ConnectionSource, as_source from ..utils import _batched from ..utils import validate_filter, retry_on_lock, normalize_filter, find_duplicates @@ -93,24 +94,37 @@ class _TxState: the outermost transaction is about to release; a rollback truncates the buffer instead, leaving the two stores in step. - `owner` is the thread id holding the transaction. A writer that is not - the owner is by definition not inside the transaction — its rows are - already committed — so it must apply its vectors immediately rather - than buffer them into a transaction that may roll back and discard - them. Vector writes happen outside the DB lock, so identity is what - decides this, not `depth` alone. + The object is shared by every catalog in a database, but everything it + stores is per thread: each thread owns a separate SQLite connection and + therefore a separate transaction. That also means a writer on another + thread reads `depth == 0` and applies its vectors immediately, rather + than buffering them into a transaction that may roll back and drop them. """ - __slots__ = ("depth", "index_ops", "owner") + __slots__ = ("_local",) def __init__(self) -> None: - self.depth: int = 0 - self.index_ops: list[Callable[[], None]] = [] - self.owner: int | None = None + self._local = threading.local() + + @property + def depth(self) -> int: + return getattr(self._local, "depth", 0) + + @depth.setter + def depth(self, value: int) -> None: + self._local.depth = value + + @property + def index_ops(self) -> list[Callable[[], None]]: + ops = getattr(self._local, "index_ops", None) + if ops is None: + ops = [] + self._local.index_ops = ops + return ops def owned_by_current_thread(self) -> bool: - """True when the calling thread is inside this transaction.""" - return self.depth > 0 and self.owner == threading.get_ident() + """True when the calling thread is inside a transaction.""" + return self.depth > 0 class _CatalogWritable: @@ -268,7 +282,7 @@ class CatalogManager: def __init__( self, - conn: sqlite3.Connection, + conn: "sqlite3.Connection | ConnectionSource", table_name: str, fts_table_name: str, lock: threading.RLock | None = None, @@ -282,7 +296,7 @@ def __init__( _validate_table_name(table_name) _validate_table_name(fts_table_name) - self.conn = conn + self._source = as_source(conn) self._table_name = table_name self._fts_table_name = fts_table_name self._fts_enabled = False @@ -311,6 +325,21 @@ def _writable(self): """ return _CatalogWritable(self._lock, self.conn, self._tx_state) + @property + def conn(self) -> "sqlite3.Connection": + """This thread's SQLite connection. + + A property rather than a stored handle: each thread owns a separate + connection, which is what keeps one thread's open transaction + invisible to another's reads. + """ + return self._source.conn + + @conn.setter + def conn(self, value: Any) -> None: + """Replace the connection source (tests and connection injection).""" + self._source = as_source(value) + def create_tables(self) -> None: """Create metadata and FTS tables if they don't exist.""" self.conn.execute( diff --git a/src/simplevecdb/engine/connection.py b/src/simplevecdb/engine/connection.py new file mode 100644 index 0000000..6fb2f14 --- /dev/null +++ b/src/simplevecdb/engine/connection.py @@ -0,0 +1,229 @@ +"""Per-thread SQLite connections for SimpleVecDB. + +SQLite gives no isolation between operations on a *single* connection: a +reader sharing one connection with a writer sees that writer's uncommitted +rows, and can act on data that later rolls back. Isolation is a property of +having separate connections — with them, "the reader is only able to see +complete committed transactions from the writer... regardless of whether the +two database connections are in the same thread, in different threads of the +same process, or in different processes." + +So each thread gets its own connection to the same database. In WAL mode that +also buys snapshot isolation and lets readers run while a writer commits. +""" + +from __future__ import annotations + +import logging +import sqlite3 +import threading +import uuid +from typing import Any, Protocol + +from .. import constants + +_logger = logging.getLogger("simplevecdb.engine.connection") + + +class ConnectionSource(Protocol): + """What the catalog and collection layers need from a connection owner.""" + + @property + def conn(self) -> sqlite3.Connection: + """The connection this thread should use.""" + ... + + def close_all(self) -> None: + """Close every connection handed out.""" + ... + + +def is_in_memory(path: str) -> bool: + """Whether `path` names an in-memory database.""" + return path == ":memory:" + + +def shared_memory_dsn() -> str: + """A named shared-cache DSN for an in-memory database. + + Not used by default, and the reason is worth recording. Pooling an + in-memory database requires this form — *"opening two database + connections each with the filename ':memory:' will create two + independent in-memory databases"* — but shared cache takes **table-level** + write locks, and a reader on another connection then fails with + ``SQLITE_LOCKED`` ("database table is locked"), which `busy_timeout` does + not wait out. That trades a harmless stale read for a hard error, so + in-memory databases keep a single shared connection instead. + """ + return f"file:simplevecdb_{uuid.uuid4().hex}?mode=memory&cache=shared" + + +def apply_pragmas(conn: sqlite3.Connection) -> None: + """Configure a freshly opened connection. + + Every one of these is connection-scoped, not database-scoped, so a pool + must apply them per connection — `foreign_keys` silently defaults back to + off otherwise. (`journal_mode=WAL` is the exception, persisting in the + file, but setting it again is harmless and keeps this in one place.) + """ + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + # Native lock-wait window so SQLite blocks the caller in C rather than + # surfacing 'database is locked' immediately under multi-writer load. + conn.execute(f"PRAGMA busy_timeout={constants.SQLITE_BUSY_TIMEOUT_MS}") + conn.execute("PRAGMA foreign_keys=ON") + + +class SingleConnection: + """Adapter presenting one already-open connection as a connection source. + + Used when a caller supplies its own `sqlite3.Connection` directly instead + of a path, so those call sites keep working unchanged. It cannot provide + cross-thread isolation — that is the point of `ConnectionPool`. + """ + + __slots__ = ("_conn",) + + def __init__(self, conn: sqlite3.Connection) -> None: + self._conn = conn + + @property + def conn(self) -> sqlite3.Connection: + return self._conn + + def close_all(self) -> None: + self._conn.close() + + +class ConnectionPool: + """One SQLite connection per thread, all against the same database. + + The first connection is opened eagerly and kept for the pool's lifetime: + a shared-cache in-memory database is reclaimed when its *last* connection + closes, so without an anchor the data would vanish whenever a worker + thread's connection happened to be the last one out. + """ + + __slots__ = ( + "_dsn", + "_uri", + "_encryption_key", + "_timeout", + "_local", + "_all", + "_all_lock", + "_closed", + "_anchor", + ) + + def __init__( + self, + path: str, + *, + encryption_key: str | bytes | None = None, + timeout: float = 30.0, + ) -> None: + if is_in_memory(path): + raise ValueError( + "ConnectionPool cannot pool an in-memory database; " + "use open_source(), which gives it a shared connection." + ) + self._dsn, self._uri = path, False + self._encryption_key = encryption_key + self._timeout = timeout + self._local = threading.local() + self._all: list[sqlite3.Connection] = [] + self._all_lock = threading.Lock() + self._closed = False + self._anchor = self._open() + self._local.conn = self._anchor + + def _open(self) -> sqlite3.Connection: + """Open and configure one connection.""" + if self._encryption_key is not None: + from ..encryption import create_encrypted_connection + + conn = create_encrypted_connection( + self._dsn, + self._encryption_key, + check_same_thread=False, + timeout=self._timeout, + ) + else: + conn = sqlite3.connect( + self._dsn, + uri=self._uri, + check_same_thread=False, + timeout=self._timeout, + ) + apply_pragmas(conn) + with self._all_lock: + self._all.append(conn) + return conn + + @property + def conn(self) -> sqlite3.Connection: + """This thread's connection, opened on first use.""" + conn: sqlite3.Connection | None = getattr(self._local, "conn", None) + if conn is None: + if self._closed: + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + conn = self._open() + self._local.conn = conn + return conn + + def close_all(self) -> None: + """Close every connection this pool opened, from any thread. + + Connections belonging to other threads are closed here too. That is + safe because they were opened with `check_same_thread=False`, and it + is necessary because a worker thread may never run again to close its + own. + """ + self._closed = True + with self._all_lock: + connections, self._all = self._all, [] + for conn in connections: + try: + conn.close() + except Exception: # pragma: no cover - close is best-effort + _logger.debug("Failed to close a pooled connection", exc_info=True) + self._local = threading.local() + + def __repr__(self) -> str: + return f"ConnectionPool(dsn={self._dsn!r}, open={len(self._all)})" + + +def open_source( + path: str, + *, + encryption_key: str | bytes | None = None, + timeout: float = 30.0, +) -> ConnectionSource: + """Open the right kind of connection source for `path`. + + File-backed databases get one connection per thread, which is what makes + a transaction on one thread invisible to reads on another. In-memory + databases get a single shared connection: pooling one requires shared + cache, whose table-level locks turn concurrent readers into + ``SQLITE_LOCKED`` errors (see `shared_memory_dsn`). They therefore keep + the old behaviour, cross-thread dirty reads included — acceptable for a + database that cannot outlive the process. + """ + if is_in_memory(path): + conn = sqlite3.connect(path, check_same_thread=False, timeout=timeout) + apply_pragmas(conn) + return SingleConnection(conn) + return ConnectionPool(path, encryption_key=encryption_key, timeout=timeout) + + +def as_source(conn_or_source: Any) -> ConnectionSource: + """Accept either a raw connection or something already pool-shaped. + + Tested by identity rather than by `isinstance(..., sqlite3.Connection)`: + sqlcipher connections and test doubles are connection-like without being + instances of it, and wrapping those is exactly the intent. + """ + if isinstance(conn_or_source, (ConnectionPool, SingleConnection)): + return conn_or_source + return SingleConnection(conn_or_source) diff --git a/tests/unit/test_concurrency_hardening.py b/tests/unit/test_concurrency_hardening.py index dba5f2a..2e006ff 100644 --- a/tests/unit/test_concurrency_hardening.py +++ b/tests/unit/test_concurrency_hardening.py @@ -232,6 +232,78 @@ def body(sync_coll): await db.close() +class TestCrossThreadIsolation: + """A reader thread must not see another thread's uncommitted rows. + + SQLite gives no isolation between operations on one connection, so while + every collection shared a single connection, a raw read could observe a + transaction that later rolled back. File-backed databases now open one + connection per thread, which is what buys the isolation. + """ + + def test_file_backed_reader_cannot_see_an_open_transaction(self, tmp_path): + db = VectorDB(str(tmp_path / "v.db")) + coll = db.collection("docs") + coll.add_texts(["committed"], embeddings=[[1.0, 0.0]]) + + started, released = threading.Event(), threading.Event() + seen: dict[str, int] = {} + + def reader(): + assert started.wait(5) + # Raw read, bypassing the collection API entirely. + seen["count"] = db.conn.execute( + f"SELECT COUNT(*) FROM {coll._table_name}" + ).fetchone()[0] + released.set() + + thread = threading.Thread(target=reader) + thread.start() + try: + with coll.tx() as tx: + tx.add_texts(["doomed"], embeddings=[[0.0, 1.0]]) + started.set() + assert released.wait(5) + raise RuntimeError("boom") + except RuntimeError: + pass + thread.join(5) + + assert seen["count"] == 1 + assert coll.count() == 1 + db.close() + + def test_each_thread_gets_its_own_connection(self, tmp_path): + db = VectorDB(str(tmp_path / "v.db")) + seen: dict[str, int] = {} + + def other(): + seen["id"] = id(db.conn) + + thread = threading.Thread(target=other) + thread.start() + thread.join(5) + + assert seen["id"] != id(db.conn) + db.close() + + def test_in_memory_keeps_one_shared_connection(self): + """Documented carve-out: pooling :memory: needs shared cache, whose + table-level locks turn concurrent readers into SQLITE_LOCKED errors.""" + db = VectorDB(":memory:") + seen: dict[str, int] = {} + + def other(): + seen["id"] = id(db.conn) + + thread = threading.Thread(target=other) + thread.start() + thread.join(5) + + assert seen["id"] == id(db.conn) + db.close() + + class TestRetryInsideTransaction: """`@retry_on_lock` must not re-run a body inside a caller's transaction. @@ -265,7 +337,6 @@ def test_retries_when_no_transaction_is_open(self): def test_does_not_retry_inside_the_owning_thread_transaction(self): fake, calls = self._locking_op() - fake._tx_state.owner = threading.get_ident() fake._tx_state.depth = 1 # The raw error must surface so the transaction rolls back as a unit. @@ -275,12 +346,23 @@ def test_does_not_retry_inside_the_owning_thread_transaction(self): assert len(calls) == 1 def test_still_retries_for_a_thread_outside_the_transaction(self): + """Transaction state is per-thread, so another thread's depth is 0.""" fake, calls = self._locking_op() - # Another thread owns the transaction; this caller is not in it. - fake._tx_state.owner = threading.get_ident() + 1 - fake._tx_state.depth = 1 + fake._tx_state.depth = 1 # this thread is in a transaction - with pytest.raises(DatabaseLockedError): - fake.op() + outcome: dict[str, object] = {} + + def other_thread(): + # Owns no transaction of its own: must still retry. + assert fake._tx_state.depth == 0 + try: + fake.op() + except BaseException as exc: # noqa: BLE001 - recorded below + outcome["error"] = exc + + thread = threading.Thread(target=other_thread) + thread.start() + thread.join(30) + assert isinstance(outcome.get("error"), DatabaseLockedError) assert len(calls) == 4 From 9a4d3be0554cf3dfdd0079fe3138c0500b19124c Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 09:29:23 -0500 Subject: [PATCH 5/6] stand down db lock when connections are pooled; begin transactions immediate --- CHANGELOG.md | 27 +++++ docs/CHANGELOG.md | 27 +++++ src/simplevecdb/core.py | 68 ++++++++++--- src/simplevecdb/engine/catalog.py | 19 +++- src/simplevecdb/engine/connection.py | 61 ++++++++++++ tests/unit/test_concurrency_hardening.py | 119 +++++++++++++++++++++++ tests/unit/test_v26_review_pass_3.py | 18 +++- 7 files changed, 320 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bcec794..1d75a3e 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -80,6 +80,33 @@ the database file. `db.conn` and `collection.conn` resolve to the calling thread's connection. Both remain assignable for injecting a connection. +**The database-wide lock now stands down when connections are pooled.** It +was held for the entire lifetime of every transaction, which was necessary +only because all threads shared one connection's transaction context. With a +connection per thread there is nothing left for it to guard, and SQLite +serializes writers itself — blocking in C for `busy_timeout` rather than +failing. A read on another thread no longer waits out a transaction: in a +timing check it returned in 1 ms against a transaction held open for 400 ms, +and still saw the committed state rather than the open transaction's rows. + +It remains a real lock where the context genuinely is shared — in-memory +databases and injected connections — and the structural lock guarding the +collections cache and the `rebuild_index` swap is untouched, since neither +has anything to do with how connections are opened. + +A rolled-back transaction now ends its enclosing transaction as well as its +savepoint. Releasing a savepoint does not end the transaction it sits inside, +so an explicit `BEGIN IMMEDIATE` left open on the rollback path would pin +SQLite's write lock for the life of that connection and make every other +connection fail with "database is locked". + +**Transactions begin with `BEGIN IMMEDIATE`.** Across connections, a +transaction that reads before it writes can find another connection has +committed in between; SQLite will not fork history, so the upgrade fails +with `SQLITE_BUSY_SNAPSHOT`, which no `busy_timeout` waits out. Starting as +a writer takes the write lock up front, after which no operation in the +transaction fails with `SQLITE_BUSY`. + ### Crash and concurrency fixes - **Deleting from a memory-mapped index segfaulted the process.** `add()` diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index bcec794..1d75a3e 100755 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -80,6 +80,33 @@ the database file. `db.conn` and `collection.conn` resolve to the calling thread's connection. Both remain assignable for injecting a connection. +**The database-wide lock now stands down when connections are pooled.** It +was held for the entire lifetime of every transaction, which was necessary +only because all threads shared one connection's transaction context. With a +connection per thread there is nothing left for it to guard, and SQLite +serializes writers itself — blocking in C for `busy_timeout` rather than +failing. A read on another thread no longer waits out a transaction: in a +timing check it returned in 1 ms against a transaction held open for 400 ms, +and still saw the committed state rather than the open transaction's rows. + +It remains a real lock where the context genuinely is shared — in-memory +databases and injected connections — and the structural lock guarding the +collections cache and the `rebuild_index` swap is untouched, since neither +has anything to do with how connections are opened. + +A rolled-back transaction now ends its enclosing transaction as well as its +savepoint. Releasing a savepoint does not end the transaction it sits inside, +so an explicit `BEGIN IMMEDIATE` left open on the rollback path would pin +SQLite's write lock for the life of that connection and make every other +connection fail with "database is locked". + +**Transactions begin with `BEGIN IMMEDIATE`.** Across connections, a +transaction that reads before it writes can find another connection has +committed in between; SQLite will not fork history, so the upgrade fails +with `SQLITE_BUSY_SNAPSHOT`, which no `busy_timeout` waits out. Starting as +a writer takes the write lock up front, after which no operation in the +transaction fails with `SQLITE_BUSY`. + ### Crash and concurrency fixes - **Deleting from a memory-mapped index segfaulted the process.** `add()` diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index 996da4f..a92cada 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -39,7 +39,12 @@ from .engine.quantization import QuantizationStrategy from .engine.search import SearchEngine from .engine.catalog import CatalogManager, _TxState -from .engine.connection import ConnectionSource, as_source, open_source +from .engine.connection import ( + ConnectionLock, + ConnectionSource, + as_source, + open_source, +) from .engine.usearch_index import UsearchIndex from .engine.clustering import ClusterEngine, ClusterAlgorithm from . import constants @@ -186,6 +191,7 @@ def __init__( store_embeddings: bool = False, lock: threading.RLock | None = None, tx_state: _TxState | None = None, + conn_lock: "ConnectionLock | None" = None, ): self._source = as_source(conn) self._db_path = db_path @@ -198,7 +204,14 @@ def __init__( # Connection-level lock shared with the parent VectorDB so all # collections sharing the same sqlite3.Connection serialize their # transactional access from Python. + # Structural lock: guards the in-memory index swap in rebuild_index, + # which is shared across threads no matter how connections are opened. self._lock: threading.RLock = lock if lock is not None else threading.RLock() + # Connection lock: guards the Python-level transaction context, and so + # engages only when threads share one connection. + self._conn_lock: Any = ( + conn_lock if conn_lock is not None else ConnectionLock(self._source.shared) + ) # Serializes rebuild_index() calls against each other WITHOUT holding the # shared DB lock during the (slow) HNSW build. Distinct from self._lock. self._rebuild_lock = threading.Lock() @@ -234,7 +247,7 @@ def __init__( conn=self._source, table_name=self._table_name, fts_table_name=self._fts_table_name, - lock=self._lock, + lock=self._conn_lock, tx_state=getattr(self, "_tx_state", None), ) self._catalog.create_tables() @@ -2282,11 +2295,22 @@ def __init__(self, db: "VectorDB") -> None: self._index_mark = 0 def __enter__(self) -> "_DBTransaction": - self._db._lock.acquire() + self._db._conn_lock.acquire() try: depth = self._db._tx_state.depth + conn = self._db.conn + if depth == 0 and not conn.in_transaction: + # Take the write lock up front. With a connection per thread, + # a transaction that reads and only later writes can find + # another connection has committed in between, and SQLite + # refuses to fork history: the upgrade fails with + # SQLITE_BUSY_SNAPSHOT. Starting as a writer means "no + # subsequent operations in that transaction will ever fail + # with an SQLITE_BUSY error", and busy_timeout covers the + # wait for the lock itself. + conn.execute("BEGIN IMMEDIATE") name = f"simplevecdb_tx_{depth + 1}" - self._db.conn.execute(f"SAVEPOINT {name}") + conn.execute(f"SAVEPOINT {name}") self._db._tx_state.depth = depth + 1 self._savepoint_name = name self._entered = True @@ -2294,7 +2318,7 @@ def __enter__(self) -> "_DBTransaction": # is dropped if it rolls back. self._index_mark = len(self._db._tx_state.index_ops) except Exception: - self._db._lock.release() + self._db._conn_lock.release() raise return self @@ -2350,14 +2374,26 @@ def __exit__(self, exc_type, exc, tb) -> None: self._db._tx_state.depth = max(0, self._db._tx_state.depth - 1) # Outermost commit: if depth fell to 0, finalize the # implicit Python sqlite3 transaction so changes flush. - if self._db._tx_state.depth == 0 and exc_type is None: + if self._db._tx_state.depth == 0: + # Terminate the explicit BEGIN IMMEDIATE, on both paths. + # Releasing the savepoint does not end the enclosing + # transaction: leaving it open would hold SQLite's write + # lock for the life of the connection, and every other + # connection would fail with "database is locked". try: - self._db.conn.commit() + if exc_type is None: + self._db.conn.commit() + else: + self._db.conn.rollback() except Exception: - _logger.error("outer transaction commit failed", exc_info=True) + _logger.error( + "outer transaction %s failed", + "commit" if exc_type is None else "rollback", + exc_info=True, + ) raise finally: - self._db._lock.release() + self._db._conn_lock.release() def __getitem__(self, name: str) -> "VectorCollection": return self._db.collection(name) @@ -2385,10 +2421,10 @@ def __enter__(self) -> "VectorCollection": # type: ignore[override] class _CollectionTxShim: """Minimal proxy emulating the VectorDB attributes _DBTransaction reads.""" - __slots__ = ("_lock", "_tx_state", "_collection") + __slots__ = ("_conn_lock", "_tx_state", "_collection") def __init__(self, collection: "VectorCollection") -> None: - self._lock = collection._lock + self._conn_lock = collection._conn_lock self._tx_state = collection._tx_state self._collection = collection @@ -2615,6 +2651,9 @@ def __init__( # check-then-insert TOCTOU) and the shared sqlite3.Connection's # Python-level transaction context. Shared with every VectorCollection # and CatalogManager constructed by this VectorDB. + # Structural lock: the _collections cache and delete_collection's + # check-then-drop. Always real — nothing about connection pooling + # makes those safe. self._lock = threading.RLock() # Shared transaction-depth counter. Bumped by VectorDB.transaction() # so all catalogs in this DB suspend per-call commits. @@ -2630,6 +2669,12 @@ def __init__( "Use a file path for encrypted databases." ) self._source: Any = open_source(self.path, encryption_key=encryption_key) + # Engages only for a shared connection. With one per thread there is + # no shared transaction context to guard, and SQLite serializes + # writers itself — blocking in C for busy_timeout rather than + # failing, provided every transaction is actually terminated (see + # _DBTransaction.__exit__). + self._conn_lock = ConnectionLock(self._source.shared) self._encrypted = encryption_key is not None if self._encrypted: _logger.info("Opened encrypted database: %s", self.path) @@ -3014,6 +3059,7 @@ def collection( if cache_key not in self._collections: self._collections[cache_key] = VectorCollection( conn=self._source, + conn_lock=self._conn_lock, db_path=self.path, name=name, distance_strategy=distance_strategy or self.distance_strategy, diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index a4c3fd6..4578dd7 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -17,7 +17,7 @@ from ..constants import SQLITE_MAX_BOUND_PARAMS from ..types import ON_CONFLICT_POLICIES, OnConflict -from .connection import ConnectionSource, as_source +from .connection import ConnectionLock, ConnectionSource, as_source from ..utils import _batched from ..utils import validate_filter, retry_on_lock, normalize_filter, find_duplicates @@ -158,6 +158,14 @@ def __enter__(self): if self._tx.depth == 0: self._conn.__enter__() self._owns_conn = True + if not self._conn.in_transaction: + # Start as a writer. sqlite3 would otherwise open a + # *deferred* transaction, and a block that reads before it + # writes (add_documents checks for colliding ids first) + # then has to upgrade — which fails with + # SQLITE_BUSY_SNAPSHOT if another connection committed in + # between, and no busy_timeout waits that out. + self._conn.execute("BEGIN IMMEDIATE") except BaseException: # __exit__ is not called if __enter__ raises; release the lock # ourselves so a connection-level error cannot leak it. @@ -285,7 +293,7 @@ def __init__( conn: "sqlite3.Connection | ConnectionSource", table_name: str, fts_table_name: str, - lock: threading.RLock | None = None, + lock: "threading.RLock | ConnectionLock | None" = None, tx_state: "_TxState | None" = None, ): # Defense-in-depth: validate table names. After this point every @@ -307,7 +315,12 @@ def __init__( # safe under WAL, but Python's `with conn:` transaction context is not # — two threads entering it simultaneously interleave their writes # under one implicit transaction. The lock prevents that. - self._lock: threading.RLock = lock if lock is not None else threading.RLock() + # A ConnectionLock when the owning VectorDB supplied one: it engages + # only while threads share a connection, which is the only case this + # interleaving can happen in. + self._lock: "threading.RLock | ConnectionLock" = ( + lock if lock is not None else threading.RLock() + ) # Optional shared cross-collection transaction state. When the # state's depth > 0, _writable() suppresses inner conn commits so # the outer SAVEPOINT controls atomicity. diff --git a/src/simplevecdb/engine/connection.py b/src/simplevecdb/engine/connection.py index 6fb2f14..cce28dc 100644 --- a/src/simplevecdb/engine/connection.py +++ b/src/simplevecdb/engine/connection.py @@ -33,11 +33,62 @@ def conn(self) -> sqlite3.Connection: """The connection this thread should use.""" ... + @property + def shared(self) -> bool: + """Whether threads share one connection. + + Decides whether the Python-level connection lock has anything to + protect: a shared connection has one transaction context that all + threads would trample, a pooled one does not. + """ + ... + def close_all(self) -> None: """Close every connection handed out.""" ... +class ConnectionLock: + """RLock that engages only when threads share one connection. + + With a connection per thread there is nothing left for it to guard: + each thread has its own transaction context, and SQLite serializes + writers itself (blocking in C for `busy_timeout` rather than failing). + Holding a process-wide lock across every write — and for the whole + lifetime of every transaction — would serialize threads that the + database is perfectly happy to run concurrently. + + It stays a real lock for a shared connection, where the transaction + context genuinely is shared: in-memory databases, and any caller that + injected its own connection. + """ + + __slots__ = ("_lock", "engaged") + + def __init__(self, engaged: bool) -> None: + self._lock = threading.RLock() + self.engaged = engaged + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + if not self.engaged: + return True + return self._lock.acquire(blocking, timeout) + + def release(self) -> None: + if self.engaged: + self._lock.release() + + def __enter__(self) -> "ConnectionLock": + self.acquire() + return self + + def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.release() + + def __repr__(self) -> str: + return f"ConnectionLock(engaged={self.engaged})" + + def is_in_memory(path: str) -> bool: """Whether `path` names an in-memory database.""" return path == ":memory:" @@ -91,6 +142,11 @@ def __init__(self, conn: sqlite3.Connection) -> None: def conn(self) -> sqlite3.Connection: return self._conn + @property + def shared(self) -> bool: + """One connection for every thread, so the lock must engage.""" + return True + def close_all(self) -> None: self._conn.close() @@ -161,6 +217,11 @@ def _open(self) -> sqlite3.Connection: self._all.append(conn) return conn + @property + def shared(self) -> bool: + """Each thread has its own connection, so the lock can stand down.""" + return False + @property def conn(self) -> sqlite3.Connection: """This thread's connection, opened on first use.""" diff --git a/tests/unit/test_concurrency_hardening.py b/tests/unit/test_concurrency_hardening.py index 2e006ff..66c54b9 100644 --- a/tests/unit/test_concurrency_hardening.py +++ b/tests/unit/test_concurrency_hardening.py @@ -304,6 +304,125 @@ def other(): db.close() +class TestTransactionConcurrency: + """Transactions must not serialize the whole database. + + The DB-wide lock used to be held for a transaction's entire lifetime, + because every thread shared one connection. With a connection per thread + it stands down, and SQLite serializes writers itself — so a reader runs + straight through, and concurrent write transactions queue in C rather + than behind a Python lock. + """ + + def test_a_reader_is_not_blocked_by_an_open_transaction(self, tmp_path): + db = VectorDB(str(tmp_path / "v.db")) + coll = db.collection("docs") + coll.add_texts(["committed"], embeddings=[[1.0, 0.0]]) + + started = threading.Event() + seen: dict[str, object] = {} + hold = 0.4 + + def reader(): + assert started.wait(5) + begun = time.perf_counter() + seen["count"] = coll.count() + seen["waited"] = time.perf_counter() - begun + + thread = threading.Thread(target=reader) + thread.start() + with coll.tx() as tx: + tx.add_texts(["in-tx"], embeddings=[[0.0, 1.0]]) + started.set() + time.sleep(hold) + thread.join(10) + + # Generous margin: the point is "did not wait for the transaction", + # not a precise timing assertion. + assert seen["waited"] < hold / 2 + # And it still saw the committed state, not the open transaction's row. + assert seen["count"] == 1 + db.close() + + def test_concurrent_write_transactions_do_not_raise(self, tmp_path): + """BEGIN IMMEDIATE takes the write lock up front. + + Without it, a transaction that reads before it writes can find + another connection committed in between and fail the upgrade with + SQLITE_BUSY_SNAPSHOT, which no busy_timeout waits out. + """ + db = VectorDB(str(tmp_path / "v.db")) + coll = db.collection("docs") + errors: list[str] = [] + per_thread = 10 + threads = 6 + + def writer(n: int): + try: + for i in range(per_thread): + with coll.tx() as tx: + tx.add_texts([f"w{n}-{i}"], embeddings=[[0.5, 0.5]]) + except Exception as exc: # noqa: BLE001 - recorded for the assert + errors.append(repr(exc)) + + workers = [threading.Thread(target=writer, args=(n,)) for n in range(threads)] + for worker in workers: + worker.start() + for worker in workers: + worker.join(120) + + assert not errors + assert coll.count() == threads * per_thread + assert coll._index.size == threads * per_thread + db.close() + + def test_rolled_back_transactions_do_not_hold_the_write_lock(self, tmp_path): + """A rolled-back transaction must end its enclosing BEGIN IMMEDIATE. + + Releasing the savepoint does not end the transaction the savepoint + sits inside. Leaving it open pins SQLite's write lock for the life of + that connection, and every other connection then fails with + "database is locked" — which only shows up once transactions actually + roll back, not on the happy path. + """ + db = VectorDB(str(tmp_path / "v.db")) + coll = db.collection("docs") + errors: list[str] = [] + + def worker(n: int): + try: + for i in range(10): + try: + with coll.tx() as tx: + tx.add_texts([f"w{n}-{i}"], embeddings=[[0.5, 0.5]]) + if i % 3 == 0: + raise RuntimeError("rollback") + except RuntimeError: + pass + except Exception as exc: # noqa: BLE001 - recorded for the assert + errors.append(repr(exc)) + + workers = [threading.Thread(target=worker, args=(n,)) for n in range(6)] + for w in workers: + w.start() + for w in workers: + w.join(120) + + assert not errors + assert coll.count() == coll._index.size + db.close() + + def test_connection_lock_engages_only_for_a_shared_connection(self, tmp_path): + file_db = VectorDB(str(tmp_path / "v.db")) + memory_db = VectorDB(":memory:") + try: + assert not file_db._conn_lock.engaged # pooled + assert memory_db._conn_lock.engaged # shared + finally: + file_db.close() + memory_db.close() + + class TestRetryInsideTransaction: """`@retry_on_lock` must not re-run a body inside a caller's transaction. diff --git a/tests/unit/test_v26_review_pass_3.py b/tests/unit/test_v26_review_pass_3.py index 4398ff0..3d9b18e 100644 --- a/tests/unit/test_v26_review_pass_3.py +++ b/tests/unit/test_v26_review_pass_3.py @@ -87,16 +87,24 @@ def boom(path, *a, **kw): class TestSharedRLock: def test_vectordb_lock_is_shared_with_catalog(self, tmp_path): - """The VectorDB-level RLock must be the same object as every - CatalogManager._lock so transactions on the shared connection do - not interleave between collections.""" + """Every CatalogManager must hold the *same* connection lock object. + + The lock guards the Python-level transaction context, so a per-catalog + lock would let collections interleave on a shared connection. It is + now `db._conn_lock` rather than `db._lock`: the latter guards + structural state (the collections cache, the rebuild_index swap) and + is always real, while this one stands down when each thread has its + own connection and there is no shared context to protect. + """ db = VectorDB(str(tmp_path / "shared.db")) col_a = db.collection("alpha") col_b = db.collection("beta") try: - assert col_a._catalog._lock is db._lock - assert col_b._catalog._lock is db._lock + assert col_a._catalog._lock is db._conn_lock + assert col_b._catalog._lock is db._conn_lock assert col_a._catalog._lock is col_b._catalog._lock + # File-backed: pooled connections, so the lock stands down. + assert not db._conn_lock.engaged finally: db.close() From 51c2ef9c9658825f1e09a74fd34ea47bfbc31b01 Mon Sep 17 00:00:00 2001 From: CoderDayton Date: Wed, 29 Jul 2026 09:44:21 -0500 Subject: [PATCH 6/6] tighten comments and docs --- README.md | 4 +- src/simplevecdb/async_core.py | 70 +++++------- src/simplevecdb/core.py | 140 ++++++++++-------------- src/simplevecdb/engine/catalog.py | 65 +++++------ src/simplevecdb/engine/connection.py | 92 ++++++---------- src/simplevecdb/engine/quantization.py | 5 +- src/simplevecdb/engine/usearch_index.py | 2 - src/simplevecdb/utils.py | 14 +-- tests/unit/test_async_parity.py | 7 +- tests/unit/test_v26_review_pass_3.py | 4 +- 10 files changed, 160 insertions(+), 243 deletions(-) diff --git a/README.md b/README.md index d3a7511..cbbb4ad 100755 --- a/README.md +++ b/README.md @@ -199,7 +199,7 @@ End-to-end notebooks (including a fully local Ollama RAG) live in the ## Feature Highlights A few of the things SimpleVecDB does well — see -[`docs/Features.md`](docs/Features.md) for the comprehensive list. +[`docs/Features.md`](docs/Features.md) for the full list. - **Vector + keyword + hybrid search** — cosine / L2 similarity, BM25 via SQLite FTS5, and Reciprocal Rank Fusion in one collection. @@ -247,7 +247,7 @@ or the [API reference](https://coderdayton.github.io/SimpleVecDB/api/core). ## Documentation -- **[Features](docs/Features.md)** — Comprehensive list of every capability, grouped by area +- **[Features](docs/Features.md)** — Every capability, grouped by area - **[Setup Guide](https://coderdayton.github.io/SimpleVecDB/ENV_SETUP)** — Environment variables, server configuration, authentication - **[API Reference](https://coderdayton.github.io/SimpleVecDB/api/core)** — Complete class/method documentation with type signatures - **[Benchmarks](https://coderdayton.github.io/SimpleVecDB/benchmarks)** — Quantization strategies, batch sizes, hardware optimization diff --git a/src/simplevecdb/async_core.py b/src/simplevecdb/async_core.py index 2236b7d..5946775 100755 --- a/src/simplevecdb/async_core.py +++ b/src/simplevecdb/async_core.py @@ -44,11 +44,8 @@ class _AsyncNamespace: """Awaitable mirror of a sync sub-namespace (`collection.edges`, …). Every public callable on the wrapped namespace is re-exposed as a - coroutine that runs the sync call in the executor. Async code therefore - reads exactly like sync code with `await` in front — - `await coll.edges.upsert(...)` against `coll.edges.upsert(...)` — and a - method added to a sync namespace is reachable from async immediately, - with no wrapper to write and no way for the two surfaces to drift. + coroutine that runs the sync call in the executor, so a method added to a + sync namespace is reachable from async with no wrapper to write. Non-callable attributes pass through unchanged. """ @@ -150,8 +147,8 @@ def name(self) -> str: def conn(self) -> Any: """This thread's SQLite connection on the underlying collection. - A handle, not an operation, so it is not awaitable. Note that reads - issued directly on it bypass the collection API entirely. + A handle, not an operation, so it is not awaitable. Reads issued + directly on it bypass the collection API. """ return self._collection.conn @@ -241,25 +238,24 @@ async def tx(self) -> AsyncIterator["AsyncVectorCollection"]: await coll.delete_by_ids([1]) await coll.add_texts(["replacement"], embeddings=[vec]) - A transaction holds a `threading.RLock` for its lifetime, so its - enter and exit must happen on one thread. The shared pool cannot - promise that — two executor tasks may land on different workers, and - releasing an RLock from a thread that never acquired it raises - `RuntimeError: cannot release un-acquired lock`. So the transaction - gets a private single-worker executor: one thread, every step on it. + A transaction holds a `threading.RLock` for its lifetime, so its enter + and exit must happen on one thread; two tasks on the shared pool may + land on different workers, and releasing an RLock from a thread that + never acquired it raises `RuntimeError: cannot release un-acquired + lock`. Each transaction therefore gets a private single-worker + executor. **Operate through the yielded handle.** It is bound to the pinned - thread; the outer collection is not. Awaiting work on the outer - handle inside the block sends it to the shared pool, where it blocks - on the DB lock this transaction holds — and that lock is only - released when the block exits, which cannot happen while it is - awaiting. Use `atomic()` if you want that mistake to be - unrepresentable: its body is synchronous and cannot await at all. + thread; the outer collection is not. Awaiting work on the outer handle + inside the block sends it to the shared pool, where it blocks on the DB + lock this transaction holds until the block exits — which it cannot do + while it is awaiting. `atomic()` takes a synchronous callback and so + cannot express that. """ loop = asyncio.get_running_loop() - # A nested tx must run on the thread that already holds the lock: the - # RLock is reentrant per thread, so pinning a second thread here - # would block forever waiting on the outer transaction. + # A nested tx must run on the thread that already holds the lock. The + # RLock is reentrant per thread, so a second pinned thread would block + # forever on the outer transaction. reuse = self._tx_pinned pinned = ( self._executor @@ -279,13 +275,11 @@ def _on_pinned(fn: Callable[..., Any], *args: Any) -> Any: try: yield scoped except (GeneratorExit, asyncio.CancelledError) as exc: - # Teardown paths where awaiting is unsafe. Suspending while a + # Teardown paths where awaiting is unsafe: suspending while a # GeneratorExit is in flight raises "async generator ignored # GeneratorExit", and a cancelled task's next await can be - # cancelled again — either way the savepoint would stay open - # and the DB lock would never be released, wedging every other - # writer. Drive the exit on the pinned thread without - # suspending. + # cancelled again. Either leaves the savepoint open and the DB + # lock held, so drive the exit without suspending. pinned.submit( manager.__exit__, type(exc), exc, exc.__traceback__ ).result() @@ -296,10 +290,9 @@ def _on_pinned(fn: Callable[..., Any], *args: Any) -> Any: await _on_pinned(manager.__exit__, None, None, None) finally: if not reuse: - # Non-blocking: the transaction is over, so waiting here would - # only stall the event loop on work the body queued and never - # awaited — work that must not run now that the savepoint has - # closed. cancel_futures drops exactly that. + # Non-blocking: the transaction is over, so anything still + # queued is work the body never awaited and must not run now + # that the savepoint has closed. cancel_futures drops it. pinned.shutdown(wait=False, cancel_futures=True) async def atomic(self, fn: Callable[[VectorCollection], T]) -> T: @@ -309,22 +302,15 @@ async def atomic(self, fn: Callable[[VectorCollection], T]) -> T: `VectorCollection`; everything it does — catalog writes and vector writes alike — commits or rolls back as one unit. - async def swap(coll): + def swap(coll): coll.delete_by_ids([1]) coll.add_texts(["replacement"], embeddings=[vec]) await collection.atomic(swap) - This is a callback rather than `async with collection.tx()` on - purpose. The transaction holds a `threading.RLock` for its whole - lifetime, and an `async with` would enter and exit in two separate - executor tasks: the pool is free to run them on different threads, - and releasing an RLock from a thread that did not acquire it raises - `RuntimeError: cannot release un-acquired lock`. Landing on the same - thread by luck would be no better — every other call on this - collection would sit on that lock across each `await` in the body, - which starves a pool this small. Running the whole body in one - executor task keeps acquire and release paired on one thread. + Running the whole body in one executor task keeps the transaction's + `threading.RLock` acquired and released on one thread, and holds it + across no `await`. `tx()` is the context-manager form. Because `fn` runs off the event loop, it must not await; use the sync collection API inside it. diff --git a/src/simplevecdb/core.py b/src/simplevecdb/core.py index a92cada..1bd6b25 100755 --- a/src/simplevecdb/core.py +++ b/src/simplevecdb/core.py @@ -201,14 +201,11 @@ def __init__( self._quantizer = QuantizationStrategy(quantization) self._encryption_key = encryption_key self._store_embeddings = store_embeddings - # Connection-level lock shared with the parent VectorDB so all - # collections sharing the same sqlite3.Connection serialize their - # transactional access from Python. # Structural lock: guards the in-memory index swap in rebuild_index, - # which is shared across threads no matter how connections are opened. + # which is shared across threads however connections are opened. self._lock: threading.RLock = lock if lock is not None else threading.RLock() - # Connection lock: guards the Python-level transaction context, and so - # engages only when threads share one connection. + # Connection lock: guards the transaction context, so it engages only + # when threads share one connection. self._conn_lock: Any = ( conn_lock if conn_lock is not None else ConnectionLock(self._source.shared) ) @@ -314,10 +311,9 @@ def reserve_ids(self, count: int) -> list[int]: """Reserve `count` document ids without writing any rows. The ids are burned out of the auto-increment sequence, so no later - insert can be handed one of them. Use this when a document's own id - has to appear in its metadata — a self-reference, or a shared group - key across a batch — so the whole group goes in with a single - `add_texts` call instead of an insert followed by a patch-up write. + insert can be handed one of them. Use this when a document's own id has + to appear in its metadata (a self-reference, or a group key shared + across a batch) so the whole group goes in with one `add_texts` call. ids = collection.reserve_ids(3) group = {"group_id": ids[0]} @@ -346,18 +342,16 @@ def _index_add( ) -> None: """Add vectors to the HNSW index, deferring the write inside a tx. - A SQLite SAVEPOINT cannot roll back usearch, so when a transaction - is open the write is buffered on the shared transaction state and - applied just before the outermost savepoint releases. A rollback - discards the buffer, so the catalog and the index stay in step. + A SQLite SAVEPOINT cannot roll back usearch, so when a transaction is + open the write is buffered and applied just before the outermost + savepoint releases; a rollback discards the buffer. - The buffered arrays are copied: they may alias caller-owned memory - that is free to change before the transaction commits. + The buffered arrays are copied: they may alias caller-owned memory that + is free to change before the transaction commits. - Only the thread that owns the transaction defers. Another thread's - write is not part of it — its rows are already committed — so - buffering would hand its vectors to a transaction that can roll - back and drop them. + Only the thread that owns the transaction defers. Another thread's rows + are already committed, so buffering its vectors would expose them to a + rollback that has nothing to do with them. """ if self._tx_state.owned_by_current_thread(): keys_buf = np.array(keys, dtype=np.uint64, copy=True) @@ -445,11 +439,9 @@ def add_texts( f"parent_ids length ({len(parent_ids)}) must match texts length ({len(texts)})" ) - # Check explicit ids once, for the whole call. The catalog re-checks - # per batch, but by then an earlier batch may already have committed — - # only an up-front check can promise that a rejected call wrote - # nothing. Runs before embedding resolution so a doomed call does not - # pay for embeddings first. + # Check explicit ids once for the whole call: the catalog re-checks per + # batch, but by then an earlier batch may already have committed. Runs + # before embedding resolution so a rejected call pays for nothing. explicit_ids = [i for i in ids if i is not None] if ids is not None else [] if explicit_ids: duplicates = find_duplicates(explicit_ids) @@ -518,11 +510,10 @@ def add_texts( ) # Add to usearch index. Inside a transaction this is buffered and - # applied with the commit, so a rollback drops both stores - # together. Outside one the catalog rows are already committed, so - # a failure here diverges the two (rows present, vectors missing); - # log it rather than fail silently. Recovery is rebuild_index() - # (needs store_embeddings=True). + # applied on commit, so a rollback drops both stores. Outside one + # the catalog rows are already committed, so a failure here leaves + # rows without vectors; rebuild_index() recovers that when + # store_embeddings=True. try: self._index_add( np.asarray(doc_ids, dtype=np.uint64), emb_np, threads=threads @@ -2275,12 +2266,11 @@ class _DBTransaction: owns skips its per-call commit. On success the SAVEPOINT is released; on exception it's rolled back and the depth is reset. - Vector writes are held too: collections buffer their HNSW mutations - on the shared transaction state (see `VectorCollection._index_add`) - and the outermost transaction applies them just before releasing, so - a rollback leaves neither store changed. The cost of deferring is - that a search inside the transaction cannot see the transaction's - own vector writes. + Vector writes are held too: collections buffer their HNSW mutations on + the transaction state (see `VectorCollection._index_add`) and the + outermost transaction applies them just before releasing, so a rollback + leaves neither store changed. A search inside the transaction does not + see the transaction's own vector writes. Yields a mapping-like object so callers can do `tx["collection_name"]` to operate on individual collections. @@ -2300,14 +2290,10 @@ def __enter__(self) -> "_DBTransaction": depth = self._db._tx_state.depth conn = self._db.conn if depth == 0 and not conn.in_transaction: - # Take the write lock up front. With a connection per thread, - # a transaction that reads and only later writes can find - # another connection has committed in between, and SQLite - # refuses to fork history: the upgrade fails with - # SQLITE_BUSY_SNAPSHOT. Starting as a writer means "no - # subsequent operations in that transaction will ever fail - # with an SQLITE_BUSY error", and busy_timeout covers the - # wait for the lock itself. + # Start as a writer. A transaction that reads and only later + # writes can find another connection committed in between, and + # the upgrade then fails with SQLITE_BUSY_SNAPSHOT. busy_timeout + # covers the wait for the write lock itself. conn.execute("BEGIN IMMEDIATE") name = f"simplevecdb_tx_{depth + 1}" conn.execute(f"SAVEPOINT {name}") @@ -2325,11 +2311,10 @@ def __enter__(self) -> "_DBTransaction": def _apply_index_ops(self, savepoint: str) -> None: """Flush buffered HNSW writes before the outermost savepoint releases. - Running ahead of the RELEASE is what makes the two stores atomic in - the normal case: if a vector write raises, the catalog side can still - be rolled back. usearch has no undo, so a failure partway through the - buffer can still leave vectors applied for rows that are about to - disappear — that case is reported loudly rather than papered over. + Running ahead of the RELEASE leaves the catalog side still rollable + when a vector write raises. usearch has no undo, so a failure partway + through the buffer leaves the already-applied vectors in place for rows + that are about to disappear; that is logged as an error. """ ops = self._db._tx_state.index_ops if not ops: @@ -2372,14 +2357,11 @@ def __exit__(self, exc_type, exc, tb) -> None: self._db.conn.execute(f"RELEASE SAVEPOINT {name}") finally: self._db._tx_state.depth = max(0, self._db._tx_state.depth - 1) - # Outermost commit: if depth fell to 0, finalize the - # implicit Python sqlite3 transaction so changes flush. if self._db._tx_state.depth == 0: - # Terminate the explicit BEGIN IMMEDIATE, on both paths. + # Terminate the explicit BEGIN IMMEDIATE on both paths. # Releasing the savepoint does not end the enclosing - # transaction: leaving it open would hold SQLite's write - # lock for the life of the connection, and every other - # connection would fail with "database is locked". + # transaction, and leaving it open holds SQLite's write + # lock for the life of the connection. try: if exc_type is None: self._db.conn.commit() @@ -2647,33 +2629,26 @@ def __init__( self.quantization = quantization self._encryption_key = encryption_key self._collections: dict[tuple, VectorCollection] = {} - # Single RLock serializing both the _collections cache (avoid - # check-then-insert TOCTOU) and the shared sqlite3.Connection's - # Python-level transaction context. Shared with every VectorCollection - # and CatalogManager constructed by this VectorDB. - # Structural lock: the _collections cache and delete_collection's - # check-then-drop. Always real — nothing about connection pooling - # makes those safe. + # Structural lock: the _collections cache (check-then-insert TOCTOU) + # and delete_collection's check-then-drop. Shared with every + # VectorCollection and CatalogManager constructed by this VectorDB. self._lock = threading.RLock() # Shared transaction-depth counter. Bumped by VectorDB.transaction() # so all catalogs in this DB suspend per-call commits. self._tx_state: _TxState = _TxState() - # One connection per thread, all against this database. A single - # shared connection cannot isolate transactions — SQLite provides no - # isolation between operations on one connection, so a reader would - # see another thread's uncommitted rows. + # One connection per thread, all against this database: SQLite gives no + # isolation between operations on one connection, so a reader sharing + # one sees another thread's uncommitted rows. if encryption_key is not None and self.path == ":memory:": raise ValueError( "In-memory databases cannot be encrypted. " "Use a file path for encrypted databases." ) self._source: Any = open_source(self.path, encryption_key=encryption_key) - # Engages only for a shared connection. With one per thread there is - # no shared transaction context to guard, and SQLite serializes - # writers itself — blocking in C for busy_timeout rather than - # failing, provided every transaction is actually terminated (see - # _DBTransaction.__exit__). + # Engages only for a shared connection. With one per thread there is no + # shared transaction context to guard, and SQLite serializes writers + # itself provided every transaction is terminated (_DBTransaction). self._conn_lock = ConnectionLock(self._source.shared) self._encrypted = encryption_key is not None if self._encrypted: @@ -2690,10 +2665,9 @@ def __init__( def conn(self) -> sqlite3.Connection: """This thread's SQLite connection, opened on first use. - Each thread gets its own, so one thread's open transaction is - invisible to another's reads — a single shared connection cannot - provide that, because SQLite has no isolation between operations on - one connection. + Each thread gets its own, so one thread's open transaction is invisible + to another's reads. In-memory databases share a single connection and + do not get that isolation. """ return self._source.conn @@ -2703,7 +2677,7 @@ def conn(self, value: Any) -> None: self._source = as_source(value) def transaction(self) -> "_DBTransaction": - """Atomic write context spanning all collections (gap 2). + """Atomic write context spanning all collections. Wraps the work in a single SQLite SAVEPOINT and bumps the shared transaction-depth counter so every catalog method skips its @@ -2832,9 +2806,9 @@ def delete_collection(self, name: str) -> None: # Close any cached collection's open index before removing the file for cached_key, cached_col in list(self._collections.items()): if cached_key[0] == name: - # Stop the TTL sweeper first: left running it keeps - # querying tables this method is about to drop, logging a - # failure every interval for the life of the process. + # Stop the TTL sweeper first: left running it queries the + # tables this method is about to drop, logging a failure + # every interval. ttl_ns = cached_col.__dict__.get("_ttl_ns") if ttl_ns is not None: try: @@ -3123,11 +3097,9 @@ def close(self) -> None: if getattr(self, "_closed", False): return self._closed = True - # Stop opt-in TTL sweepers first. They are daemon threads holding a - # reference to this connection: left running, they wake on their - # interval and fail against a closed database, logging a warning - # every cycle until the process exits — and a sweep in flight can - # race the conn.close() below. + # Stop opt-in TTL sweepers first: they are daemon threads that would + # otherwise wake against a closed database every interval, and a sweep + # in flight can race the conn.close() below. for col in self._collections.values(): ttl_ns = col.__dict__.get("_ttl_ns") if ttl_ns is not None: diff --git a/src/simplevecdb/engine/catalog.py b/src/simplevecdb/engine/catalog.py index 4578dd7..7bb44a6 100755 --- a/src/simplevecdb/engine/catalog.py +++ b/src/simplevecdb/engine/catalog.py @@ -82,23 +82,21 @@ def _check_finite_edge_field(value: Any, field: str) -> None: class _TxState: - """Shared per-VectorDB transaction state (gap 2). + """Shared per-VectorDB transaction state. Used by VectorDB.transaction() to mark all collections/catalogs as operating inside an outer SAVEPOINT. Catalog write helpers consult `depth` to decide whether to commit on exit. - `index_ops` buffers HNSW mutations deferred by the collections taking - part in the transaction. SQLite SAVEPOINTs cannot roll back the - usearch index, so vector writes are held here and applied only when - the outermost transaction is about to release; a rollback truncates - the buffer instead, leaving the two stores in step. - - The object is shared by every catalog in a database, but everything it - stores is per thread: each thread owns a separate SQLite connection and - therefore a separate transaction. That also means a writer on another - thread reads `depth == 0` and applies its vectors immediately, rather - than buffering them into a transaction that may roll back and drop them. + `index_ops` buffers HNSW mutations deferred by the collections taking part + in the transaction. SQLite SAVEPOINTs cannot roll back the usearch index, + so vector writes are held here and applied only when the outermost + transaction is about to release; a rollback truncates the buffer. + + The object is shared by every catalog in a database, but its contents are + per thread: each thread owns a separate connection and so a separate + transaction. A writer on another thread therefore reads `depth == 0` and + applies its vectors immediately. """ __slots__ = ("_local",) @@ -159,12 +157,11 @@ def __enter__(self): self._conn.__enter__() self._owns_conn = True if not self._conn.in_transaction: - # Start as a writer. sqlite3 would otherwise open a - # *deferred* transaction, and a block that reads before it - # writes (add_documents checks for colliding ids first) - # then has to upgrade — which fails with - # SQLITE_BUSY_SNAPSHOT if another connection committed in - # between, and no busy_timeout waits that out. + # Start as a writer. sqlite3 opens a deferred transaction, + # so a block that reads before it writes (add_documents + # checks for colliding ids first) has to upgrade, which + # fails with SQLITE_BUSY_SNAPSHOT if another connection + # committed in between; busy_timeout does not wait it out. self._conn.execute("BEGIN IMMEDIATE") except BaseException: # __exit__ is not called if __enter__ raises; release the lock @@ -312,12 +309,11 @@ def __init__( self._cluster_table_ready = False # Serializes Python-level access to the shared sqlite3.Connection. The # connection is opened with check_same_thread=False; SQLite itself is - # safe under WAL, but Python's `with conn:` transaction context is not - # — two threads entering it simultaneously interleave their writes - # under one implicit transaction. The lock prevents that. - # A ConnectionLock when the owning VectorDB supplied one: it engages - # only while threads share a connection, which is the only case this - # interleaving can happen in. + # safe under WAL, but Python's `with conn:` transaction context is not: + # two threads entering it simultaneously interleave their writes under + # one implicit transaction. The VectorDB supplies a ConnectionLock, + # which engages only while threads share a connection — the only case + # that interleaving can happen in. self._lock: "threading.RLock | ConnectionLock" = ( lock if lock is not None else threading.RLock() ) @@ -342,9 +338,8 @@ def _writable(self): def conn(self) -> "sqlite3.Connection": """This thread's SQLite connection. - A property rather than a stored handle: each thread owns a separate - connection, which is what keeps one thread's open transaction - invisible to another's reads. + A property rather than a stored handle, because each thread owns a + separate connection. """ return self._source.conn @@ -695,9 +690,8 @@ def add_documents( real_ids: list[int] = [-1] * len(ids_list) if explicit_rows: - # A repeat inside one call is always an error: the two rows would - # race for the same id whatever the policy, and under "replace" - # the loser would be silently discarded. + # A repeat inside one call is an error under any policy: the two + # rows target the same id, and "replace" would discard one of them. explicit_ids = [int(r[0]) for r in explicit_rows] duplicates = find_duplicates(explicit_ids) if duplicates: @@ -802,14 +796,13 @@ def existing_ids(self, ids: Sequence[int]) -> list[int]: def reserve_ids(self, count: int) -> list[int]: """Reserve `count` document ids without inserting any rows. - Advances the table's AUTOINCREMENT high-water mark so the returned - ids can never be handed out again by a later auto-id insert. This - lets a caller stamp ids into metadata (self-referential rows, a - shared group key) and then write the whole group in one - `add_texts` call, instead of inserting and patching afterwards. + Advances the table's AUTOINCREMENT high-water mark so the returned ids + can never be handed out again by a later auto-id insert. Lets a caller + stamp ids into metadata (self-referential rows, a shared group key) and + write the whole group in one `add_texts` call. Reserved ids are not rows: nothing is stored until they are passed - back as `ids=`. Ids that are never used simply stay unallocated. + back as `ids=`. Ids that are never used stay unallocated. Args: count: How many ids to reserve. Must be positive. diff --git a/src/simplevecdb/engine/connection.py b/src/simplevecdb/engine/connection.py index cce28dc..b40ea0a 100644 --- a/src/simplevecdb/engine/connection.py +++ b/src/simplevecdb/engine/connection.py @@ -1,15 +1,10 @@ """Per-thread SQLite connections for SimpleVecDB. -SQLite gives no isolation between operations on a *single* connection: a -reader sharing one connection with a writer sees that writer's uncommitted -rows, and can act on data that later rolls back. Isolation is a property of -having separate connections — with them, "the reader is only able to see -complete committed transactions from the writer... regardless of whether the -two database connections are in the same thread, in different threads of the -same process, or in different processes." - -So each thread gets its own connection to the same database. In WAL mode that -also buys snapshot isolation and lets readers run while a writer commits. +SQLite gives no isolation between operations on a single connection: a reader +sharing one connection with a writer sees that writer's uncommitted rows. +Separate connections are isolated from each other, so each thread gets its own +connection to the same database. In WAL mode readers also run concurrently +with a writer. """ from __future__ import annotations @@ -35,12 +30,7 @@ def conn(self) -> sqlite3.Connection: @property def shared(self) -> bool: - """Whether threads share one connection. - - Decides whether the Python-level connection lock has anything to - protect: a shared connection has one transaction context that all - threads would trample, a pooled one does not. - """ + """Whether threads share one connection, and so one transaction context.""" ... def close_all(self) -> None: @@ -51,16 +41,10 @@ def close_all(self) -> None: class ConnectionLock: """RLock that engages only when threads share one connection. - With a connection per thread there is nothing left for it to guard: - each thread has its own transaction context, and SQLite serializes - writers itself (blocking in C for `busy_timeout` rather than failing). - Holding a process-wide lock across every write — and for the whole - lifetime of every transaction — would serialize threads that the - database is perfectly happy to run concurrently. - - It stays a real lock for a shared connection, where the transaction - context genuinely is shared: in-memory databases, and any caller that - injected its own connection. + A shared connection has one transaction context for every thread, so the + lock is real for in-memory databases and injected connections. With a + connection per thread there is no shared context to guard and SQLite + serializes writers itself, so acquire and release become no-ops. """ __slots__ = ("_lock", "engaged") @@ -97,14 +81,10 @@ def is_in_memory(path: str) -> bool: def shared_memory_dsn() -> str: """A named shared-cache DSN for an in-memory database. - Not used by default, and the reason is worth recording. Pooling an - in-memory database requires this form — *"opening two database - connections each with the filename ':memory:' will create two - independent in-memory databases"* — but shared cache takes **table-level** - write locks, and a reader on another connection then fails with - ``SQLITE_LOCKED`` ("database table is locked"), which `busy_timeout` does - not wait out. That trades a harmless stale read for a hard error, so - in-memory databases keep a single shared connection instead. + Two connections opened on `":memory:"` get two independent databases, so + pooling one requires this form. Unused by default: shared cache takes + table-level write locks, and a reader on another connection then fails + with ``SQLITE_LOCKED``, which `busy_timeout` does not wait out. """ return f"file:simplevecdb_{uuid.uuid4().hex}?mode=memory&cache=shared" @@ -112,10 +92,9 @@ def shared_memory_dsn() -> str: def apply_pragmas(conn: sqlite3.Connection) -> None: """Configure a freshly opened connection. - Every one of these is connection-scoped, not database-scoped, so a pool - must apply them per connection — `foreign_keys` silently defaults back to - off otherwise. (`journal_mode=WAL` is the exception, persisting in the - file, but setting it again is harmless and keeps this in one place.) + These are connection-scoped rather than database-scoped, so a pool must + apply them per connection; `foreign_keys` silently defaults back to off + otherwise. `journal_mode=WAL` is the exception and persists in the file. """ conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") @@ -128,9 +107,8 @@ def apply_pragmas(conn: sqlite3.Connection) -> None: class SingleConnection: """Adapter presenting one already-open connection as a connection source. - Used when a caller supplies its own `sqlite3.Connection` directly instead - of a path, so those call sites keep working unchanged. It cannot provide - cross-thread isolation — that is the point of `ConnectionPool`. + Used when a caller supplies its own `sqlite3.Connection` instead of a + path. It provides no cross-thread isolation. """ __slots__ = ("_conn",) @@ -154,10 +132,9 @@ def close_all(self) -> None: class ConnectionPool: """One SQLite connection per thread, all against the same database. - The first connection is opened eagerly and kept for the pool's lifetime: - a shared-cache in-memory database is reclaimed when its *last* connection - closes, so without an anchor the data would vanish whenever a worker - thread's connection happened to be the last one out. + The constructing thread's connection is opened eagerly, so an unusable + path or key fails at construction rather than on first use, and is held + for the pool's lifetime. """ __slots__ = ( @@ -236,10 +213,9 @@ def conn(self) -> sqlite3.Connection: def close_all(self) -> None: """Close every connection this pool opened, from any thread. - Connections belonging to other threads are closed here too. That is - safe because they were opened with `check_same_thread=False`, and it - is necessary because a worker thread may never run again to close its - own. + Other threads' connections are closed here too; they were opened with + `check_same_thread=False`, and a worker thread may never run again to + close its own. """ self._closed = True with self._all_lock: @@ -263,13 +239,11 @@ def open_source( ) -> ConnectionSource: """Open the right kind of connection source for `path`. - File-backed databases get one connection per thread, which is what makes - a transaction on one thread invisible to reads on another. In-memory - databases get a single shared connection: pooling one requires shared - cache, whose table-level locks turn concurrent readers into - ``SQLITE_LOCKED`` errors (see `shared_memory_dsn`). They therefore keep - the old behaviour, cross-thread dirty reads included — acceptable for a - database that cannot outlive the process. + File-backed databases get one connection per thread, so a transaction on + one thread is invisible to reads on another. In-memory databases get a + single shared connection and so keep cross-thread dirty reads: pooling one + requires shared cache, whose table-level locks turn concurrent readers + into ``SQLITE_LOCKED`` errors (see `shared_memory_dsn`). """ if is_in_memory(path): conn = sqlite3.connect(path, check_same_thread=False, timeout=timeout) @@ -281,9 +255,9 @@ def open_source( def as_source(conn_or_source: Any) -> ConnectionSource: """Accept either a raw connection or something already pool-shaped. - Tested by identity rather than by `isinstance(..., sqlite3.Connection)`: - sqlcipher connections and test doubles are connection-like without being - instances of it, and wrapping those is exactly the intent. + Anything that is not already a source is treated as a raw connection. + Checked this way round because sqlcipher connections and test doubles are + connection-like without being `sqlite3.Connection` instances. """ if isinstance(conn_or_source, (ConnectionPool, SingleConnection)): return conn_or_source diff --git a/src/simplevecdb/engine/quantization.py b/src/simplevecdb/engine/quantization.py index f453994..fe924cb 100755 --- a/src/simplevecdb/engine/quantization.py +++ b/src/simplevecdb/engine/quantization.py @@ -110,9 +110,8 @@ def deserialize(self, blob: bytes, dim: int | None) -> np.ndarray: """ if self.quantization == Quantization.FLOAT: # np.frombuffer aliases the blob and is read-only. Every other - # branch below returns a fresh writable array (astype/where copy), - # so copy here too rather than hand back a vector whose - # mutability depends on the quantization mode. + # branch returns a fresh writable array (astype/where copy), so + # copy here too and keep the return type consistent. return np.frombuffer(blob, dtype=np.float32).copy() elif self.quantization == Quantization.INT8: diff --git a/src/simplevecdb/engine/usearch_index.py b/src/simplevecdb/engine/usearch_index.py index ef2f659..2a2cdcf 100755 --- a/src/simplevecdb/engine/usearch_index.py +++ b/src/simplevecdb/engine/usearch_index.py @@ -398,8 +398,6 @@ def remove(self, keys: NDArray[np.uint64] | list[int]) -> int: return 0 with self._write_lock: - # Mapped read-only indexes must be reloaded writable first, or the - # remove below segfaults rather than raising. self._ensure_writable("remove") # Re-read after taking the lock: close() also takes it, so the diff --git a/src/simplevecdb/utils.py b/src/simplevecdb/utils.py index 3b80536..2560305 100755 --- a/src/simplevecdb/utils.py +++ b/src/simplevecdb/utils.py @@ -37,8 +37,7 @@ def _batched(iterable: Iterable[Any], n: int) -> Iterable[Sequence[Any]]: def find_duplicates(values: Sequence[int]) -> list[int]: """Return the values appearing more than once, in ascending order. - Single pass over the input — callers hand this whole batches of - caller-supplied document ids, so a quadratic scan is not an option. + Single pass; callers pass whole batches of document ids. """ seen: set[int] = set() repeated: set[int] = set() @@ -113,12 +112,11 @@ def retry_on_lock( def decorator(func: F) -> F: @wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: - # Never retry inside the caller's transaction. There, the write - # helper suppresses the per-call commit, so a statement that - # failed halfway through leaves its earlier statements applied - # and un-rolled-back; re-running the body would insert the - # auto-id rows a second time. The transaction owns atomicity, so - # let the error reach it and be rolled back as a unit. + # Never retry inside the caller's transaction: the write helper + # suppresses the per-call commit there, so a body that failed + # partway leaves its earlier statements applied, and re-running it + # would insert the auto-id rows twice. The transaction owns + # atomicity, so let the error reach it. tx_state = getattr(args[0], "_tx_state", None) if args else None if tx_state is not None and tx_state.owned_by_current_thread(): return func(*args, **kwargs) diff --git a/tests/unit/test_async_parity.py b/tests/unit/test_async_parity.py index 99e2828..1970c04 100644 --- a/tests/unit/test_async_parity.py +++ b/tests/unit/test_async_parity.py @@ -21,11 +21,8 @@ from simplevecdb.async_core import AsyncVectorCollection, AsyncVectorDB from simplevecdb.core import VectorCollection, VectorDB -# Nothing is sync-only any more. `tx()` used to be: a transaction holds a -# threading.RLock for its lifetime, and on the shared pool an `async with` -# could enter and exit on different workers, releasing a lock the thread -# never acquired. Giving each transaction a private single-worker executor -# removes that, so the async surface now mirrors sync completely. +# Nothing is sync-only: every public sync method has an async counterpart. +# Add a name here only with the reason it cannot be mirrored. INTENTIONALLY_SYNC_ONLY: set[str] = set() NAMESPACES = ("edges", "events", "ttl", "pending", "maintenance", "counters") diff --git a/tests/unit/test_v26_review_pass_3.py b/tests/unit/test_v26_review_pass_3.py index 3d9b18e..ea9cddf 100644 --- a/tests/unit/test_v26_review_pass_3.py +++ b/tests/unit/test_v26_review_pass_3.py @@ -184,8 +184,8 @@ def test_filter_does_not_inflate_vector_rrf_score(self, tmp_path): col.add_texts(texts, metadatas=metas, embeddings=embs) # Hybrid search with a filter that drops the top 9 vector hits. - # If rank symmetry is broken, the surviving "drop" → wait, - # all dropped — so we use a more nuanced setup: keep just one. + # Exactly one document survives the filter, so a broken rank + # symmetry shows up as an empty or misordered result. results = col.hybrid_search( query="doc number 0", query_vector=[0.1, 9.9, 0.0, 0.0],