mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-03 19:57:35 +00:00
fix precommit errors
This commit is contained in:
parent
d32d8ec94b
commit
82ca72bde6
2 changed files with 9 additions and 7 deletions
|
@ -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),
|
||||
|
|
|
@ -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:")
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue