如何在SQL Expression语句中调整自定义类型？

Posted 2021-04-25

我使用SQLAlchemy ORM映射器在应用程序中有一个自定义类型。对于一些复杂的查询，我需要使用SQL表达式模块，但这使得自定义类型的处理不透明。如何在不使用ORM时告诉SQLAlchemy使用我的自定义类型进行映射？

以下是演示此问题的快速示例。

请注意，第一个查询有效，但我必须首先手动将其转换为Python中的str，然后将INET放在PostgreSQL旁边，即使我已定义自定义类型。

我知道SQL表达式模块不知道自定义类型，因为它在ORM中定义在它上面的一层。但是我想知道是否我无法以某种方式将自定义类型连接到SQL层中，从而使类型和值的使用更加透明。此外，无论使用哪个SA层，都要确保自定义类型中定义的任何操作（清理等）始终如一。

from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql.expression import any_
from sqlalchemy.types import TypeDecorator

Base = declarative_base()


class PgIpInterface(TypeDecorator):
    """
    A codec for :py:mod:`ipaddress` interfaces.
    """

    impl = INET

    def process_bind_param(self, value, dialect):
        return str(value) if value else None

    def process_result_value(self, value, dialect):
        return ip_interface(value) if value else None

    def process_literal_param(self, value, dialect):
        raise NotImplementedError('Not yet implemented')


class Network(Base):
    __tablename__ = 'example_table'
    cidr = Column(PgIpInterface, primary_key=True)


def execute(query):
    import logging
    LOG = logging.getLogger()
    try:
        print(query)
        print(query.all())
    except:
        LOG.exception('!!! failed')


engine = create_engine('postgresql://malbert@/malbert')
Base.metadata.create_all(engine)
Session = sessionmaker(bind=engine)

session = Session()

ranges = [
    ip_interface('192.168.1.0/24'),
    ip_interface('192.168.3.0/24'),
]


# Query with manual casting
print(' Manual Casting via "str" '.center(80, '-'))
arr = array([cast(str(_), INET) for _ in ranges])
query1 = session.query(Network).filter(Network.cidr.op("<<=")(any_(arr)))
execute(query1)


print(' Manual Casting '.center(80, '-'))
arr = array([cast(_, INET) for _ in ranges])
query2 = session.query(Network).filter(Network.cidr.op("<<=")(any_(arr)))
execute(query2)


# Query without casting
print(' No Casting '.center(80, '-'))
query3 = session.query(Network).filter(Network.cidr.op("<<=")(any_(ranges)))
execute(query3)

答案

要使第二个查询起作用，只需转换为自定义类型：

arr = array([cast(_, PgIpInterface) for _ in ranges])

要使您的第三个查询起作用，您需要在psycopg2中更深入一级。 psycopg2为builtin support类型有ipaddress，但不幸的是它似乎不完整。 （ipaddress类型转换为字符串而没有显式转换。）

register_ipaddress()  # register ipaddress handling globally
arr = [ip_interface('192.168.1.0/24'), ip_interface('192.168.3.0/24')]
session.query(Foo).filter(Foo.cidr.op("<<=")(any_(arr))).all()

这呈现出类似的东西

WHERE foo.cidr <<= ANY (ARRAY['192.168.1.0/24', '192.168.3.0/24'])

因operator does not exist: inet <<= text错误而失败。幸运的是，它很容易修复;我们将自己重写register_ipaddress：

import ipaddress

from psycopg2.extensions import (
    AsIs,
    new_array_type,
    new_type,
    register_adapter,
    register_type
)


def register_ipaddress():
    def cast_interface(s, cur=None):
        if s is None:
            return None
        return ipaddress.ip_interface(s)
    inet = new_type((869,), 'INET', cast_interface)
    ainet = new_array_type((1041,), 'INET[]', inet)

    def cast_network(s, cur=None):
        if s is None:
            return None
        return ipaddress.ip_network(s)
    cidr = new_type((650,), 'CIDR', cast_network)
    acidr = new_array_type((651,), 'CIDR[]', cidr)

    for caster in [inet, ainet, cidr, acidr]:
        register_type(caster)

    def adapt_interface(obj):
        return AsIs("'{}'::inet".format(obj))
    for t in [ipaddress.IPv4Interface, ipaddress.IPv6Interface]:
        register_adapter(t, adapt_interface)

    def adapt_network(obj):
        return AsIs("'{}'::cidr".format(obj))
    for t in [ipaddress.IPv4Network, ipaddress.IPv6Network]:
        register_adapter(t, adapt_network)

这将呈现您的查询

WHERE foo.cidr <<= ANY (ARRAY['192.168.1.0/24'::inet, '192.168.3.0/24'::inet])

注意使用之间的区别

arr = array([ip_interface...])

和

arr = [ip_interface...]

在前一种情况下，数组由SQLAlchemy处理，因此您将获得列表中n项的n绑定参数;在后一种情况下，数组由psycopg2处理，因此您将获得整个数组的一个绑定参数。