Python: Validate PostgreSQL collection field types

python/semantic_kernel/connectors/postgres.py

@@ -90,53 +90,116 @@
     IndexKind.IVF_FLAT: "ivfflat",
     IndexKind.DEFAULT: "hnsw",
 }
+POSTGRES_TYPE_ALIASES: Final[dict[str, str]] = {
+    "str": "TEXT",
+    "text": "TEXT",
+    "int": "INTEGER",
+    "integer": "INTEGER",
+    "bigint": "BIGINT",
+    "smallint": "SMALLINT",
+    "float": "DOUBLE PRECISION",
+    "double precision": "DOUBLE PRECISION",
+    "real": "REAL",
+    "bool": "BOOLEAN",
+    "boolean": "BOOLEAN",
+    "dict": "JSONB",
+    "jsonb": "JSONB",
+    "date": "DATE",
+    "datetime": "TIMESTAMP",
+    "timestamp": "TIMESTAMP",
+    "timestamptz": "TIMESTAMP WITH TIME ZONE",
+    "timestamp with time zone": "TIMESTAMP WITH TIME ZONE",
+    "uuid": "UUID",
+    "decimal": "DECIMAL",
+    "numeric": "NUMERIC",
+    "bytes": "BYTEA",
+    "bytea": "BYTEA",
+}
+LENGTH_TYPE_PATTERN: Final[re.Pattern[str]] = re.compile(
+    r"(varchar|character varying)\s*\(\s*(\d{1,8})\s*\)", re.IGNORECASE
+)
+PRECISION_TYPE_PATTERN: Final[re.Pattern[str]] = re.compile(
+    r"(decimal|numeric)\s*\(\s*(\d{1,4})(?:\s*,\s*(\d{1,4}))?\s*\)", re.IGNORECASE
+)
 
 # region: Helpers
 
 
-def _python_type_to_postgres(python_type_str: str) -> str | None:
-    """Convert a string representation of a Python type to a PostgreSQL data type.
+def _get_python_generic_arguments(type_name: str, generic_name: str, argument_count: int) -> tuple[str, ...] | None:
+    """Get validated arguments from a Python generic type annotation."""
+    try:
+        expression = ast.parse(type_name, mode="eval").body
+    except SyntaxError:
+        return None
+
+    if not isinstance(expression, ast.Subscript):
+        return None
+
+    if isinstance(expression.value, ast.Name):
+        expression_name = expression.value.id
+    else:
+        return None
+
+    arguments = expression.slice.elts if isinstance(expression.slice, ast.Tuple) else [expression.slice]
+    if expression_name.casefold() != generic_name or len(arguments) != argument_count:
+        return None
+
+    def is_valid_annotation(node: ast.expr) -> bool:
+        if isinstance(node, ast.Name):
+            return True
+        if isinstance(node, ast.Attribute):
+            return is_valid_annotation(node.value)
+        if isinstance(node, ast.Subscript):
+            values = node.slice.elts if isinstance(node.slice, ast.Tuple) else [node.slice]
+            return is_valid_annotation(node.value) and all(is_valid_annotation(value) for value in values)
+        return False
+
+    if not all(is_valid_annotation(argument) for argument in arguments):
+        return None
+
+    return tuple(ast.unparse(argument) for argument in arguments)
+
+
+def _get_postgres_type(type_name: str) -> sql.Composable:
+    """Convert a supported field type to a trusted PostgreSQL type expression.
 
     Args:
-        python_type_str: The string representation of the Python type (e.g., "int", "List[str]").
+        type_name: The Python or PostgreSQL type name.
 
     Returns:
-        Corresponding PostgreSQL data type as a string, if found. If the type is not found, return None.
+        A trusted PostgreSQL type expression.
+
+    Raises:
+        VectorStoreModelValidationError: If the type is not supported for automatic table creation.
     """
-    # Basic type mapping from Python types (in string form) to PostgreSQL types
-    type_mapping = {
-        "str": "TEXT",
-        "int": "INTEGER",
-        "float": "DOUBLE PRECISION",
-        "bool": "BOOLEAN",
-        "dict": "JSONB",
-        "datetime": "TIMESTAMP",
-        "bytes": "BYTEA",
-        "NoneType": "NULL",
-    }
-
-    # Regular expression to detect lists, e.g., "List[str]" or "List[int]"
-    list_pattern = re.compile(r"(?i)List\[(.*)\]")
-
-    # Check if the type is a list
-    match = list_pattern.match(python_type_str)
-    if match:
-        # Extract the inner type of the list and convert it to a PostgreSQL array type
-        element_type_str = match.group(1)
-        postgres_element_type = _python_type_to_postgres(element_type_str)
-        return f"{postgres_element_type}[]"
-
-    # Check if the type is a dictionary
-    dict_pattern = re.compile(r"(?i)Dict\[(.*), (.*)\]")
-    match = dict_pattern.match(python_type_str)
-    if match:
-        return "JSONB"
-
-    # Handle basic types
-    if python_type_str in type_mapping:
-        return type_mapping[python_type_str]
-
-    return None
+    normalized_type = " ".join(type_name.strip().split())
+    if postgres_type := POSTGRES_TYPE_ALIASES.get(normalized_type.casefold()):
+        return sql.SQL(postgres_type)
+
+    if normalized_type.endswith("[]"):
+        return sql.SQL("{}[]").format(_get_postgres_type(normalized_type[:-2]))
+
+    if match := LENGTH_TYPE_PATTERN.fullmatch(normalized_type):
+        postgres_type = "VARCHAR" if match.group(1).casefold() == "varchar" else "CHARACTER VARYING"
+        return sql.SQL("{}({})").format(sql.SQL(postgres_type), sql.Literal(int(match.group(2))))
+
+    if match := PRECISION_TYPE_PATTERN.fullmatch(normalized_type):
+        postgres_type = POSTGRES_TYPE_ALIASES[match.group(1).casefold()]
+        modifiers = [sql.Literal(int(match.group(2)))]
+        if scale := match.group(3):
+            modifiers.append(sql.Literal(int(scale)))
+        return sql.SQL("{}({})").format(sql.SQL(postgres_type), sql.SQL(", ").join(modifiers))
+
+    if list_arguments := _get_python_generic_arguments(normalized_type, "list", 1):
+        return sql.SQL("{}[]").format(_get_postgres_type(list_arguments[0]))
+
+    if _get_python_generic_arguments(normalized_type, "dict", 2):
+        return sql.SQL("JSONB")
+
+    raise VectorStoreModelValidationError(
+        f"PostgreSQL field type {type_name!r} is not supported for automatic collection creation. "
+        "Use a supported Python or PostgreSQL type, or create the table before using the collection."
+    )
 
 
 def _convert_row_to_dict(row: tuple[Any, ...], fields: Sequence[tuple[str, VectorStoreField | None]]) -> dict[str, Any]:
@@ -562,12 +625,18 @@ def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: A
 
     @override
     async def ensure_collection_exists(self, **kwargs: Any) -> None:
-        """Create a PostgreSQL table based on a dictionary of VectorStoreRecordField.
+        """Create a PostgreSQL table based on the collection definition.
+
+        Automatic creation supports built-in Python field types, their corresponding PostgreSQL aliases,
+        validated length and precision modifiers, and arrays of those types. For custom domains, enums,
+        extension types, or other type expressions, create the table separately and use the collection without
+        calling this method.
 
         Args:
-            table_name: Name of the table to be created
-            fields: A dictionary where keys are column names and values are VectorStoreRecordField instances
-            **kwargs: Additional arguments
+            **kwargs: Additional arguments.
+
+        Raises:
+            VectorStoreModelValidationError: If a field type is not supported for automatic table creation.
         """
         if self.connection_pool is None:
             raise VectorStoreOperationException(
@@ -581,8 +650,7 @@ async def ensure_collection_exists(self, **kwargs: Any) -> None:
             if not field.type_:
                 raise ValueError(f"Property type is not defined for field '{field.name}'")
 
-            # If the property type represents a Python type, convert it to a PostgreSQL type
-            property_type = _python_type_to_postgres(field.type_) or field.type_.upper()
+            property_type = _get_postgres_type(field.type_)
 
             # For Vector fields with dimensions, use pgvector's VECTOR type
             # Note that other vector types are supported in pgvector (e.g. halfvec),
@@ -595,17 +663,15 @@ async def ensure_collection_exists(self, **kwargs: Any) -> None:
                     )
                 )
             elif field.field_type == FieldTypes.KEY:
-                # Use the property_type directly for key fields
                 column_definitions.append(
                     sql.SQL("{name} {col_type} PRIMARY KEY").format(
-                        name=sql.Identifier(field.storage_name or field.name), col_type=sql.SQL(property_type)
+                        name=sql.Identifier(field.storage_name or field.name), col_type=property_type
                     )
                 )
             else:
-                # Use the property_type directly for other types
                 column_definitions.append(
                     sql.SQL("{name} {col_type}").format(
-                        name=sql.Identifier(field.storage_name or field.name), col_type=sql.SQL(property_type)
+                        name=sql.Identifier(field.storage_name or field.name), col_type=property_type
                     )
                 )
 

python/tests/integration/memory/postgres/test_postgres_int.py

@@ -8,6 +8,7 @@
 import pandas as pd
 import pytest
 import pytest_asyncio
+from psycopg import sql
 from pydantic import BaseModel
 
 from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore
@@ -19,6 +20,7 @@
     VectorStoreField,
     vectorstoremodel,
 )
+from semantic_kernel.exceptions import VectorStoreModelValidationError
 from semantic_kernel.exceptions.memory_connector_exceptions import (
     MemoryConnectorConnectionException,
     MemoryConnectorInitializationError,
@@ -143,6 +145,62 @@ async def test_ensure_collection_exists_exists_and_delete(vector_store: Postgres
     assert does_exist_3 is False
 
 
+async def test_ensure_collection_exists_rejects_unsupported_type_without_database_side_effects(
+    vector_store: PostgresStore,
+):
+    suffix = str(uuid.uuid4()).replace("-", "")[:8]
+    sentinel_name = f"sentinel_{suffix}"
+    collection_name = f"test_collection_{suffix}"
+    connection_pool = vector_store.connection_pool
+
+    async with connection_pool.connection() as conn, conn.cursor() as cur:
+        await cur.execute(
+            sql.SQL("CREATE TABLE {schema}.{table} (id INTEGER)").format(
+                schema=sql.Identifier(vector_store.db_schema),
+                table=sql.Identifier(sentinel_name),
+            )
+        )
+        await conn.commit()
+
+    definition = VectorStoreCollectionDefinition(
+        fields=[
+            VectorStoreField("key", name="id", type="int"),
+            VectorStoreField(
+                "data",
+                name="content",
+                type=f"TEXT); DROP TABLE {sentinel_name}; --",
+            ),
+        ]
+    )
+    collection = vector_store.get_collection(
+        collection_name=collection_name,
+        record_type=dict,
+        definition=definition,
+    )
+
+    try:
+        with pytest.raises(VectorStoreModelValidationError, match="not supported for automatic collection creation"):
+            await collection.ensure_collection_exists()
+
+        async with connection_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                "SELECT to_regclass(%s)",
+                (f"{vector_store.db_schema}.{sentinel_name}",),
+            )
+            row = await cur.fetchone()
+            assert row is not None
+            assert row[0] == sentinel_name
+    finally:
+        async with connection_pool.connection() as conn, conn.cursor() as cur:
+            await cur.execute(
+                sql.SQL("DROP TABLE IF EXISTS {schema}.{table}").format(
+                    schema=sql.Identifier(vector_store.db_schema),
+                    table=sql.Identifier(sentinel_name),
+                )
+            )
+            await conn.commit()
+
+
 async def test_list_collection_names(vector_store):
     async with create_simple_collection(vector_store) as simple_collection:
         simple_collection_id = simple_collection.collection_name