add validator for table name

This commit is contained in:
Dinesh Yeduguru 2024-11-05 11:38:42 -08:00
parent a4c8b3ec0a
commit 5567243553

View file

@ -7,8 +7,9 @@
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
import re
from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
@ -53,6 +54,19 @@ class PostgresKVStoreConfig(CommonConfig):
password: Optional[str] = None
table_name: str = "llamastack_kvstore"
@field_validator('table_name')
def validate_table_name(cls, 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[
Union[RedisKVStoreConfig, SqliteKVStoreConfig, PostgresKVStoreConfig],