From 82ca72bde629a8dee452a8cdf27698feab94d3bf Mon Sep 17 00:00:00 2001
From: Raghotham Murthy <rsm@meta.com>
Date: Thu, 2 Oct 2025 11:34:10 -0700
Subject: [PATCH] fix precommit errors

---
 llama_stack/providers/utils/kvstore/sqlite/sqlite.py | 11 +++++++++--
 tests/unit/utils/kvstore/test_sqlite_memory.py       |  5 -----
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py
index 5372d2981..e1c0332fe 100644
--- a/llama_stack/providers/utils/kvstore/sqlite/sqlite.py
+++ b/llama_stack/providers/utils/kvstore/sqlite/sqlite.py
@@ -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),
diff --git a/tests/unit/utils/kvstore/test_sqlite_memory.py b/tests/unit/utils/kvstore/test_sqlite_memory.py
index 8017f9bdd..6d8267ce8 100644
--- a/tests/unit/utils/kvstore/test_sqlite_memory.py
+++ b/tests/unit/utils/kvstore/test_sqlite_memory.py
@@ -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:")