address comments

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.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
+        await self.conn.commit()
-        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),
         )

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:")