add postgres kvstoreimpl (#374)

* add postgres kvstoreimpl

* make table name configurable

* add validator for table name

* linter fix

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>
This commit is contained in:
Dinesh Yeduguru 2024-11-05 11:42:21 -08:00 committed by GitHub
parent 8de845a96d
commit dcd8cfe0f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 131 additions and 2 deletions

View file

@ -4,10 +4,11 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
import re
from enum import Enum
from typing import Literal, Optional, Union
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, field_validator
from typing_extensions import Annotated
from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
@ -51,6 +52,22 @@ class PostgresKVStoreConfig(CommonConfig):
db: str = "llamastack"
user: str
password: Optional[str] = None
table_name: str = "llamastack_kvstore"
@field_validator("table_name")
def validate_table_name(self, v: str) -> str:
# PostgreSQL identifiers rules:
# - Must start with a letter or underscore
# - Can contain letters, numbers, and underscores
# - Maximum length is 63 bytes
pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$"
if not re.match(pattern, v):
raise ValueError(
"Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores"
)
if len(v) > 63:
raise ValueError("Table name must be less than 63 characters")
return v
KVStoreConfig = Annotated[