phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-30 23:51:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

make table name configurable

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-05 11:31:08 -08:00
parent 726fd15dc4
commit a4c8b3ec0a
2 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

1
llama_stack/providers/utils/kvstore/config.py
View file

@ -51,6 +51,7 @@ class PostgresKVStoreConfig(CommonConfig):
db: str = "llamastack" db: str = "llamastack"
user: str user: str
password: Optional[str] = None password: Optional[str] = None
table_name: str = "llamastack_kvstore"
KVStoreConfig = Annotated[ KVStoreConfig = Annotated[

11
llama_stack/providers/utils/kvstore/postgres/postgres.py
View file

@ -17,7 +17,6 @@ from ..config import PostgresKVStoreConfig
class PostgresKVStoreImpl(KVStore): class PostgresKVStoreImpl(KVStore):
def __init__(self, config: PostgresKVStoreConfig): def __init__(self, config: PostgresKVStoreConfig):
self.config = config self.config = config
self.table_name = "kvstore"
self.conn = None self.conn = None
self.cursor = None self.cursor = None
@ -36,7 +35,7 @@ class PostgresKVStoreImpl(KVStore):
# Create table if it doesn't exist # Create table if it doesn't exist
self.cursor.execute( self.cursor.execute(
f""" f"""
CREATE TABLE IF NOT EXISTS {self.table_name} ( CREATE TABLE IF NOT EXISTS {self.config.table_name} (
key TEXT PRIMARY KEY, key TEXT PRIMARY KEY,
value TEXT, value TEXT,
expiration TIMESTAMP expiration TIMESTAMP
@ -60,7 +59,7 @@ class PostgresKVStoreImpl(KVStore):
key = self._namespaced_key(key) key = self._namespaced_key(key)
self.cursor.execute( self.cursor.execute(
f""" f"""
INSERT INTO {self.table_name} (key, value, expiration) INSERT INTO {self.config.table_name} (key, value, expiration)
VALUES (%s, %s, %s) VALUES (%s, %s, %s)
ON CONFLICT (key) DO UPDATE ON CONFLICT (key) DO UPDATE
SET value = EXCLUDED.value, expiration = EXCLUDED.expiration SET value = EXCLUDED.value, expiration = EXCLUDED.expiration
@ -72,7 +71,7 @@ class PostgresKVStoreImpl(KVStore):
key = self._namespaced_key(key) key = self._namespaced_key(key)
self.cursor.execute( self.cursor.execute(
f""" f"""
SELECT value FROM {self.table_name} SELECT value FROM {self.config.table_name}
WHERE key = %s WHERE key = %s
AND (expiration IS NULL OR expiration > NOW()) AND (expiration IS NULL OR expiration > NOW())
""", """,
@ -84,7 +83,7 @@ class PostgresKVStoreImpl(KVStore):
async def delete(self, key: str) -> None: async def delete(self, key: str) -> None:
key = self._namespaced_key(key) key = self._namespaced_key(key)
self.cursor.execute( self.cursor.execute(
f"DELETE FROM {self.table_name} WHERE key = %s", f"DELETE FROM {self.config.table_name} WHERE key = %s",
(key,), (key,),
) )
@ -94,7 +93,7 @@ class PostgresKVStoreImpl(KVStore):
self.cursor.execute( self.cursor.execute(
f""" f"""
SELECT value FROM {self.table_name} SELECT value FROM {self.config.table_name}
WHERE key >= %s AND key < %s WHERE key >= %s AND key < %s
AND (expiration IS NULL OR expiration > NOW()) AND (expiration IS NULL OR expiration > NOW())
ORDER BY key ORDER BY key