perf: route native JSON columns through JSONCodec instead of stdlib json (#28396)

JSONField serializes with JSONCodec, but columns declared as SQLAlchemy's own JSON
type go through the engine's serializer instead, and no engine set one. That left
Chat.chat - the largest blob the app stores - on stdlib json.dumps/loads no matter
what ENABLE_ORJSON was set to, while the rest of the app used the codec. SQLAlchemy
invokes it once per write and once per read, so every chat read and write paid a
full stdlib pass over the whole conversation on top of whatever the caller did.

Both engine constructors are now wrapped so the codec is wired in by default and
cannot be missed by a call site that forgets it; an explicit json_serializer still
wins. The 10 create_engine/create_async_engine calls in this module go through the
wrappers. Vector-store engines (pgvector, mariadb, opengauss) are separate databases
and are left alone.

Serializing and deserializing chat-shaped blobs, median of 11 runs:

| chat blob | write | read |
| --- | --- | --- |
| 600 msgs (2.8 MB) | 10.1 -> 1.7 ms | 8.4 -> 3.7 ms |
| 3000 msgs (14.2 MB) | 51.9 -> 8.0 ms | 48.5 -> 27.8 ms |
| 6000 msgs (28.5 MB) | 105.9 -> 29.5 ms | 112.7 -> 80.5 ms |

With ENABLE_ORJSON off JSONCodec is stdlib json, so this is a no-op until the flag
is set - the change cannot regress a default deployment.

With it on, a round-trip probe through a native JSON column returns objects equal to
the stdlib ones on all 12 shapes tried: ASCII, CJK, emoji, astral-plane, unicode
keys, null bytes, lone surrogates, floats, ints above 2**63 and 2**64, line
separators, empty and deeply nested. Stored text changes for non-ASCII, which is
written as raw UTF-8 rather than backslash-uXXXX escapes and is correspondingly
smaller. Nothing queries that text by escape except two Postgres safety filters in
chats.py, and both still hold: a null byte is escaped identically by both codecs,
and the title filter reads a text column rather than JSON. The ->> and json_extract
searches decode the string before matching, so escaping cannot reach them.

Two differences are inherent to JSONCodec and already apply to every JSONField
column: ints beyond 2**64-1 come back as float, and NaN/Infinity serialize to null
rather than the bare literals stdlib emits - the latter being invalid JSON that a
Postgres json column rejects today. Neither shape occurs in chat blobs. Alembic
builds its own engine and stays on stdlib, which is fine in both directions since
each codec reads the other's output.


Claude-Session: https://claude.ai/code/session_014BXoM6QiFJKisxcxKAXii8

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Classic298
2026-08-10 19:39:49 -05:00
committed by GitHub
co-authored by Claude
parent 8fbfd14a8b
commit 8d6a7c8308
+31 -10
View File
@@ -232,6 +232,27 @@ def _make_async_url(url: str) -> str:
return url
def _json_codec_kwargs(kwargs: dict) -> dict:
"""Default an engine to JSONCodec for native ``JSON`` columns.
Unlike ``JSONField``, those serialize through the engine, which otherwise uses
stdlib ``json``. With ``ENABLE_ORJSON`` off JSONCodec is stdlib ``json`` anyway.
"""
kwargs.setdefault('json_serializer', JSONCodec.dumps)
kwargs.setdefault('json_deserializer', JSONCodec.loads)
return kwargs
def _create_engine(*args, **kwargs):
"""``create_engine`` with the app JSON codec wired in."""
return create_engine(*args, **_json_codec_kwargs(kwargs))
def _create_async_engine(*args, **kwargs):
"""``create_async_engine`` with the app JSON codec wired in."""
return create_async_engine(*args, **_json_codec_kwargs(kwargs))
# ============================================================
# SYNC ENGINE (used only for: startup migrations, config loading,
# Alembic, peewee migration, health checks)
@@ -260,7 +281,7 @@ if SQLALCHEMY_DATABASE_URL.startswith('sqlite+sqlcipher://'):
# in the native sqlcipher3 C library. Use NullPool by default for safety,
# or QueuePool if DATABASE_POOL_SIZE is explicitly configured.
if isinstance(DATABASE_POOL_SIZE, int) and DATABASE_POOL_SIZE > 0:
engine = create_engine(
engine = _create_engine(
'sqlite://',
creator=create_sqlcipher_connection,
pool_size=DATABASE_POOL_SIZE,
@@ -272,7 +293,7 @@ if SQLALCHEMY_DATABASE_URL.startswith('sqlite+sqlcipher://'):
echo=False,
)
else:
engine = create_engine(
engine = _create_engine(
'sqlite://',
creator=create_sqlcipher_connection,
poolclass=NullPool,
@@ -282,7 +303,7 @@ if SQLALCHEMY_DATABASE_URL.startswith('sqlite+sqlcipher://'):
log.info('Connected to encrypted SQLite database using SQLCipher')
elif 'sqlite' in SQLALCHEMY_DATABASE_URL:
engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={'check_same_thread': False})
engine = _create_engine(SQLALCHEMY_DATABASE_URL, connect_args={'check_same_thread': False})
def _apply_sqlite_pragmas(dbapi_connection):
"""Apply all configured SQLite PRAGMAs to a raw DBAPI connection."""
@@ -314,7 +335,7 @@ elif 'sqlite' in SQLALCHEMY_DATABASE_URL:
else:
if isinstance(DATABASE_POOL_SIZE, int):
if DATABASE_POOL_SIZE > 0:
engine = create_engine(
engine = _create_engine(
SQLALCHEMY_DATABASE_URL,
pool_size=DATABASE_POOL_SIZE,
max_overflow=DATABASE_POOL_MAX_OVERFLOW,
@@ -324,9 +345,9 @@ else:
poolclass=QueuePool,
)
else:
engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, poolclass=NullPool)
engine = _create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, poolclass=NullPool)
else:
engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
engine = _create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
enable_iam_token_auth(engine)
@@ -373,7 +394,7 @@ if 'sqlite' in ASYNC_SQLALCHEMY_DATABASE_URL:
# No pool_pre_ping: a local SQLite file cannot drop connections, and the
# ping costs a worker-thread hop plus a SELECT 1 on every checkout.
_sqlite_pool_size = DATABASE_POOL_SIZE if isinstance(DATABASE_POOL_SIZE, int) and DATABASE_POOL_SIZE > 0 else 512
async_engine = create_async_engine(
async_engine = _create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
connect_args={'check_same_thread': False},
pool_size=_sqlite_pool_size,
@@ -387,7 +408,7 @@ if 'sqlite' in ASYNC_SQLALCHEMY_DATABASE_URL:
else:
if isinstance(DATABASE_POOL_SIZE, int):
if DATABASE_POOL_SIZE > 0:
async_engine = create_async_engine(
async_engine = _create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
pool_size=DATABASE_POOL_SIZE,
max_overflow=DATABASE_POOL_MAX_OVERFLOW,
@@ -396,13 +417,13 @@ else:
pool_pre_ping=True,
)
else:
async_engine = create_async_engine(
async_engine = _create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
pool_pre_ping=True,
poolclass=NullPool,
)
else:
async_engine = create_async_engine(
async_engine = _create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
pool_pre_ping=True,
)