mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 19:29:43 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8225701511 | ||
|
|
93e4c8cc1f |
@@ -57,6 +57,17 @@ MIGRATIONS = [
|
||||
PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
|
||||
);""",
|
||||
"ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
|
||||
"""
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
|
||||
""",
|
||||
"""
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
|
||||
""",
|
||||
]
|
||||
|
||||
SELECT_SQL = f"""
|
||||
|
||||
@@ -6,7 +6,6 @@ from typing import Any, Callable, Optional, Union, cast
|
||||
|
||||
import orjson
|
||||
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
@@ -19,9 +18,7 @@ from langgraph.store.base import (
|
||||
Result,
|
||||
SearchOp,
|
||||
)
|
||||
from langgraph.store.base.batch import (
|
||||
BatchedBaseStore,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
from langgraph.store.postgres.base import (
|
||||
_PLACEHOLDER,
|
||||
BasePostgresStore,
|
||||
@@ -38,7 +35,7 @@ from langgraph.store.postgres.base import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
"""Asynchronous Postgres-backed store with optional vector search using pgvector.
|
||||
|
||||
!!! example "Examples"
|
||||
@@ -159,7 +156,12 @@ class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
return results
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
|
||||
futures = []
|
||||
for op in ops:
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue[fut] = op
|
||||
futures.append(fut)
|
||||
return [fut.result() for fut in asyncio.as_completed(futures)]
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
@@ -221,22 +223,19 @@ class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
|
||||
"""
|
||||
|
||||
async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
|
||||
try:
|
||||
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
|
||||
row = await cur.fetchone()
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
await cur.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
await cur.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
|
||||
row = cast(dict, await cur.fetchone())
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
return version
|
||||
|
||||
async with self._cursor() as cur:
|
||||
|
||||
@@ -21,7 +21,6 @@ from typing import (
|
||||
|
||||
import orjson
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.errors import UndefinedTable
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
@@ -30,6 +29,7 @@ from typing_extensions import TypedDict
|
||||
from langgraph.checkpoint.postgres import _ainternal as _ainternal
|
||||
from langgraph.checkpoint.postgres import _internal as _pg_internal
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
GetOp,
|
||||
IndexConfig,
|
||||
Item,
|
||||
@@ -43,7 +43,6 @@ from langgraph.store.base import (
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
)
|
||||
from langgraph.store.base.batch import SyncBatchedBaseStore
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_core.embeddings import Embeddings
|
||||
@@ -73,7 +72,7 @@ CREATE TABLE IF NOT EXISTS store (
|
||||
""",
|
||||
"""
|
||||
-- For faster lookups by prefix
|
||||
CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
|
||||
""",
|
||||
]
|
||||
|
||||
@@ -107,7 +106,7 @@ CREATE TABLE IF NOT EXISTS store_vectors (
|
||||
),
|
||||
Migration(
|
||||
"""
|
||||
CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
|
||||
USING %(index_type)s (embedding %(ops)s)%(index_params)s;
|
||||
""",
|
||||
condition=lambda store: bool(
|
||||
@@ -533,7 +532,7 @@ class BasePostgresStore(Generic[C]):
|
||||
raise ValueError(f"Unsupported operator: {op}")
|
||||
|
||||
|
||||
class PostgresStore(SyncBatchedBaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
"""Postgres-backed store with optional vector search using pgvector.
|
||||
|
||||
!!! example "Examples"
|
||||
@@ -847,22 +846,19 @@ class PostgresStore(SyncBatchedBaseStore, BasePostgresStore[_pg_internal.Conn]):
|
||||
"""
|
||||
|
||||
def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
|
||||
try:
|
||||
cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
|
||||
row = cast(dict, cur.fetchone())
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
except UndefinedTable:
|
||||
version = -1
|
||||
cur.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
cur.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {table} (
|
||||
v INTEGER PRIMARY KEY
|
||||
)
|
||||
"""
|
||||
)
|
||||
cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
|
||||
row = cast(dict, cur.fetchone())
|
||||
if row is None:
|
||||
version = -1
|
||||
else:
|
||||
version = row["v"]
|
||||
return version
|
||||
|
||||
with self._cursor() as cur:
|
||||
|
||||
Generated
+13
-14
@@ -13,24 +13,24 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "anyio"
|
||||
version = "4.6.2.post1"
|
||||
version = "4.7.0"
|
||||
description = "High level compatibility layer for multiple asynchronous event loop implementations"
|
||||
optional = false
|
||||
python-versions = ">=3.9"
|
||||
files = [
|
||||
{file = "anyio-4.6.2.post1-py3-none-any.whl", hash = "sha256:6d170c36fba3bdd840c73d3868c1e777e33676a69c3a72cf0a0d5d6d8009b61d"},
|
||||
{file = "anyio-4.6.2.post1.tar.gz", hash = "sha256:4c8bc31ccdb51c7f7bd251f51c609e038d63e34219b44aa86e47576389880b4c"},
|
||||
{file = "anyio-4.7.0-py3-none-any.whl", hash = "sha256:ea60c3723ab42ba6fff7e8ccb0488c898ec538ff4df1f1d5e642c3601d07e352"},
|
||||
{file = "anyio-4.7.0.tar.gz", hash = "sha256:2f834749c602966b7d456a7567cafcb309f96482b5081d14ac93ccd457f9dd48"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""}
|
||||
idna = ">=2.8"
|
||||
sniffio = ">=1.1"
|
||||
typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""}
|
||||
typing_extensions = {version = ">=4.5", markers = "python_version < \"3.13\""}
|
||||
|
||||
[package.extras]
|
||||
doc = ["Sphinx (>=7.4,<8.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
|
||||
test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21.0b1)"]
|
||||
doc = ["Sphinx (>=7.4,<8.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx_rtd_theme"]
|
||||
test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21)"]
|
||||
trio = ["trio (>=0.26.1)"]
|
||||
|
||||
[[package]]
|
||||
@@ -244,13 +244,13 @@ trio = ["trio (>=0.22.0,<1.0)"]
|
||||
|
||||
[[package]]
|
||||
name = "httpx"
|
||||
version = "0.27.2"
|
||||
version = "0.28.0"
|
||||
description = "The next generation HTTP client."
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0"},
|
||||
{file = "httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2"},
|
||||
{file = "httpx-0.28.0-py3-none-any.whl", hash = "sha256:dc0b419a0cfeb6e8b34e85167c0da2671206f5095f1baa9663d23bcfd6b535fc"},
|
||||
{file = "httpx-0.28.0.tar.gz", hash = "sha256:0858d3bab51ba7e386637f22a61d8ccddaeec5f3fe4209da3a6168dbb91573e0"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -258,7 +258,6 @@ anyio = "*"
|
||||
certifi = "*"
|
||||
httpcore = "==1.*"
|
||||
idna = "*"
|
||||
sniffio = "*"
|
||||
|
||||
[package.extras]
|
||||
brotli = ["brotli", "brotlicffi"]
|
||||
@@ -342,7 +341,7 @@ typing-extensions = ">=4.7"
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "2.0.7"
|
||||
version = "2.0.8"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
optional = false
|
||||
python-versions = "^3.9.0,<4.0"
|
||||
@@ -741,13 +740,13 @@ typing-extensions = ">=4.6"
|
||||
|
||||
[[package]]
|
||||
name = "pydantic"
|
||||
version = "2.10.2"
|
||||
version = "2.10.3"
|
||||
description = "Data validation using Python type hints"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "pydantic-2.10.2-py3-none-any.whl", hash = "sha256:cfb96e45951117c3024e6b67b25cdc33a3cb7b2fa62e239f7af1378358a1d99e"},
|
||||
{file = "pydantic-2.10.2.tar.gz", hash = "sha256:2bc2d7f17232e0841cbba4641e65ba1eb6fafb3a08de3a091ff3ce14a197c4fa"},
|
||||
{file = "pydantic-2.10.3-py3-none-any.whl", hash = "sha256:be04d85bbc7b65651c5f8e6b9976ed9c6f41782a55524cef079a34a0bb82144d"},
|
||||
{file = "pydantic-2.10.3.tar.gz", hash = "sha256:cb5ac360ce894ceacd69c403187900a02c4b20b693a9dd1d643e1effab9eadf9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "2.0.7"
|
||||
version = "2.0.8"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -63,9 +63,8 @@ async def _pipe_saver():
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
async with conn.pipeline() as pipe:
|
||||
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
|
||||
await checkpointer.setup()
|
||||
checkpointer = AsyncPostgresSaver(conn)
|
||||
await checkpointer.setup()
|
||||
async with conn.pipeline() as pipe:
|
||||
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
|
||||
yield checkpointer
|
||||
|
||||
@@ -65,7 +65,7 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
|
||||
await conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
def test_large_batches(store: AsyncPostgresStore) -> None:
|
||||
async def test_large_batches(store: AsyncPostgresStore) -> None:
|
||||
N = 1000
|
||||
M = 10
|
||||
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
# type: ignore
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Optional
|
||||
from uuid import uuid4
|
||||
@@ -19,7 +17,11 @@ from langgraph.store.base import (
|
||||
SearchOp,
|
||||
)
|
||||
from langgraph.store.postgres import PostgresStore
|
||||
from tests.conftest import DEFAULT_URI, VECTOR_TYPES, CharacterEmbeddings
|
||||
from tests.conftest import (
|
||||
DEFAULT_URI,
|
||||
VECTOR_TYPES,
|
||||
CharacterEmbeddings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
|
||||
@@ -57,96 +59,6 @@ def store(request) -> PostgresStore:
|
||||
conn.execute(f"DROP DATABASE {database}")
|
||||
|
||||
|
||||
def test_large_batches(store: PostgresStore) -> None:
|
||||
N = 1000
|
||||
M = 10
|
||||
|
||||
with ThreadPoolExecutor(max_workers=10) as executor:
|
||||
for m in range(M):
|
||||
for i in range(N):
|
||||
_ = [
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
),
|
||||
executor.submit(
|
||||
store.get,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
),
|
||||
executor.submit(
|
||||
store.list_namespaces,
|
||||
prefix=None,
|
||||
max_depth=m + 1,
|
||||
),
|
||||
executor.submit(
|
||||
store.search,
|
||||
("test",),
|
||||
),
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
),
|
||||
executor.submit(
|
||||
store.put,
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
None,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
async def test_large_batches_async(store: PostgresStore) -> None:
|
||||
N = 1000
|
||||
M = 10
|
||||
coros = []
|
||||
for m in range(M):
|
||||
for i in range(N):
|
||||
coros.append(
|
||||
store.aput(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.aget(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.alist_namespaces(
|
||||
prefix=None,
|
||||
max_depth=m + 1,
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.asearch(
|
||||
("test",),
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.aput(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
value={"foo": "bar" + str(i)},
|
||||
)
|
||||
)
|
||||
coros.append(
|
||||
store.adelete(
|
||||
("test", "foo", "bar", "baz", str(m % 2)),
|
||||
f"key{i}",
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.gather(*coros)
|
||||
|
||||
|
||||
def test_batch_order(store: PostgresStore) -> None:
|
||||
# Setup test data
|
||||
store.put(("test", "foo"), "key1", {"data": "value1"})
|
||||
|
||||
@@ -57,9 +57,8 @@ def _pipe_saver():
|
||||
prepare_threshold=0,
|
||||
row_factory=dict_row,
|
||||
) as conn:
|
||||
with conn.pipeline() as pipe:
|
||||
checkpointer = PostgresSaver(conn, pipe=pipe)
|
||||
checkpointer.setup()
|
||||
checkpointer = PostgresSaver(conn)
|
||||
checkpointer.setup()
|
||||
with conn.pipeline() as pipe:
|
||||
checkpointer = PostgresSaver(conn, pipe=pipe)
|
||||
yield checkpointer
|
||||
|
||||
@@ -808,8 +808,6 @@ class BaseStore(ABC):
|
||||
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
if max_depth is not None and max_depth <= 0:
|
||||
raise ValueError("If provided, max_depth must be greater than 0")
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
@@ -1006,8 +1004,6 @@ class BaseStore(ABC):
|
||||
# Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
if max_depth is not None and max_depth <= 0:
|
||||
raise ValueError("If provided, max_depth must be greater than 0")
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from concurrent.futures import Future
|
||||
from typing import Any, Iterable, Literal, Optional, Union
|
||||
from typing import Any, Literal, Optional, Union
|
||||
|
||||
from langgraph.store.base import (
|
||||
BaseStore,
|
||||
@@ -14,19 +11,24 @@ from langgraph.store.base import (
|
||||
NamespacePath,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchItem,
|
||||
SearchOp,
|
||||
_validate_namespace,
|
||||
)
|
||||
|
||||
|
||||
class AsyncBatchedBaseStoreMixin:
|
||||
class AsyncBatchedBaseStore(BaseStore):
|
||||
"""Efficiently batch operations in a background task."""
|
||||
|
||||
_loop: asyncio.AbstractEventLoop
|
||||
_aqueue: dict[asyncio.Future, Op]
|
||||
_task: asyncio.Task
|
||||
__slots__ = ("_loop", "_aqueue", "_task")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._aqueue: dict[asyncio.Future, Op] = {}
|
||||
self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
|
||||
|
||||
def __del__(self) -> None:
|
||||
self._task.cancel()
|
||||
|
||||
async def aget(
|
||||
self,
|
||||
@@ -98,29 +100,6 @@ class AsyncBatchedBaseStoreMixin:
|
||||
return await fut
|
||||
|
||||
|
||||
class AsyncBatchedBaseStore(AsyncBatchedBaseStoreMixin, BaseStore):
|
||||
"""Efficiently batch operations in a background task."""
|
||||
|
||||
__slots__ = ("_loop", "_aqueue", "_task")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._aqueue: dict[asyncio.Future, Op] = {}
|
||||
self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
|
||||
|
||||
def __del__(self) -> None:
|
||||
self._task.cancel()
|
||||
|
||||
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
futures = []
|
||||
for op in ops:
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue[fut] = op
|
||||
futures.append(fut)
|
||||
return [fut.result() for fut in asyncio.as_completed(futures)]
|
||||
|
||||
|
||||
def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
|
||||
"""Dedupe operations while preserving order for results.
|
||||
|
||||
@@ -195,181 +174,3 @@ async def _run(
|
||||
break
|
||||
# remove strong ref to store
|
||||
del s
|
||||
|
||||
|
||||
class SyncBatchedBaseStoreMixin(BaseStore):
|
||||
"""Efficiently batch operations in a background thread."""
|
||||
|
||||
_sync_queue: dict[Future, Op]
|
||||
_sync_thread: threading.Thread
|
||||
|
||||
def get(
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
) -> Optional[Item]:
|
||||
fut: Future[Optional[Item]] = Future()
|
||||
self._sync_queue[fut] = GetOp(namespace, key)
|
||||
return fut.result()
|
||||
|
||||
def search(
|
||||
self,
|
||||
namespace_prefix: tuple[str, ...],
|
||||
/,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
filter: Optional[dict[str, Any]] = None,
|
||||
limit: int = 10,
|
||||
offset: int = 0,
|
||||
) -> list[SearchItem]:
|
||||
fut: Future[list[SearchItem]] = Future()
|
||||
self._sync_queue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query)
|
||||
return fut.result()
|
||||
|
||||
def put(
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
value: dict[str, Any],
|
||||
index: Optional[Union[Literal[False], list[str]]] = None,
|
||||
) -> None:
|
||||
_validate_namespace(namespace)
|
||||
fut: Future[None] = Future()
|
||||
self._sync_queue[fut] = PutOp(namespace, key, value, index)
|
||||
return fut.result()
|
||||
|
||||
def delete(
|
||||
self,
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
) -> None:
|
||||
fut: Future[None] = Future()
|
||||
self._sync_queue[fut] = PutOp(namespace, key, None)
|
||||
return fut.result()
|
||||
|
||||
def list_namespaces(
|
||||
self,
|
||||
*,
|
||||
prefix: Optional[NamespacePath] = None,
|
||||
suffix: Optional[NamespacePath] = None,
|
||||
max_depth: Optional[int] = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> list[tuple[str, ...]]:
|
||||
fut: Future[list[tuple[str, ...]]] = Future()
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
if suffix:
|
||||
match_conditions.append(MatchCondition(match_type="suffix", path=suffix))
|
||||
|
||||
op = ListNamespacesOp(
|
||||
match_conditions=tuple(match_conditions),
|
||||
max_depth=max_depth,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
self._sync_queue[fut] = op
|
||||
return fut.result()
|
||||
|
||||
|
||||
class SyncBatchedBaseStore(SyncBatchedBaseStoreMixin, BaseStore):
|
||||
"""Efficiently batch operations in a background thread."""
|
||||
|
||||
__slots__ = ("_sync_queue", "_sync_thread")
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._sync_queue: dict[Future, Op] = {}
|
||||
self._sync_thread = threading.Thread(
|
||||
target=_sync_run,
|
||||
args=(self._sync_queue, weakref.ref(self)),
|
||||
daemon=True,
|
||||
)
|
||||
self._sync_thread.start()
|
||||
|
||||
def __del__(self) -> None:
|
||||
# Signal the thread to stop
|
||||
if self._sync_thread.is_alive():
|
||||
empty_future: Future = Future()
|
||||
self._sync_queue[empty_future] = None # type: ignore
|
||||
self._sync_thread.join()
|
||||
|
||||
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
||||
futures = []
|
||||
for op in ops:
|
||||
fut: Future[Result] = Future()
|
||||
self._sync_queue[fut] = op
|
||||
futures.append(fut)
|
||||
return [fut.result() for fut in futures]
|
||||
|
||||
|
||||
class BatchedBaseStore(
|
||||
AsyncBatchedBaseStoreMixin, SyncBatchedBaseStoreMixin, BaseStore
|
||||
):
|
||||
__slots__ = (
|
||||
"_sync_queue",
|
||||
"_sync_thread",
|
||||
"_task",
|
||||
"_loop",
|
||||
"_aqueue",
|
||||
)
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
# Setup async processing
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._aqueue: dict[asyncio.Future, Op] = {}
|
||||
self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
|
||||
|
||||
self._sync_queue: dict[Future, Op] = {}
|
||||
self._sync_thread = threading.Thread(
|
||||
target=_sync_run,
|
||||
args=(self._sync_queue, weakref.ref(self)),
|
||||
daemon=True,
|
||||
)
|
||||
self._sync_thread.start()
|
||||
|
||||
def __del__(self) -> None:
|
||||
# Signal the thread to stop
|
||||
if self._sync_thread.is_alive():
|
||||
empty_future: Future[None] = Future()
|
||||
self._sync_queue[empty_future] = None # type: ignore
|
||||
self._sync_thread.join()
|
||||
|
||||
# Signal the thread to stop
|
||||
if self._task is not None:
|
||||
self._task.cancel()
|
||||
|
||||
|
||||
def _sync_run(queue: dict[Future, Op], store: weakref.ReferenceType[BaseStore]) -> None:
|
||||
while True:
|
||||
time.sleep(0.001) # Yield to other threads
|
||||
if not queue:
|
||||
continue
|
||||
if s := store():
|
||||
# get the operations to run
|
||||
taken = queue.copy()
|
||||
# action each operation
|
||||
try:
|
||||
values = list(taken.values())
|
||||
if None in values: # Exit signal
|
||||
break
|
||||
listen, dedupped = _dedupe_ops(values)
|
||||
results = s.batch(dedupped) # Note: Using sync batch here
|
||||
if listen is not None:
|
||||
results = [results[ix] for ix in listen]
|
||||
|
||||
# set the results of each operation
|
||||
for fut, result in zip(taken, results):
|
||||
fut.set_result(result)
|
||||
except Exception as e:
|
||||
for fut in taken:
|
||||
fut.set_exception(e)
|
||||
# remove the operations from the queue
|
||||
for fut in taken:
|
||||
del queue[fut]
|
||||
else:
|
||||
break
|
||||
# remove strong ref to store
|
||||
del s
|
||||
|
||||
Reference in New Issue
Block a user