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})
    )
""")

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}>
    )
""")

#!/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.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 main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-root", type=Path, default=Path.cwd())
    args = parser.parse_args()
    source_root = args.source_root.resolve()

    output = {
        "source": {
            "repository": "MervinPraison/PraisonAI",
            "head": git_value(source_root, "rev-parse", "HEAD"),
            "describe": git_value(source_root, "describe", "--tags", "--always", "--dirty"),
        },
        "pgvector": run_pgvector(source_root),
        "cassandra": run_cassandra(source_root),
    }

    assert output["pgvector"]["invalid_collection_control"].startswith("ValueError:"), output
    assert output["cassandra"]["invalid_collection_control"].startswith("ValueError:"), output
    assert output["pgvector"]["safe_contains_drop_table"] is False, output
    assert output["cassandra"]["safe_contains_drop_table"] is False, output
    assert output["pgvector"]["malicious_contains_drop_table"] is True, output
    assert output["cassandra"]["malicious_contains_drop_table"] is True, output

    print(json.dumps(output, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()

def validate_vector_dimension(value: object) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError("dimension must be an integer")
    if value <= 0 or value > 200000:
        raise ValueError("dimension is outside the supported range")
    return value