fix: Replace raw SQL string interpolation with proper SQLAlchemy parameterized APIs in PostgresKVStore (#20104)

* fix: Eliminate SQL injection vulnerabilities in PostgresKVStore

This commit addresses multiple SQL injection vulnerabilities in the
PostgresKVStore integration by replacing unsafe string interpolation
with proper SQLAlchemy parameterized APIs.

## Vulnerabilities Fixed

1. **_create_schema_if_not_exists()** (lines 223-231)
   - Replaced f-string interpolation in SELECT query
   - Replaced f-string interpolation in CREATE SCHEMA statement
   - Now uses sqlalchemy.schema.CreateSchema with if_not_exists parameter

2. **put_all()** (lines 305-310)
   - Replaced raw SQL text() with f-string table/schema names
   - Now uses sqlalchemy.dialects.postgresql.insert() with proper
     parameterization for both identifiers and values

3. **aput_all()** (lines 347-352)
   - Same vulnerabilities and fixes as put_all() for async version

## Changes

- Import CreateSchema at module level for cleaner code
- Replace text(f"SELECT ... WHERE schema_name = '{self.schema_name}'")
  with CreateSchema(self.schema_name, if_not_exists=True)
- Replace text(f"INSERT INTO {self.schema_name}.{tablename} ...")
  with insert(self._table_class).values().on_conflict_do_update()
- All user data continues to be properly parameterized

## Security Impact

Before: Attackers could inject arbitrary SQL via schema_name or
indirectly through table_name, potentially leading to:
- Data exfiltration
- Data manipulation
- Privilege escalation
- Database schema manipulation

After: All SQL identifiers and values are properly handled through
SQLAlchemy's parameterization APIs, eliminating SQL injection vectors.

## Testing

Added 4 new security-focused tests:
- test_schema_creation_uses_safe_api: Verifies CreateSchema usage
- test_put_all_uses_safe_insert: Verifies parameterized insert
- test_aput_all_uses_safe_insert: Verifies async parameterized insert
- test_schema_name_with_special_characters: Tests injection attempts

All existing and new tests pass (6/6 tests passing).

* fix: Remove double JSON serialization in insert operations

SQLAlchemy automatically handles JSON serialization for JSON/JSONB column
types. Manually calling json.dumps() was causing double serialization,
returning JSON strings instead of dicts.

This fixes the failing integration tests that expected dict values.

* bump pyproject.toml version to 0.4.2

Author: ColeMurray

llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/llama_index/storage/kvstore/postgres/base.py

@@ -1,11 +1,12 @@
-import json
 from typing import Any, Dict, List, Optional, Tuple, Type
 from urllib.parse import urlparse
 from llama_index.core.bridge.pydantic import PrivateAttr
 
 try:
     import sqlalchemy
     import sqlalchemy.ext.asyncio  # noqa
+    from sqlalchemy import Column, Index, Integer, UniqueConstraint
+    from sqlalchemy.schema import CreateSchema
 except ImportError:
     raise ImportError("`sqlalchemy[asyncio]` package should be pre installed")
 
@@ -27,7 +28,6 @@ def get_data_model(
     """
     This part create a dynamic sqlalchemy model with a new table.
     """
-    from sqlalchemy import Column, Index, Integer, UniqueConstraint
     from sqlalchemy.dialects.postgresql import JSON, JSONB, VARCHAR
 
     tablename = "data_%s" % index_name  # dynamic table name
@@ -216,21 +216,7 @@ def _connect(self) -> Any:
 
     def _create_schema_if_not_exists(self) -> None:
         with self._session() as session, session.begin():
-            from sqlalchemy import text
-
-            # Check if the specified schema exists with "CREATE" statement
-            check_schema_statement = text(
-                f"SELECT schema_name FROM information_schema.schemata WHERE schema_name = '{self.schema_name}'"
-            )
-            result = session.execute(check_schema_statement).fetchone()
-
-            # If the schema does not exist, then create it
-            if not result:
-                create_schema_statement = text(
-                    f"CREATE SCHEMA IF NOT EXISTS {self.schema_name}"
-                )
-                session.execute(create_schema_statement)
-
+            session.execute(CreateSchema(self.schema_name, if_not_exists=True))
             session.commit()
 
     def _create_tables_if_not_exists(self) -> None:
@@ -285,40 +271,29 @@ def put_all(
         collection: str = DEFAULT_COLLECTION,
         batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> None:
-        from sqlalchemy import text
+        from sqlalchemy.dialects.postgresql import insert
 
         self._initialize()
         with self._session() as session:
             for i in range(0, len(kv_pairs), batch_size):
                 batch = kv_pairs[i : i + batch_size]
 
-                # Prepare the VALUES part of the SQL statement
-                values_clause = ", ".join(
-                    f"(:key_{i}, :namespace_{i}, :value_{i})"
-                    for i, _ in enumerate(batch)
+                values_to_insert = [
+                    {
+                        "key": key,
+                        "namespace": collection,
+                        "value": value,
+                    }
+                    for key, value in batch
+                ]
+
+                stmt = insert(self._table_class).values(values_to_insert)
+                stmt = stmt.on_conflict_do_update(
+                    index_elements=["key", "namespace"],
+                    set_={"value": stmt.excluded.value},
                 )
 
-                # Prepare the raw SQL for bulk upsert
-                # Note: This SQL is PostgreSQL-specific. Adjust for other databases.
-                stmt = text(
-                    f"""
-                INSERT INTO {self.schema_name}.{self._table_class.__tablename__} (key, namespace, value)
-                VALUES {values_clause}
-                ON CONFLICT (key, namespace)
-                DO UPDATE SET
-                value = EXCLUDED.value;
-                """
-                )
-
-                # Flatten the list of tuples for execute parameters
-                params = {}
-                for i, (key, value) in enumerate(batch):
-                    params[f"key_{i}"] = key
-                    params[f"namespace_{i}"] = collection
-                    params[f"value_{i}"] = json.dumps(value)
-
-                # Execute the bulk upsert
-                session.execute(stmt, params)
+                session.execute(stmt)
                 session.commit()
 
     async def aput_all(
@@ -327,40 +302,29 @@ async def aput_all(
         collection: str = DEFAULT_COLLECTION,
         batch_size: int = DEFAULT_BATCH_SIZE,
     ) -> None:
-        from sqlalchemy import text
+        from sqlalchemy.dialects.postgresql import insert
 
         self._initialize()
         async with self._async_session() as session:
             for i in range(0, len(kv_pairs), batch_size):
                 batch = kv_pairs[i : i + batch_size]
 
-                # Prepare the VALUES part of the SQL statement
-                values_clause = ", ".join(
-                    f"(:key_{i}, :namespace_{i}, :value_{i})"
-                    for i, _ in enumerate(batch)
+                values_to_insert = [
+                    {
+                        "key": key,
+                        "namespace": collection,
+                        "value": value,
+                    }
+                    for key, value in batch
+                ]
+
+                stmt = insert(self._table_class).values(values_to_insert)
+                stmt = stmt.on_conflict_do_update(
+                    index_elements=["key", "namespace"],
+                    set_={"value": stmt.excluded.value},
                 )
 
-                # Prepare the raw SQL for bulk upsert
-                # Note: This SQL is PostgreSQL-specific. Adjust for other databases.
-                stmt = text(
-                    f"""
-                INSERT INTO {self.schema_name}.{self._table_class.__tablename__} (key, namespace, value)
-                VALUES {values_clause}
-                ON CONFLICT (key, namespace)
-                DO UPDATE SET
-                value = EXCLUDED.value;
-                """
-                )
-
-                # Flatten the list of tuples for execute parameters
-                params = {}
-                for i, (key, value) in enumerate(batch):
-                    params[f"key_{i}"] = key
-                    params[f"namespace_{i}"] = collection
-                    params[f"value_{i}"] = json.dumps(value)
-
-                # Execute the bulk upsert
-                await session.execute(stmt, params)
+                await session.execute(stmt)
                 await session.commit()
 
     def get(self, key: str, collection: str = DEFAULT_COLLECTION) -> Optional[dict]:

llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/pyproject.toml

@@ -28,7 +28,7 @@ dev = [
 
 [project]
 name = "llama-index-storage-kvstore-postgres"
-version = "0.4.1"
+version = "0.4.2"
 description = "llama-index kvstore postgres integration"
 authors = [{name = "Your Name", email = "you@example.com"}]
 requires-python = ">=3.9,<4.0"

llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-postgres/tests/test_storage_kvstore_postgres.py

@@ -1,5 +1,6 @@
 import pytest
 from importlib.util import find_spec
+from unittest.mock import MagicMock, patch
 from llama_index.core.storage.kvstore.types import BaseKVStore
 from llama_index.storage.kvstore.postgres import PostgresKVStore
 
@@ -51,3 +52,201 @@ def test_initialization():
     assert sum(errors) == 3
     assert pgstore4._engine is None
     assert pgstore4._async_engine is None
+
+
+@pytest.mark.skipif(
+    no_packages, reason="asyncpg, pscopg2-binary and sqlalchemy not installed"
+)
+def test_schema_creation_uses_safe_api():
+    import sqlalchemy
+
+    mock_engine = MagicMock()
+    mock_async_engine = MagicMock()
+
+    mock_session_instance = MagicMock()
+    mock_session_ctx = MagicMock()
+    mock_session_ctx.__enter__.return_value = mock_session_instance
+    mock_session_ctx.__exit__.return_value = None
+
+    mock_begin_ctx = MagicMock()
+    mock_begin_ctx.__enter__.return_value = MagicMock()
+    mock_begin_ctx.__exit__.return_value = None
+    mock_session_instance.begin.return_value = mock_begin_ctx
+
+    mock_session_factory = MagicMock(return_value=mock_session_ctx)
+
+    with (
+        patch.object(sqlalchemy, "create_engine", return_value=mock_engine),
+        patch.object(
+            sqlalchemy.ext.asyncio,
+            "create_async_engine",
+            return_value=mock_async_engine,
+        ),
+        patch("sqlalchemy.orm.sessionmaker", return_value=mock_session_factory),
+    ):
+        pgstore = PostgresKVStore(
+            table_name="test_table",
+            connection_string="postgresql://user:pass@localhost/db",
+            async_connection_string="postgresql+asyncpg://user:pass@localhost/db",
+            schema_name="test_schema",
+            perform_setup=False,
+        )
+
+        pgstore._connect()
+        pgstore._create_schema_if_not_exists()
+
+        execute_calls = mock_session_instance.execute.call_args_list
+        assert len(execute_calls) == 1
+
+        from sqlalchemy.schema import CreateSchema
+
+        executed_statement = execute_calls[0][0][0]
+        assert isinstance(executed_statement, CreateSchema)
+        assert executed_statement.element == "test_schema"
+        assert executed_statement.if_not_exists is True
+
+
+@pytest.mark.skipif(
+    no_packages, reason="asyncpg, pscopg2-binary and sqlalchemy not installed"
+)
+def test_put_all_uses_safe_insert():
+    import sqlalchemy
+
+    mock_engine = MagicMock()
+    mock_async_engine = MagicMock()
+
+    mock_session_instance = MagicMock()
+    mock_session_ctx = MagicMock()
+    mock_session_ctx.__enter__.return_value = mock_session_instance
+    mock_session_ctx.__exit__.return_value = None
+
+    mock_session_factory = MagicMock(return_value=mock_session_ctx)
+
+    with (
+        patch.object(sqlalchemy, "create_engine", return_value=mock_engine),
+        patch.object(
+            sqlalchemy.ext.asyncio,
+            "create_async_engine",
+            return_value=mock_async_engine,
+        ),
+        patch("sqlalchemy.orm.sessionmaker", return_value=mock_session_factory),
+    ):
+        pgstore = PostgresKVStore(
+            table_name="test_table",
+            connection_string="postgresql://user:pass@localhost/db",
+            async_connection_string="postgresql+asyncpg://user:pass@localhost/db",
+            schema_name="test_schema",
+            perform_setup=False,
+        )
+        pgstore._connect()
+        pgstore._is_initialized = True
+
+        test_data = [("key1", {"value": "data1"}), ("key2", {"value": "data2"})]
+        pgstore.put_all(test_data)
+
+        execute_calls = mock_session_instance.execute.call_args_list
+        assert len(execute_calls) >= 1
+
+        executed_statement = execute_calls[-1][0][0]
+        assert hasattr(executed_statement, "compile")
+
+
+@pytest.mark.skipif(
+    no_packages, reason="asyncpg, pscopg2-binary and sqlalchemy not installed"
+)
+@pytest.mark.asyncio
+async def test_aput_all_uses_safe_insert():
+    import sqlalchemy
+    from unittest.mock import AsyncMock
+
+    mock_engine = MagicMock()
+    mock_async_engine = MagicMock()
+
+    mock_session_instance = MagicMock()
+    mock_session_ctx_manager = MagicMock()
+    mock_session_ctx_manager.__aenter__ = AsyncMock(return_value=mock_session_instance)
+    mock_session_ctx_manager.__aexit__ = AsyncMock(return_value=None)
+
+    mock_async_session_factory = MagicMock(return_value=mock_session_ctx_manager)
+
+    with (
+        patch.object(sqlalchemy, "create_engine", return_value=mock_engine),
+        patch.object(
+            sqlalchemy.ext.asyncio,
+            "create_async_engine",
+            return_value=mock_async_engine,
+        ),
+        patch("sqlalchemy.orm.sessionmaker", return_value=mock_async_session_factory),
+    ):
+        pgstore = PostgresKVStore(
+            table_name="test_table",
+            connection_string="postgresql://user:pass@localhost/db",
+            async_connection_string="postgresql+asyncpg://user:pass@localhost/db",
+            schema_name="test_schema",
+            perform_setup=False,
+        )
+        pgstore._connect()
+        pgstore._is_initialized = True
+
+        mock_session_instance.execute = AsyncMock()
+        mock_session_instance.commit = AsyncMock()
+
+        test_data = [("key1", {"value": "data1"}), ("key2", {"value": "data2"})]
+        await pgstore.aput_all(test_data)
+
+        execute_calls = mock_session_instance.execute.call_args_list
+        assert len(execute_calls) >= 1
+
+        executed_statement = execute_calls[-1][0][0]
+        assert hasattr(executed_statement, "compile")
+
+
+@pytest.mark.skipif(
+    no_packages, reason="asyncpg, pscopg2-binary and sqlalchemy not installed"
+)
+def test_schema_name_with_special_characters():
+    import sqlalchemy
+
+    mock_engine = MagicMock()
+    mock_async_engine = MagicMock()
+
+    mock_session_instance = MagicMock()
+    mock_session_ctx = MagicMock()
+    mock_session_ctx.__enter__.return_value = mock_session_instance
+    mock_session_ctx.__exit__.return_value = None
+
+    mock_begin_ctx = MagicMock()
+    mock_begin_ctx.__enter__.return_value = MagicMock()
+    mock_begin_ctx.__exit__.return_value = None
+    mock_session_instance.begin.return_value = mock_begin_ctx
+
+    mock_session_factory = MagicMock(return_value=mock_session_ctx)
+
+    with (
+        patch.object(sqlalchemy, "create_engine", return_value=mock_engine),
+        patch.object(
+            sqlalchemy.ext.asyncio,
+            "create_async_engine",
+            return_value=mock_async_engine,
+        ),
+        patch("sqlalchemy.orm.sessionmaker", return_value=mock_session_factory),
+    ):
+        special_schema = "test'schema"
+        pgstore = PostgresKVStore(
+            table_name="test_table",
+            connection_string="postgresql://user:pass@localhost/db",
+            async_connection_string="postgresql+asyncpg://user:pass@localhost/db",
+            schema_name=special_schema,
+            perform_setup=False,
+        )
+
+        pgstore._connect()
+        pgstore._create_schema_if_not_exists()
+
+        execute_calls = mock_session_instance.execute.call_args_list
+        assert len(execute_calls) == 1
+
+        from sqlalchemy.schema import CreateSchema
+
+        executed_statement = execute_calls[0][0][0]
+        assert isinstance(executed_statement, CreateSchema)