WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore: pgqs (#6567)

Browse Source 
Add more argument sanitization
This commit is contained in:
William FH
2025-12-09 14:51:29 -08:00
committed by GitHub
parent b70d5aac0e
commit e86b5f4da2
5 changed files with 85 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py
+5 -3
View File
@@ -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
libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py
+5 -3
View File
@@ -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
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
+10 -6
View File
@@ -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"],
libs/checkpoint-postgres/langgraph/store/postgres/aio.py
+21
View File
@@ -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(
libs/checkpoint-postgres/langgraph/store/postgres/base.py
+44 -14
View File
@@ -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,))