From d32d8ec94b2fc664e72d508c7e0de1966bba623e Mon Sep 17 00:00:00 2001 From: Raghotham Murthy Date: Thu, 2 Oct 2025 11:17:41 -0700 Subject: [PATCH] feat: Add support for memory-only kvstore --- llama_stack/providers/utils/kvstore/config.py | 2 +- .../providers/utils/kvstore/sqlite/sqlite.py | 106 ++++++++------- .../unit/utils/kvstore/test_sqlite_memory.py | 121 ++++++++++++++++++ 3 files changed, 181 insertions(+), 48 deletions(-) create mode 100644 tests/unit/utils/kvstore/test_sqlite_memory.py diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index 7b6a79350..4899bbc80 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -53,7 +53,7 @@ class SqliteKVStoreConfig(CommonConfig): type: Literal["sqlite"] = KVStoreType.sqlite.value db_path: str = Field( default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(), - description="File path for the sqlite database", + description="File path for the sqlite database. Use ':memory:' for an in-memory database", ) @classmethod diff --git a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py index 5b782902e..5372d2981 100644 --- a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py +++ b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py @@ -21,67 +21,79 @@ class SqliteKVStoreImpl(KVStore): def __init__(self, config: SqliteKVStoreConfig): self.db_path = config.db_path self.table_name = "kvstore" + self._conn = None def __str__(self): return f"SqliteKVStoreImpl(db_path={self.db_path}, table_name={self.table_name})" + def _is_memory_db(self) -> bool: + """Check if this is an in-memory database.""" + return self.db_path == ":memory:" or "mode=memory" in self.db_path + async def initialize(self): - os.makedirs(os.path.dirname(self.db_path), exist_ok=True) - async with aiosqlite.connect(self.db_path) as db: - await db.execute( - f""" - CREATE TABLE IF NOT EXISTS {self.table_name} ( - key TEXT PRIMARY KEY, - value TEXT, - expiration TIMESTAMP - ) - """ + # Skip directory creation for in-memory databases and file: URIs + if not self._is_memory_db() and not self.db_path.startswith("file:"): + db_dir = os.path.dirname(self.db_path) + if db_dir: # Only create if there's a directory component + os.makedirs(db_dir, exist_ok=True) + + # Create persistent connection for all databases + self._conn = await aiosqlite.connect(self.db_path) + await self._conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {self.table_name} ( + key TEXT PRIMARY KEY, + value TEXT, + expiration TIMESTAMP ) - await db.commit() + """ + ) + await self._conn.commit() + + async def close(self): + """Close the persistent connection.""" + if self._conn: + await self._conn.close() + self._conn = None async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: - async with aiosqlite.connect(self.db_path) as db: - await db.execute( - f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)", - (key, value, expiration), - ) - await db.commit() + await self._conn.execute( + f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)", + (key, value, expiration), + ) + await self._conn.commit() async def get(self, key: str) -> str | None: - async with aiosqlite.connect(self.db_path) as db: - async with db.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor: - row = await cursor.fetchone() - if row is None: - return None - value, expiration = row - if not isinstance(value, str): - logger.warning(f"Expected string value for key {key}, got {type(value)}, returning None") - return None - return value + async with self._conn.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor: + row = await cursor.fetchone() + if row is None: + return None + value, expiration = row + if not isinstance(value, str): + logger.warning(f"Expected string value for key {key}, got {type(value)}, returning None") + return None + return value async def delete(self, key: str) -> None: - async with aiosqlite.connect(self.db_path) as db: - await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,)) - await db.commit() + await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,)) + await self._conn.commit() async def values_in_range(self, start_key: str, end_key: str) -> list[str]: - async with aiosqlite.connect(self.db_path) as db: - async with db.execute( - f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?", - (start_key, end_key), - ) as cursor: - result = [] - async for row in cursor: - _, value, _ = row - result.append(value) - return result + async with self._conn.execute( + f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?", + (start_key, end_key), + ) as cursor: + result = [] + async for row in cursor: + _, value, _ = row + result.append(value) + return result async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: """Get all keys in the given range.""" - async with aiosqlite.connect(self.db_path) as db: - cursor = await db.execute( - f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?", - (start_key, end_key), - ) - rows = await cursor.fetchall() - return [row[0] for row in rows] + cursor = await self._conn.execute( + f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?", + (start_key, end_key), + ) + rows = await cursor.fetchall() + return [row[0] for row in rows] diff --git a/tests/unit/utils/kvstore/test_sqlite_memory.py b/tests/unit/utils/kvstore/test_sqlite_memory.py new file mode 100644 index 000000000..8017f9bdd --- /dev/null +++ b/tests/unit/utils/kvstore/test_sqlite_memory.py @@ -0,0 +1,121 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + +import pytest + +from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig +from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl + + +@pytest.mark.asyncio +async def test_memory_kvstore_basic_operations(): + """Test basic CRUD operations with :memory: database.""" + config = SqliteKVStoreConfig(db_path=":memory:") + kvstore = SqliteKVStoreImpl(config) + await kvstore.initialize() + + try: + # Test set and get + await kvstore.set("test_key", "test_value") + value = await kvstore.get("test_key") + assert value == "test_value" + + # Test get non-existent key + value = await kvstore.get("non_existent") + assert value is None + + # Test delete + await kvstore.delete("test_key") + value = await kvstore.get("test_key") + assert value is None + finally: + await kvstore.close() + + +@pytest.mark.asyncio +async def test_memory_kvstore_range_operations(): + """Test range operations with :memory: database.""" + config = SqliteKVStoreConfig(db_path=":memory:") + kvstore = SqliteKVStoreImpl(config) + await kvstore.initialize() + + try: + # Set multiple keys + await kvstore.set("key_1", "value_1") + await kvstore.set("key_2", "value_2") + await kvstore.set("key_3", "value_3") + await kvstore.set("key_4", "value_4") + + # Test values_in_range + values = await kvstore.values_in_range("key_1", "key_3") + assert len(values) == 3 + assert "value_1" in values + assert "value_2" in values + assert "value_3" in values + + # Test keys_in_range + keys = await kvstore.keys_in_range("key_2", "key_4") + assert len(keys) == 3 + assert "key_2" in keys + assert "key_3" in keys + assert "key_4" in keys + finally: + await kvstore.close() + + +@pytest.mark.asyncio +async def test_memory_kvstore_multiple_instances(): + """Test that multiple :memory: instances are independent.""" + config1 = SqliteKVStoreConfig(db_path=":memory:") + kvstore1 = SqliteKVStoreImpl(config1) + await kvstore1.initialize() + + config2 = SqliteKVStoreConfig(db_path=":memory:") + kvstore2 = SqliteKVStoreImpl(config2) + await kvstore2.initialize() + + try: + # Set value in first instance + await kvstore1.set("shared_key", "value_1") + + # Verify second instance doesn't have the value + value = await kvstore2.get("shared_key") + assert value is None + + # Set different value in second instance + await kvstore2.set("shared_key", "value_2") + + # Verify instances remain independent + value1 = await kvstore1.get("shared_key") + value2 = await kvstore2.get("shared_key") + assert value1 == "value_1" + assert value2 == "value_2" + finally: + await kvstore1.close() + await kvstore2.close() + + +@pytest.mark.asyncio +async def test_memory_kvstore_persistence_behavior(): + """Test that :memory: database doesn't persist across instances.""" + config = SqliteKVStoreConfig(db_path=":memory:") + + # First instance + kvstore1 = SqliteKVStoreImpl(config) + await kvstore1.initialize() + await kvstore1.set("test_key", "test_value") + await kvstore1.close() + + # Create new instance with same config + kvstore2 = SqliteKVStoreImpl(config) + await kvstore2.initialize() + + try: + # Data should not persist + value = await kvstore2.get("test_key") + assert value is None + finally: + await kvstore2.close()