fix(graphrag): compose graph SQL through psycopg instead of f-strings

Four graph-store queries formatted table and column names into the SQL
string: get_chunk_texts and delete_by_source (Bandit B608, alerts #582/#583
on main) and entity_pages/chunk_similarities from this branch (#661/#662,
dismissed). The names were validated by _safe_identifier, so none was
injectable, but each query was still a string built at runtime.

They are now fixed statements composed with psycopg.sql: identifiers go in as
sql.Identifier, values stay bound. Identifiers are lower-cased before quoting,
because PGVectorStore writes the same names unquoted and Postgres folds those
to lower case -- a quoted mixed-case name would address a different table.
Bandit reports nothing for docsgpt/graphrag now; the queries return the same
rows against a real graph as before.

docsgpt/graphrag/store.py

@@ -18,6 +18,7 @@ import uuid
 from typing import Any, Dict, List, Optional
 
 import psycopg
+from psycopg import sql
 from psycopg.types.json import Jsonb
 
 from docsgpt.core.settings import settings
@@ -52,6 +53,18 @@ def _safe_identifier(name: str) -> str:
     return name
 
 
+def _identifier(name: str) -> sql.Identifier:
+    """``name`` as a quoted identifier, folded the way Postgres folds it unquoted.
+
+    Composing identifiers through psycopg keeps every query a fixed statement
+    with bound values: nothing is formatted into the SQL string. The fold
+    matters because ``PGVectorStore`` writes these names unquoted, which
+    Postgres lower-cases, while a quoted identifier keeps its case; folding
+    first keeps both stores addressing the same table.
+    """
+    return sql.Identifier(_safe_identifier(name).lower())
+
+
 def _pgvector_identifiers() -> tuple[str, str, str, str]:
     """Resolve ``(table, text_col, metadata_col, source_col)`` from ``PGVectorStore``.
 
@@ -1228,18 +1241,25 @@ class GraphStore:
         cursor = conn.cursor()
         try:
             cursor.execute(
-                f"""
-                SELECT d.{metadata_col}, d.{text_col},
-                       (lower(n.name) = %s OR lower(n.name) LIKE %s) AS is_subject
-                FROM graph_node_chunks gc
-                JOIN graph_nodes n ON n.id = gc.node_id
-                JOIN {table} d ON d.id::text = gc.chunk_id
-                WHERE gc.source_id = %s AND d.{source_col} = %s
-                  AND (lower(n.name) = %s OR lower(n.name) LIKE %s OR n.name ILIKE %s)
-                GROUP BY d.{metadata_col}, d.{text_col}, is_subject
-                ORDER BY is_subject DESC, (d.{text_col} ILIKE %s) DESC
-                LIMIT %s;
-                """,
+                sql.SQL(
+                    """
+                    SELECT d.{metadata}, d.{text},
+                           (lower(n.name) = %s OR lower(n.name) LIKE %s) AS is_subject
+                    FROM graph_node_chunks gc
+                    JOIN graph_nodes n ON n.id = gc.node_id
+                    JOIN {table} d ON d.id::text = gc.chunk_id
+                    WHERE gc.source_id = %s AND d.{source} = %s
+                      AND (lower(n.name) = %s OR lower(n.name) LIKE %s OR n.name ILIKE %s)
+                    GROUP BY d.{metadata}, d.{text}, is_subject
+                    ORDER BY is_subject DESC, (d.{text} ILIKE %s) DESC
+                    LIMIT %s;
+                    """
+                ).format(
+                    metadata=_identifier(metadata_col),
+                    text=_identifier(text_col),
+                    table=_identifier(table),
+                    source=_identifier(source_col),
+                ),
                 (
                     clean.lower(), f"{clean.lower()} %",
                     source_id, source_id,
@@ -1277,11 +1297,17 @@ class GraphStore:
         cursor = conn.cursor()
         try:
             cursor.execute(
-                f"""
-                SELECT id::text, 1 - ({vector_col} <=> %s::vector)
-                FROM {table}
-                WHERE {source_col} = %s AND id::text = ANY(%s);
-                """,
+                sql.SQL(
+                    """
+                    SELECT id::text, 1 - ({vector} <=> %s::vector)
+                    FROM {table}
+                    WHERE {source} = %s AND id::text = ANY(%s);
+                    """
+                ).format(
+                    vector=_identifier(vector_col),
+                    table=_identifier(table),
+                    source=_identifier(source_col),
+                ),
                 (query_embedding, source_id, [str(c) for c in chunk_ids]),
             )
             return {row[0]: float(row[1]) for row in cursor.fetchall()}
@@ -1312,10 +1338,17 @@ class GraphStore:
         cursor = conn.cursor()
         try:
             cursor.execute(
-                f"""
-                SELECT id, {text_col}, {metadata_col} FROM {table}
-                WHERE {source_col} = %s AND id::text = ANY(%s);
-                """,
+                sql.SQL(
+                    """
+                    SELECT id, {text}, {metadata} FROM {table}
+                    WHERE {source} = %s AND id::text = ANY(%s);
+                    """
+                ).format(
+                    text=_identifier(text_col),
+                    metadata=_identifier(metadata_col),
+                    table=_identifier(table),
+                    source=_identifier(source_col),
+                ),
                 (source_id, [str(c) for c in chunk_ids]),
             )
             return {
@@ -1501,7 +1534,10 @@ class GraphStore:
                 "graph_ingest_progress",
             ):
                 cursor.execute(
-                    f"DELETE FROM {table} WHERE source_id = %s;", (source_id,)
+                    sql.SQL("DELETE FROM {} WHERE source_id = %s;").format(
+                        sql.Identifier(table)
+                    ),
+                    (source_id,),
                 )
             conn.commit()
         except Exception as e:

tests/graphrag/test_store.py

@@ -558,16 +558,23 @@ class TestGraphStoreParameterization:
         return store, cursor
 
     def test_delete_by_source_binds_source_id(self):
+        from psycopg import sql as pgsql
+
         store, cursor = self._store_with_mock_conn()
         sid = str(uuid.uuid4())
         store.delete_by_source(sid)
 
+        tables = []
         for call in cursor.execute.call_args_list:
-            sql = call.args[0]
+            query = call.args[0]
             params = call.args[1] if len(call.args) > 1 else None
+            assert isinstance(query, pgsql.Composable)
+            sql = query.as_string()
             assert "WHERE source_id = %s" in sql
             assert sid not in sql
             assert params == (sid,)
+            tables.append(sql.split('"')[1])
+        assert tables == ["graph_node_chunks", "graph_edges", "graph_nodes", "graph_ingest_progress"]
 
     def test_search_binds_embedding_and_source(self):
         store, cursor = self._store_with_mock_conn()
@@ -636,6 +643,14 @@ class TestGraphReadQueries:
         store._get_connection = lambda: conn
         return store, cursor, conn
 
+    def test_identifiers_are_quoted_as_postgres_folds_them_unquoted(self):
+        # PGVectorStore writes these names unquoted, which Postgres folds to
+        # lower case; quoting keeps case, so the fold happens first or the two
+        # stores would address different tables.
+        assert store_module._identifier("Documents").as_string() == '"documents"'
+        with pytest.raises(ValueError):
+            store_module._identifier('documents"; DROP TABLE graph_nodes; --')
+
     def test_fact_seeds_bind_every_value_and_read_weight_as_distance(self):
         store, cursor, _ = self._store(rows=[("n1", "Quill", "a store", 0.8), ("n2", "Alder", None, None)])
         sid = str(uuid.uuid4())
@@ -675,7 +690,9 @@ class TestGraphReadQueries:
 
         pages = store.entity_pages(sid, "Quill", limit=0)
 
-        sql, params = cursor.execute.call_args.args
+        query, params = cursor.execute.call_args.args
+        sql = query.as_string()
+        assert 'JOIN "documents" d' in sql and 'd."source_id" = %s' in sql
         assert "Quill" not in sql
         # Exact name, name plus a qualifier ("Quill Store"), substring fallback,
         # text-opens-with ordering, then the clamped limit.
@@ -689,7 +706,9 @@ class TestGraphReadQueries:
 
         scores = store.chunk_similarities(sid, [11, "12"], embedding)
 
-        sql, params = cursor.execute.call_args.args
+        query, params = cursor.execute.call_args.args
+        sql = query.as_string()
+        assert '1 - ("embedding" <=> %s::vector)' in sql and 'FROM "documents"' in sql
         assert "= ANY(%s)" in sql and sid not in sql
         assert params == (embedding, sid, ["11", "12"])
         assert scores == {"11": 0.75}

tests/retriever/test_graph_rag.py

@@ -472,11 +472,16 @@ class TestGetChunkTexts:
         sid = str(uuid.uuid4())
         store.get_chunk_texts(sid, ["1", "2"])
 
-        sql, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1]
-        assert f"FROM {table}" in sql
-        assert text_col in sql
-        assert metadata_col in sql
-        assert f"{source_col} = %s" in sql
+        from psycopg import sql as pgsql
+
+        query, params = cursor.execute.call_args.args[0], cursor.execute.call_args.args[1]
+        # Identifiers are composed and quoted by psycopg, never formatted in.
+        assert isinstance(query, pgsql.Composable)
+        sql = query.as_string()
+        assert f'FROM "{table}"' in sql
+        assert f'"{text_col}"' in sql
+        assert f'"{metadata_col}"' in sql
+        assert f'"{source_col}" = %s' in sql
         assert "id::text = ANY(%s)" in sql
         assert sid not in sql
         assert params == (sid, ["1", "2"])