From b3d7e63ae32228e6914c5543d86e23e3d260e597 Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 19 Sep 2026 20:10:54 +0100 Subject: [PATCH] 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 | 80 ++++++++++++++++++++++--------- tests/graphrag/test_store.py | 25 ++++++++-- tests/retriever/test_graph_rag.py | 15 ++++-- 3 files changed, 90 insertions(+), 30 deletions(-) diff --git a/docsgpt/graphrag/store.py b/docsgpt/graphrag/store.py index 134998ba..ab2b143a 100644 --- a/docsgpt/graphrag/store.py +++ b/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: diff --git a/tests/graphrag/test_store.py b/tests/graphrag/test_store.py index 27cfddc2..355bed09 100644 --- a/tests/graphrag/test_store.py +++ b/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} diff --git a/tests/retriever/test_graph_rag.py b/tests/retriever/test_graph_rag.py index e4d1bf46..71c996d5 100644 --- a/tests/retriever/test_graph_rag.py +++ b/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"])