mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-04 10:10:36 +00:00
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:
parent
531e00366f
commit
c2c5ab622c
9 changed files with 300 additions and 51 deletions
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue