mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-16 10:15:00 +00:00
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.
This commit is contained in:
1 parent
558f803fed
commit
b3d7e63ae3
3 files changed
+90
-30
No files matched your search
@@ -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}
|
||||
|
||||
@@ -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"])
|
||||
|
||||
Reference in new issue
Block a user