feat(storage): share sql/kv instances and add upsert support

Sql stores now share a single SqlAlchemySqlStoreImpl per backend, and kvstore_impl caches instances per (backend, namespace). This avoids spawning multiple SQLite connections for the same file, reducing lock contention and aligning the cache story for all backends.

Added an async upsert API (with SQLite/Postgres dialect inserts) and routed it through AuthorizedSqlStore, then switched conversations and responses to call it. Using native ON CONFLICT DO UPDATE eliminates the insert-then-update retry window that previously caused long WAL lock retries.

Introduced an opt-in conversation stress test that mirrors the recorded prompts from test_conversation_multi_turn_and_streaming while fanning them out across many threads. This gives us a fast local way to hammer the conversations/responses sync path when investigating lockups.
This commit is contained in:
Ashwin Bharambe 2025-11-12 11:27:35 -08:00
parent 531e00366f
commit c2c5ab622c
9 changed files with 300 additions and 51 deletions

View file

@ -9,7 +9,7 @@ from tempfile import TemporaryDirectory
import pytest
from llama_stack.providers.utils.sqlstore.api import ColumnType
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
from llama_stack.providers.utils.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
@ -65,6 +65,38 @@ async def test_sqlite_sqlstore():
assert result.has_more is False
async def test_sqlstore_upsert_support():
with TemporaryDirectory() as tmp_dir:
db_path = tmp_dir + "/upsert.db"
store = SqlAlchemySqlStoreImpl(SqliteSqlStoreConfig(db_path=db_path))
await store.create_table(
"items",
{
"id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
"value": ColumnType.STRING,
"updated_at": ColumnType.INTEGER,
},
)
await store.upsert(
table="items",
data={"id": "item_1", "value": "first", "updated_at": 1},
conflict_columns=["id"],
)
row = await store.fetch_one("items", {"id": "item_1"})
assert row == {"id": "item_1", "value": "first", "updated_at": 1}
await store.upsert(
table="items",
data={"id": "item_1", "value": "second", "updated_at": 2},
conflict_columns=["id"],
update_columns=["value", "updated_at"],
)
row = await store.fetch_one("items", {"id": "item_1"})
assert row == {"id": "item_1", "value": "second", "updated_at": 2}
async def test_sqlstore_pagination_basic():
"""Test basic pagination functionality at the SQL store level."""
with TemporaryDirectory() as tmp_dir: