mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-27 18:50:41 +00:00
d3b60507d7
# What does this PR do? Inference/Response stores now store user attributes when inserting, and respects them when fetching. ## Test Plan pytest tests/unit/utils/test_sqlstore.py
128 lines
3.6 KiB
Python
128 lines
3.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from collections.abc import Mapping
|
|
from enum import Enum
|
|
from typing import Any, Literal, Protocol
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from llama_stack.apis.common.responses import PaginatedResponse
|
|
|
|
|
|
class ColumnType(Enum):
|
|
INTEGER = "INTEGER"
|
|
STRING = "STRING"
|
|
TEXT = "TEXT"
|
|
FLOAT = "FLOAT"
|
|
BOOLEAN = "BOOLEAN"
|
|
JSON = "JSON"
|
|
DATETIME = "DATETIME"
|
|
|
|
|
|
class ColumnDefinition(BaseModel):
|
|
type: ColumnType
|
|
primary_key: bool = False
|
|
nullable: bool = True
|
|
default: Any = None
|
|
|
|
|
|
class SqlStore(Protocol):
|
|
"""
|
|
A protocol for a SQL store.
|
|
"""
|
|
|
|
async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None:
|
|
"""
|
|
Create a table.
|
|
"""
|
|
pass
|
|
|
|
async def insert(self, table: str, data: Mapping[str, Any]) -> None:
|
|
"""
|
|
Insert a row into a table.
|
|
"""
|
|
pass
|
|
|
|
async def fetch_all(
|
|
self,
|
|
table: str,
|
|
where: Mapping[str, Any] | None = None,
|
|
where_sql: str | None = None,
|
|
limit: int | None = None,
|
|
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
|
cursor: tuple[str, str] | None = None,
|
|
) -> PaginatedResponse:
|
|
"""
|
|
Fetch all rows from a table with optional cursor-based pagination.
|
|
|
|
:param table: The table name
|
|
:param where: Simple key-value WHERE conditions
|
|
:param where_sql: Raw SQL WHERE clause for complex queries
|
|
:param limit: Maximum number of records to return
|
|
:param order_by: List of (column, order) tuples for sorting
|
|
:param cursor: Tuple of (key_column, cursor_id) for pagination (None for first page)
|
|
Requires order_by with exactly one column when used
|
|
:return: PaginatedResult with data and has_more flag
|
|
|
|
Note: Cursor pagination only supports single-column ordering for simplicity.
|
|
Multi-column ordering is allowed without cursor but will raise an error with cursor.
|
|
"""
|
|
pass
|
|
|
|
async def fetch_one(
|
|
self,
|
|
table: str,
|
|
where: Mapping[str, Any] | None = None,
|
|
where_sql: str | None = None,
|
|
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
|
) -> dict[str, Any] | None:
|
|
"""
|
|
Fetch one row from a table.
|
|
"""
|
|
pass
|
|
|
|
async def update(
|
|
self,
|
|
table: str,
|
|
data: Mapping[str, Any],
|
|
where: Mapping[str, Any],
|
|
) -> None:
|
|
"""
|
|
Update a row in a table.
|
|
"""
|
|
pass
|
|
|
|
async def delete(
|
|
self,
|
|
table: str,
|
|
where: Mapping[str, Any],
|
|
) -> None:
|
|
"""
|
|
Delete a row from a table.
|
|
"""
|
|
pass
|
|
|
|
async def add_column_if_not_exists(
|
|
self,
|
|
table: str,
|
|
column_name: str,
|
|
column_type: ColumnType,
|
|
nullable: bool = True,
|
|
) -> None:
|
|
"""
|
|
Add a column to an existing table if the column doesn't already exist.
|
|
|
|
This is useful for table migrations when adding new functionality.
|
|
If the table doesn't exist, this method should do nothing.
|
|
If the column already exists, this method should do nothing.
|
|
|
|
:param table: Table name
|
|
:param column_name: Name of the column to add
|
|
:param column_type: Type of the column to add
|
|
:param nullable: Whether the column should be nullable (default: True)
|
|
"""
|
|
pass
|