fix precommit errors

This commit is contained in:
Raghotham Murthy 2025-10-02 11:34:10 -07:00
parent d32d8ec94b
commit 82ca72bde6
2 changed files with 9 additions and 7 deletions

View file

@ -21,7 +21,7 @@ class SqliteKVStoreImpl(KVStore):
def __init__(self, config: SqliteKVStoreConfig):
self.db_path = config.db_path
self.table_name = "kvstore"
self._conn = None
self._conn: aiosqlite.Connection | None = None
def __str__(self):
return f"SqliteKVStoreImpl(db_path={self.db_path}, table_name={self.table_name})"
@ -57,6 +57,7 @@ class SqliteKVStoreImpl(KVStore):
self._conn = None
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(
f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
(key, value, expiration),
@ -64,7 +65,10 @@ class SqliteKVStoreImpl(KVStore):
await self._conn.commit()
async def get(self, key: str) -> str | None:
async with self._conn.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
assert self._conn is not None, "Connection not initialized. Call initialize() first."
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
@ -75,10 +79,12 @@ 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()
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(
f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
(start_key, end_key),
@ -91,6 +97,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(
f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
(start_key, end_key),

View file

@ -4,13 +4,11 @@
# 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:")
@ -35,7 +33,6 @@ async def test_memory_kvstore_basic_operations():
await kvstore.close()
@pytest.mark.asyncio
async def test_memory_kvstore_range_operations():
"""Test range operations with :memory: database."""
config = SqliteKVStoreConfig(db_path=":memory:")
@ -66,7 +63,6 @@ async def test_memory_kvstore_range_operations():
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:")
@ -98,7 +94,6 @@ async def test_memory_kvstore_multiple_instances():
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:")