From ed78090b8e6841dceb8d24b89924105a2fab3bdd Mon Sep 17 00:00:00 2001 From: Raghotham Murthy Date: Tue, 7 Oct 2025 14:19:29 -0700 Subject: [PATCH] address comments --- .../providers/utils/kvstore/sqlite/sqlite.py | 26 +++--- .../unit/utils/kvstore/test_sqlite_memory.py | 85 ------------------- 2 files changed, 14 insertions(+), 97 deletions(-) diff --git a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py index e1c0332fe..4759bfa88 100644 --- a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py +++ b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py @@ -56,17 +56,22 @@ class SqliteKVStoreImpl(KVStore): await self._conn.close() self._conn = None + @property + def conn(self) -> aiosqlite.Connection: + """Get the connection, raising an error if not initialized.""" + if self._conn is None: + raise RuntimeError("Connection not initialized. Call initialize() first.") + return self._conn + async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: - assert self._conn is not None, "Connection not initialized. Call initialize() first." - await self._conn.execute( + await self.conn.execute( f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)", (key, value, expiration), ) - await self._conn.commit() + await self.conn.commit() async def get(self, key: str) -> str | None: - assert self._conn is not None, "Connection not initialized. Call initialize() first." - async with self._conn.execute( + async with self.conn.execute( f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,) ) as cursor: row = await cursor.fetchone() @@ -79,13 +84,11 @@ class SqliteKVStoreImpl(KVStore): return value async def delete(self, key: str) -> None: - assert self._conn is not None, "Connection not initialized. Call initialize() first." - await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,)) - await self._conn.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]: - assert self._conn is not None, "Connection not initialized. Call initialize() first." - async with self._conn.execute( + async with self.conn.execute( f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?", (start_key, end_key), ) as cursor: @@ -97,8 +100,7 @@ class SqliteKVStoreImpl(KVStore): async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: """Get all keys in the given range.""" - assert self._conn is not None, "Connection not initialized. Call initialize() first." - cursor = await self._conn.execute( + cursor = await self.conn.execute( f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?", (start_key, end_key), ) diff --git a/tests/unit/utils/kvstore/test_sqlite_memory.py b/tests/unit/utils/kvstore/test_sqlite_memory.py index 16d9cc101..942ad1087 100644 --- a/tests/unit/utils/kvstore/test_sqlite_memory.py +++ b/tests/unit/utils/kvstore/test_sqlite_memory.py @@ -9,91 +9,6 @@ from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig from llama_stack.providers.utils.kvstore.sqlite.sqlite import SqliteKVStoreImpl -async def test_memory_kvstore_basic_operations(): - """Test basic CRUD operations with :memory: database.""" - config = SqliteKVStoreConfig(db_path=":memory:") - store = SqliteKVStoreImpl(config) - await store.initialize() - - # Test set and get - await store.set("key1", "value1") - result = await store.get("key1") - assert result == "value1" - - # Test get non-existent key - result = await store.get("nonexistent") - assert result is None - - # Test update - await store.set("key1", "updated_value") - result = await store.get("key1") - assert result == "updated_value" - - # Test delete - await store.delete("key1") - result = await store.get("key1") - assert result is None - - await store.close() - - -async def test_memory_kvstore_range_operations(): - """Test range query operations with :memory: database.""" - config = SqliteKVStoreConfig(db_path=":memory:") - store = SqliteKVStoreImpl(config) - await store.initialize() - - # Set up test data - await store.set("key_a", "value_a") - await store.set("key_b", "value_b") - await store.set("key_c", "value_c") - await store.set("key_d", "value_d") - - # Test values_in_range - values = await store.values_in_range("key_b", "key_c") - assert len(values) == 2 - assert "value_b" in values - assert "value_c" in values - - # Test keys_in_range - keys = await store.keys_in_range("key_a", "key_c") - assert len(keys) == 3 - assert "key_a" in keys - assert "key_b" in keys - assert "key_c" in keys - - await store.close() - - -async def test_memory_kvstore_multiple_instances(): - """Test that multiple :memory: instances are independent.""" - config1 = SqliteKVStoreConfig(db_path=":memory:") - config2 = SqliteKVStoreConfig(db_path=":memory:") - - store1 = SqliteKVStoreImpl(config1) - store2 = SqliteKVStoreImpl(config2) - - await store1.initialize() - await store2.initialize() - - # Set data in store1 - await store1.set("shared_key", "value_from_store1") - - # Verify store2 doesn't see store1's data - result = await store2.get("shared_key") - assert result is None - - # Set different value in store2 - await store2.set("shared_key", "value_from_store2") - - # Verify both stores have independent data - assert await store1.get("shared_key") == "value_from_store1" - assert await store2.get("shared_key") == "value_from_store2" - - await store1.close() - await store2.close() - - async def test_memory_kvstore_persistence_behavior(): """Test that :memory: database doesn't persist across instances.""" config = SqliteKVStoreConfig(db_path=":memory:")