From e86b5f4da21f0fa00bec19fe0e86d031dba0b004 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Tue, 9 Dec 2025 14:51:29 -0800 Subject: [PATCH] chore: pgqs (#6567) Add more argument sanitization --- .../langgraph/checkpoint/postgres/__init__.py | 8 ++- .../langgraph/checkpoint/postgres/aio.py | 8 ++- .../langgraph/checkpoint/postgres/shallow.py | 16 +++-- .../langgraph/store/postgres/aio.py | 21 +++++++ .../langgraph/store/postgres/base.py | 58 ++++++++++++++----- 5 files changed, 85 insertions(+), 26 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 01816c8f2..600127333 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -143,11 +143,13 @@ class PostgresSaver(BasePostgresSaver): """ where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" - if limit: - query += f" LIMIT {limit}" + params = list(args) + if limit is not None: + query += " LIMIT %s" + params.append(int(limit)) # if we change this to use .stream() we need to make sure to close the cursor with self._cursor() as cur: - cur.execute(query, args) + cur.execute(query, params) values = cur.fetchall() if not values: return diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 1d507b893..5cd533f8d 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -132,11 +132,13 @@ class AsyncPostgresSaver(BasePostgresSaver): """ where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC" - if limit: - query += f" LIMIT {limit}" + params = list(args) + if limit is not None: + query += " LIMIT %s" + params.append(int(limit)) # if we change this to use .stream() we need to make sure to close the cursor async with self._cursor() as cur: - await cur.execute(query, args, binary=True) + await cur.execute(query, params, binary=True) values = await cur.fetchall() if not values: return diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py index d3d5c8f0b..90fc95e74 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py @@ -272,10 +272,12 @@ class ShallowPostgresSaver(BasePostgresSaver): """ where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" + params = list(args) + if limit is not None: + query += " LIMIT %s" + params.append(int(limit)) with self._cursor() as cur: - cur.execute(self.SELECT_SQL + where, args, binary=True) + cur.execute(query, params, binary=True) for value in cur: checkpoint: Checkpoint = { **value["checkpoint"], @@ -636,10 +638,12 @@ class AsyncShallowPostgresSaver(BasePostgresSaver): """ where, args = self._search_where(config, filter, before) query = self.SELECT_SQL + where - if limit: - query += f" LIMIT {limit}" + params = list(args) + if limit is not None: + query += " LIMIT %s" + params.append(int(limit)) async with self._cursor() as cur: - await cur.execute(self.SELECT_SQL + where, args, binary=True) + await cur.execute(query, params, binary=True) async for value in cur: checkpoint: Checkpoint = { **value["checkpoint"], diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index d2f38246e..763133a9a 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -266,6 +266,27 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con k: v(self) if v is not None and callable(v) else v for k, v in migration.params.items() } + if "dims" in params: + try: + params["dims"] = int(params["dims"]) + except Exception as e: + raise ValueError( + f"Invalid dims for vector index: {params['dims']}" + ) from e + if "vector_type" in params: + vt = str(params["vector_type"]) + if vt not in ("vector", "halfvec"): + raise ValueError( + f"Invalid vector_type for pgvector: {vt}" + ) + params["vector_type"] = vt + if "index_type" in params: + it = str(params["index_type"]) + if it not in ("hnsw", "ivfflat"): + raise ValueError( + f"Invalid index_type for pgvector: {it}" + ) + params["index_type"] = it sql = sql % params await cur.execute(sql) await cur.execute( diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 825590e23..7bd8f0683 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -327,31 +327,36 @@ class BasePostgresStore(Generic[C]): embedding_request: tuple[str, Sequence[tuple[str, str, str, str]]] | None = None if inserts: values = [] - insertion_params = [] + insertion_params: list[Any] = [] vector_values = [] embedding_request_params = [] # Handle TTL expiration # First handle main store insertions for op in inserts: - if op.ttl is not None: - expires_at_str = f"NOW() + INTERVAL '{op.ttl * 60} seconds'" - ttl_minutes = op.ttl - else: - expires_at_str = "NULL" - ttl_minutes = None - - values.append( - f"(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, {expires_at_str}, %s)" - ) insertion_params.extend( - [ + ( _namespace_to_text(op.namespace), op.key, Jsonb(cast(dict, op.value)), - ttl_minutes, - ] + ) ) + if op.ttl is not None: + values.append( + "(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, NOW() + %s::interval, %s)" + ) + ttl_minutes = float(op.ttl) + insertion_params.extend( + ( + f"{ttl_minutes * 60} seconds", + ttl_minutes, + ) + ) + else: + values.append( + "(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, NULL, %s)" + ) + insertion_params.append(None) # Then handle embeddings if configured if self.index_config: @@ -465,6 +470,10 @@ class BasePostgresStore(Generic[C]): cast(dict, self.index_config)["dims"], ) else: + if vector_type not in ("vector", "halfvec"): + raise ValueError( + f"Invalid vector_type for pgvector: {vector_type}" + ) score_operator = score_operator % ("%s", vector_type) vectors_per_doc_estimate = cast(dict, self.index_config)[ @@ -1122,6 +1131,27 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): k: v(self) if v is not None and callable(v) else v for k, v in migration.params.items() } + if "dims" in params: + try: + params["dims"] = int(params["dims"]) + except Exception as e: + raise ValueError( + f"Invalid dims for vector index: {params['dims']}" + ) from e + if "vector_type" in params: + vt = str(params["vector_type"]) + if vt not in ("vector", "halfvec"): + raise ValueError( + f"Invalid vector_type for pgvector: {vt}" + ) + params["vector_type"] = vt + if "index_type" in params: + it = str(params["index_type"]) + if it not in ("hnsw", "ivfflat"): + raise ValueError( + f"Invalid index_type for pgvector: {it}" + ) + params["index_type"] = it sql = sql % params cur.execute(sql) cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))