Skip to content

Commit e7bad18

Browse files
committed
fix(persistence): register asyncpg JSON/JSONB codecs on PG engine
asyncpg returns JSON / JSONB columns as raw strings by default, which breaks every LangBot code path that expects a dict (mcp server loading, bot adapter config, pipeline config, etc.) when running on PostgreSQL. The asyncpg dialect's built-in on_connect hook should install the codecs, but it does not fire reliably on the SQLAlchemy 2.0 async engine path. Monkey-patch AsyncAdapt_asyncpg_dbapi.connect so that each new connection registers json / jsonb codecs eagerly via the wrapper's run_async hook.
1 parent bb89fe7 commit e7bad18

1 file changed

Lines changed: 56 additions & 0 deletions

File tree

src/langbot/pkg/persistence/databases/postgresql.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,66 @@
11
from __future__ import annotations
22

3+
import json
4+
35
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
46

57
from .. import database
68

79

10+
async def _register_json_codecs(pg_conn) -> None:
11+
"""Register asyncpg type codecs for JSON / JSONB so values come back as
12+
Python ``dict`` / ``list`` instead of raw ``str``.
13+
"""
14+
await pg_conn.set_type_codec(
15+
'json',
16+
encoder=json.dumps,
17+
decoder=json.loads,
18+
schema='pg_catalog',
19+
format='text',
20+
)
21+
await pg_conn.set_type_codec(
22+
'jsonb',
23+
encoder=json.dumps,
24+
decoder=json.loads,
25+
schema='pg_catalog',
26+
format='text',
27+
)
28+
29+
30+
def _patch_asyncpg_dialect() -> None:
31+
"""Wrap ``AsyncAdapt_asyncpg_dbapi.connect`` so that every new
32+
asyncpg-backed SQLAlchemy connection gets our JSON codecs installed
33+
eagerly, via the wrapper's ``run_async`` hook.
34+
"""
35+
from sqlalchemy.dialects.postgresql.asyncpg import AsyncAdapt_asyncpg_dbapi
36+
37+
_orig = AsyncAdapt_asyncpg_dbapi.connect
38+
39+
def _patched(self, *args, **kwargs):
40+
wrapper = _orig(self, *args, **kwargs)
41+
42+
# `AsyncAdapt_asyncpg_connection.run_async(fn)` invokes `fn(conn)`
43+
# where `conn` is the real asyncpg.Connection. The hook must accept
44+
# that argument; some SQLAlchemy versions forward it positionally,
45+
# others don't, so we tolerate both.
46+
async def _init(conn=None):
47+
pg = conn
48+
if pg is None:
49+
adapt = wrapper.connection
50+
pg = adapt._connection if hasattr(adapt, '_connection') else adapt
51+
await _register_json_codecs(pg)
52+
53+
wrapper.run_async(_init)
54+
return wrapper
55+
56+
AsyncAdapt_asyncpg_dbapi.connect = _patched
57+
58+
59+
# Apply the patch at module import time so subsequent engine creation
60+
# inherits the codec registration.
61+
_patch_asyncpg_dialect()
62+
63+
864
@database.manager_class('postgresql')
965
class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
1066
"""PostgreSQL database manager"""

0 commit comments

Comments
 (0)