SQLAlchemy 中 JSON 列的自定义 json 序列化程序

Posted 2023-03-16

【中文标题】SQLAlchemy 中 JSON 列的自定义 json 序列化程序

【英文标题】：Custom json serializer for JSON column in SQLAlchemy

【发布时间】：2021-09-21 18:56:51

【问题描述】：

我有以下 ORM 对象（简化）：

import datetime as dt

from sqlalchemy import create_engine, Integer, Column, DateTime
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session, declarative_base

Base = declarative_base()

class Metrics(Base):
    __tablename__ = 'metrics'

    id = Column(Integer, primary_key=True)
    ts = Column(DateTime, default=dt.datetime.now())
    computed_values = Column(JSONB)
    dates = Column(JSONB)

    calculated_values = 
    entry = Metrics(computed_values='foo': 12.3, 'bar':45.6,
                    dates=[datetime.date.today()])

    engine = create_engine('postgresql://postgres:postgres@localhost:5432/my_schema')
    with Session(engine, future=True) as session:
        session.add(entry)
        session.commit()

每一行有：

id主键 ts 插入行时的时间戳 computed_values 要存储的实际 JSONB 数据 dates JSONB 存储计算数据的日期列表。

虽然我对 computed_values 列没有任何问题，但默认情况下，SQLAlchemy JSON 序列化程序无法序列化 dates 列内列表中的 datetime.date 对象。

我的想法是为该确切列重新定义 date 对象的序列化程序行为。为此，我必须定义自己的自定义 JSON 序列化程序，或者使用一些现成的序列化程序，例如 orjson。由于我可能会在项目中遇到许多其他 JSON 序列化问题，所以我更喜欢后者。

深入了解JSONB 类及其超类，我认为以下应该可以解决问题：

class Metrics(Base):
    __tablename__ = 'metrics'

    # ---%<--- snip ---%<---
    dates = Column(JSONB(json_serializer=lambda obj: orjson.dumps(obj, option=orjson.OPT_NAIVE_UTC)))

    # ---%<--- snip ---%<---

但它没有：

File "metrics.py", line 30, in Metrics
    dates = Column(JSONB(json_serializer=lambda obj: orjson.dumps(obj, option=orjson.OPT_NAIVE_UTC)))
TypeError: __init__() got an unexpected keyword argument 'json_serializer'

我做错了什么以及如何为 JSON（和 JSONB）列正确定义自定义 SQLAlchemy 序列化程序？

【参考方案1】：

看起来你应该能够通过修改你的create_engine 语句来得到你想要的。

来自the docstring in SQLAlchemy：

Custom serializers and deserializers are specified at the dialect level,
that is using :func:`_sa.create_engine`.  The reason for this is that when
using psycopg2, the DBAPI only allows serializers at the per-cursor
or per-connection level.   E.g.::
    engine = create_engine("postgresql://scott:tiger@localhost/test",
                            json_serializer=my_serialize_fn,
                            json_deserializer=my_deserialize_fn
                    )

所以生成的代码应该如下：

import datetime as dt

import orjson

from sqlalchemy import create_engine, Integer, Column, DateTime
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Session, declarative_base

Base = declarative_base()

class Metrics(Base):
    __tablename__ = 'metrics'

    id = Column(Integer, primary_key=True)
    ts = Column(DateTime, default=dt.datetime.now())
    computed_values = Column(JSONB)
    dates = Column(JSONB)

    calculated_values = 
    entry = Metrics(computed_values='foo': 12.3, 'bar':45.6,
                    dates=[datetime.date.today()])

    def orjson_serializer(obj):
        """
            Note that `orjson.dumps()` return byte array, while sqlalchemy expects string, thus `decode()` call.
        """
        return orjson.dumps(obj, option=orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NAIVE_UTC).decode()
    

    engine = create_engine('postgresql://postgres:postgres@localhost:5432/my_schema', json_serializer=orjson_serializer, json_deserializer=orjson.loads)
    with Session(engine, future=True) as session:
        session.add(entry)
        session.commit()

【讨论】：

是的，你是对的！发现它自己深入挖掘文档。我将编辑您的答案以解决orjson.dumps 可能存在的问题