From e49f3e9c02dc021a3d01f0e7dab40deaba74ff5a Mon Sep 17 00:00:00 2001 From: Kwizii Date: Sat, 22 Aug 2026 17:46:47 +0800 Subject: [PATCH] fix: support pagination and filters in Postgres export_graph export_graph silently ignored page/page_size/memory_type/status/filter (they were swallowed by **kwargs), so callers always received the full node list with total_nodes equal to the page length. This broke pagination and tag filtering for every consumer of POST /get_memory when the graph backend is postgres. - translate page/page_size into LIMIT/OFFSET (normalized like neo4j) - filter memory_type via properties->>'memory_type' = ANY(...) - default status filter excludes 'deleted'; explicit status list uses ANY - reuse _build_filter_where_clause so and/or/contains filter DSL works - count total_nodes with a separate COUNT query before pagination - resolve edges via subqueries over the full filtered node set (not page-local ids) so the edge set is complete and consistent across pages --- src/memos/graph_dbs/postgres.py | 127 +++++++++++--- tests/graph_dbs/test_postgres_export_graph.py | 156 ++++++++++++++++++ 2 files changed, 256 insertions(+), 27 deletions(-) create mode 100644 tests/graph_dbs/test_postgres_export_graph.py diff --git a/src/memos/graph_dbs/postgres.py b/src/memos/graph_dbs/postgres.py index 594f7e695..6234cd746 100644 --- a/src/memos/graph_dbs/postgres.py +++ b/src/memos/graph_dbs/postgres.py @@ -1100,48 +1100,121 @@ def clear(self, user_name: str | None = None) -> None: finally: self._put_conn(conn) - def export_graph(self, include_embedding: bool = False, **kwargs) -> dict[str, Any]: - """Export all data.""" + def export_graph( + self, + include_embedding: bool = False, + page: int | None = None, + page_size: int | None = None, + memory_type: list[str] | None = None, + status: list[str] | None = None, + filter: dict | None = None, + **kwargs, + ) -> dict[str, Any]: + """Export nodes and edges with optional pagination and filtering. + + Args: + include_embedding (bool): Whether to include embedding fields in node metadata. + page (int, optional): Page number (starts from 1). If None, exports all data without pagination. + page_size (int, optional): Number of items per page. If None, exports all data without pagination. + memory_type (list[str], optional): Only export nodes whose ``properties->>'memory_type'`` is in this list. + status (list[str], optional): If not provided, only nodes with status != 'deleted' are exported. + If a non-empty list is provided, only nodes whose status is in this list are exported. + filter (dict, optional): Filter conditions with 'and'/'or' logic, same as :meth:`get_all_memory_items`. + **kwargs: Additional keyword arguments, including: + - user_name (str, optional): User name for filtering. + + Returns: + { + "nodes": [ { "id": ..., "memory": ..., "metadata": {...} }, ... ], + "edges": [ { "source": ..., "target": ..., "type": ... }, ... ], + "total_nodes": int, # Total number of nodes matching the filters (ignores pagination) + } + + Edges are not paginated: every response includes all edges connected to any + node matching the filters (resolved via subqueries over the full filtered + node set), so iterating pages yields a complete edge set. + """ user_name = kwargs.get("user_name") or self.user_name + + use_pagination = page is not None and page_size is not None + if use_pagination: + if page < 1: + page = 1 + if page_size < 1: + page_size = 10 + offset = (page - 1) * page_size + + conditions = ["user_name = %s"] + params: list[Any] = [user_name] + + if memory_type: + conditions.append("properties->>'memory_type' = ANY(%s)") + params.append(list(memory_type)) + + if status is None: + conditions.append( + "(properties->>'status' <> 'deleted' OR properties->>'status' IS NULL)" + ) + elif len(status) > 0: + conditions.append("properties->>'status' = ANY(%s)") + params.append(list(status)) + + filter_clause = self._build_filter_where_clause(filter, params) + if filter_clause: + conditions.append(filter_clause) + + where_clause = " AND ".join(conditions) + conn = self._get_conn() try: with conn.cursor() as cur: - # Get nodes + # Count total matching nodes before pagination + cur.execute( + f""" + SELECT COUNT(*) FROM {self.schema}.memories + WHERE {where_clause} + """, + params, + ) + total_nodes = cur.fetchone()[0] + + # Get nodes (paginated) cols = "id, memory, properties, created_at, updated_at" if include_embedding: cols += ", embedding" + query = f""" + SELECT {cols} FROM {self.schema}.memories + WHERE {where_clause} + ORDER BY created_at DESC, id DESC + """ + query_params = list(params) + if use_pagination: + query += " LIMIT %s OFFSET %s" + query_params.extend([page_size, offset]) + cur.execute(query, query_params) + nodes = [self._parse_row(row, include_embedding) for row in cur.fetchall()] + + # Get all edges connected to any node matching the filters. + # Deliberately not paginated: resolving endpoints against the full + # filtered node set keeps the edge set complete and consistent across + # pages (page-local ids would drop or duplicate cross-page edges). cur.execute( f""" - SELECT {cols} FROM {self.schema}.memories - WHERE user_name = %s - ORDER BY created_at DESC + SELECT source_id, target_id, edge_type + FROM {self.schema}.edges + WHERE source_id IN (SELECT id FROM {self.schema}.memories WHERE {where_clause}) + OR target_id IN (SELECT id FROM {self.schema}.memories WHERE {where_clause}) """, - (user_name,), + [*params, *params], ) - nodes = [self._parse_row(row, include_embedding) for row in cur.fetchall()] - - # Get edges - node_ids = [n["id"] for n in nodes] - if node_ids: - cur.execute( - f""" - SELECT source_id, target_id, edge_type - FROM {self.schema}.edges - WHERE source_id = ANY(%s) OR target_id = ANY(%s) - """, - (node_ids, node_ids), - ) - edges = [ - {"source": row[0], "target": row[1], "type": row[2]} - for row in cur.fetchall() - ] - else: - edges = [] + edges = [ + {"source": row[0], "target": row[1], "type": row[2]} for row in cur.fetchall() + ] return { "nodes": nodes, "edges": edges, - "total_nodes": len(nodes), + "total_nodes": total_nodes, "total_edges": len(edges), } finally: diff --git a/tests/graph_dbs/test_postgres_export_graph.py b/tests/graph_dbs/test_postgres_export_graph.py new file mode 100644 index 000000000..e543402a1 --- /dev/null +++ b/tests/graph_dbs/test_postgres_export_graph.py @@ -0,0 +1,156 @@ +"""Unit tests for PostgresGraphDB.export_graph pagination and filtering.""" + +from __future__ import annotations + +import json + +from datetime import datetime +from typing import Any +from unittest.mock import MagicMock, patch + +from memos.graph_dbs.postgres import PostgresGraphDB + + +SCHEMA = "test_schema" +USER_NAME = "user-1" + + +def _row(node_id: str, memory: str, memory_type: str = "UserMemory"): + props = {"memory_type": memory_type, "tags": ["写作风格"]} + return ( + node_id, + memory, + json.dumps(props), + datetime(2026, 8, 22, 12, 0, 0), + datetime(2026, 8, 22, 12, 0, 0), + ) + + +class _FakeCursor: + """Cursor that serves queued results in order and records executed SQL.""" + + def __init__(self, results: list[Any]): + self.results = list(results) + self.executed: list[tuple[str, list]] = [] + + def execute(self, query: str, params=None): + self.executed.append((query, list(params or []))) + self._current = self.results.pop(0) + + def fetchone(self): + return self._current + + def fetchall(self): + return self._current + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + +class _FakeConn: + def __init__(self, results: list[Any]): + self.cursor_obj = _FakeCursor(results) + + def cursor(self): + return self.cursor_obj + + +def _build_db(results: list[Any]) -> tuple[PostgresGraphDB, _FakeCursor]: + with ( + patch("memos.graph_dbs.postgres.require_python_package", lambda **kwargs: lambda fn: fn), + patch("psycopg2.pool.ThreadedConnectionPool", MagicMock()), + patch.object(PostgresGraphDB, "_init_schema", lambda self: None), + ): + config = MagicMock() + config.schema_name = SCHEMA + config.user_name = USER_NAME + db = PostgresGraphDB(config) + conn = _FakeConn(results) + db._get_conn = lambda: conn + return db, conn.cursor_obj + + +def test_export_graph_paginates_with_total_count() -> None: + """page/page_size must translate to LIMIT/OFFSET and total comes from COUNT, not page size.""" + db, cursor = _build_db( + results=[[9], [_row("n7", "m7"), _row("n8", "m8"), _row("n9", "m9")], []] + ) + + result = db.export_graph(page=2, page_size=6) + + count_sql, _count_params = cursor.executed[0] + assert "COUNT(*)" in count_sql + + select_sql, select_params = cursor.executed[1] + assert "LIMIT %s OFFSET %s" in select_sql + assert select_params[-2:] == [6, 6] + + assert result["total_nodes"] == 9 + assert len(result["nodes"]) == 3 + + +def test_export_graph_edges_resolve_against_full_filtered_set() -> None: + """Edges must resolve via subqueries over the full filtered node set, not page-local ids.""" + db, cursor = _build_db( + results=[ + [9], + [_row("n7", "m7"), _row("n8", "m8"), _row("n9", "m9")], + [("n7", "n1", "rel")], + ] + ) + + result = db.export_graph(page=2, page_size=6, filter={"tags": {"contains": "写作风格"}}) + + edges_sql, edges_params = cursor.executed[2] + assert edges_sql.count("SELECT id FROM test_schema.memories WHERE") == 2 + assert "source_id IN" in edges_sql and "target_id IN" in edges_sql + # filter params are applied to both subqueries (user_name + tags containment, twice) + assert edges_params.count(json.dumps(["写作风格"])) == 2 + assert result["edges"] == [{"source": "n7", "target": "n1", "type": "rel"}] + assert result["total_edges"] == 1 + + +def test_export_graph_without_pagination_returns_all() -> None: + db, cursor = _build_db(results=[[2], [_row("n1", "m1"), _row("n2", "m2")], []]) + + result = db.export_graph() + + select_sql, _select_params = cursor.executed[1] + assert "LIMIT" not in select_sql + assert result["total_nodes"] == 2 + assert len(result["nodes"]) == 2 + + +def test_export_graph_applies_memory_type_and_status_filters() -> None: + db, cursor = _build_db(results=[[0], [], []]) + + db.export_graph(memory_type=["UserMemory", "LongTermMemory"]) + + count_sql, count_params = cursor.executed[0] + assert "properties->>'memory_type' = ANY(%s)" in count_sql + assert ["UserMemory", "LongTermMemory"] in count_params + assert "properties->>'status' <> 'deleted'" in count_sql + + +def test_export_graph_applies_tag_filter() -> None: + """Tag filter (contains) must reach the SQL WHERE clause as a jsonb containment check.""" + db, cursor = _build_db(results=[[1], [_row("n1", "m1")], []]) + + db.export_graph(filter={"tags": {"contains": "写作风格"}}) + + count_sql, count_params = cursor.executed[0] + assert "properties->'tags' @> %s::jsonb" in count_sql + assert json.dumps(["写作风格"]) in count_params + + +def test_export_graph_invalid_pagination_inputs_are_normalized() -> None: + db, cursor = _build_db(results=[[0], [], []]) + + db.export_graph(page=0, page_size=-3) + + select_sql, select_params = cursor.executed[1] + assert "LIMIT %s OFFSET %s" in select_sql + assert select_params[-2:] == [10, 0]