SQLAlchemy 如何在 SQLite 和 PostgreSQL 之间互换使用带有索引的 JSON 列

Question

我正在努力使用 sqlalchemy 定义一个模式，该模式可与多个引擎后端一起使用，特别是 sqlite 和 postgresql。

我遇到问题，因为我有一个带有索引的 JSON 列。这似乎适用于 sqlite，但对于 postgresql，它会抱怨索引类型不能是 btree。我看过突出显示特定于 postgres 方言的 JSONB 类型的文档，但问题是我的架构是声明性的：我不知道是否要连接到 SQLite 或 PostgreSQL 数据库。

作为示例，以下是一个玩具声明性模式：

    # from sqlalchemy.dialects.postgresql import JSONB
    from sqlalchemy import create_engine
    from sqlalchemy import inspect
    from sqlalchemy.ext.declarative import declarative_base
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.sql.schema import Column, Index
    from sqlalchemy.types import Integer, JSON
    from sqlalchemy_utils import database_exists, create_database

    CustomBase = declarative_base()

    class User(CustomBase):
        __tablename__ = 'users'
        id = Column(Integer, primary_key=True, doc='unique internal id')
        name = Column(JSON)
        loose_identifer = Column(JSON, index=True, unique=False)
        # loose_identifer = Column(JSONB, index=True, unique=False)

    uri = 'sqlite:///test_sqlite_v7.sqlite'
    # uri = 'postgresql+psycopg2://admin:admin@localhost:5432/test_postgresql_v4.postgres'

    engine = create_engine(uri)
    DBSession = sessionmaker(bind=engine)
    session = DBSession()

    if 'postgresql' in uri:
        if not database_exists(uri):
            create_database(uri)

    inspector = inspect(engine)
    table_names = inspector.get_table_names()
    if len(table_names) == 0:
        CustomBase.metadata.create_all(engine)

    user_infos = [
        {'name': 'user1', 'loose_identifer': "AA" },
        {'name': 'user2', 'loose_identifer': "33" },
        {'name': 'user3', 'loose_identifer': 33 },
        {'name': 'user4', 'loose_identifer': 33 },
        {'name': 'user5', 'loose_identifer': "AA" },
        {'name': 'user6', 'loose_identifer': None},
        {'name': 'user7', 'loose_identifer': [1, 'weird']},
    ]
    for row in user_infos:
        user = User(**row)
        session.add(user)

    session.commit()

    import pandas as pd
    import json
    table_df = pd.read_sql_table('users', con=engine)
    table_df['loose_identifer'] = table_df['loose_identifer'].apply(repr)
    print(table_df)

    query = session.query(User.name, User.loose_identifer).filter(User.loose_identifer == json.dumps(33))
    results = list(query.all())
    print(f'results={results}')

    query = session.query(User.name, User.loose_identifer).filter(User.loose_identifer == json.dumps('33'))
    results = list(query.all())
    print(f'results={results}')

该User表有一loose_identifer列，我希望允许它是相当任意的 JSON 类型，并且我想在其上添加索引。我这样的主要原因是因为我必须支持这些“松散”标识符，它们可以是整数或字符串。

当我使用 sqlite 时，使用 aColumn(JSON, index=True, unique=False)似乎工作得很好，但是当我将其切换到目标 postgresql 引擎时，我收到此错误：

ProgrammingError: (psycopg2.errors.UndefinedObject) data type json has no default operator class for access method "btree"
HINT:  You must specify an operator class for the index or define a default operator class for the data type.

[SQL: CREATE INDEX ix_users_loose_identifer ON users (loose_identifer)]
(Background on this error at: https://sqlalche.me/e/14/f405)

我尝试通过添加此类属性来显式添加索引：

    __table_args__ =  (
        # /sf/ask/2162009251/
        Index(
            "ix_users_loose_identifer", loose_identifer,
            postgresql_using="gin",
        ),
    )

但这似乎不起作用。我可能在该声明中做错了什么。

如果我在上面的模式中更改JSON为，它确实可以工作，但与 sqlite 不兼容，所以我的问题是：如何使用 json 列声明我的模式，这些列将使用 sqlite 和 postgresql 后端之间兼容的语法进行索引？JSONBJSONB

Answer 1

我遇到了类似的问题。这是该问题的更简洁的演示：

from sqlalchemy import JSON, Column, Integer, Unicode, create_engine
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import declarative_base

Base = declarative_base()


class Node(Base):
    __tablename__ = "nodes"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    key = Column(Unicode(1023), index=True, nullable=False)
    ancestors = Column(JSON, index=True, nullable=True)


# This works:
engine = create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)
print("SQLite works")
Base.metadata.drop_all(engine)

# One way to get a postgres database to test against:
# docker run --name test-postgres -e POSTGRES_PASSWORD=secret -d docker.io/postgres

# This fails:
engine = create_engine("postgresql://postgres:secret@localhost:5432")
Base.metadata.create_all(engine)
print("PostgreSQL works")
Base.metadata.drop_all(engine)

GitHub 上的 SQLAlchemy 讨论中为我提供了解决方案。@Erotemic 和我所缺少的功能是TypeEngine.with_variant。下面，我将其应用到我的简单示例中来演示修复：

from sqlalchemy import JSON, Column, Integer, Unicode, create_engine
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import declarative_base

Base = declarative_base()
# Use JSON with SQLite and JSONB with PostgreSQL.
JSONVariant = JSON().with_variant(JSONB(), "postgresql")


class Node(Base):
    __tablename__ = "nodes"

    id = Column(Integer, primary_key=True, index=True, autoincrement=True)
    key = Column(Unicode(1023), index=True, nullable=False)
    ancestors = Column(JSONVariant, index=True, nullable=True)


# Both of these now work:
engine = create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)
print("SQLite works")
Base.metadata.drop_all(engine)

engine = create_engine("postgresql://postgres:secret@localhost:5432")
Base.metadata.create_all(engine)
print("PostgreSQL works")
Base.metadata.drop_all(engine)

感谢 GitHub 上的 @CaselIT 提供的快速帮助！