From 55672435532acaf335af313d5a9be7f1f8305fbb Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Tue, 5 Nov 2024 11:38:42 -0800 Subject: [PATCH] add validator for table name --- llama_stack/providers/utils/kvstore/config.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index b4c5cd7b0..5829650d6 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -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],