如何将习俗类型数组传递给PostgreSQL函数

debugcn 发表于 Dev

施穆托夫

我在Postgres中有一个自定义数组类型：

CREATE TYPE core.arr_message_input AS (
    idmessage uuid,
    idplugin integer,
    versionplugin numeric,
    ttl integer
);

有一个简单的函数可将记录添加到表中：

CREATE OR REPLACE FUNCTION queue_push(
    arr_message_input[])
    RETURNS Bool
    LANGUAGE 'plpgsql'
AS $$
    BEGIN
        INSERT INTO queue(idmessage, idplugin, versionplugin, queuetime, ttl)
        SELECT idmessage, idplugin, versionplugin, now(), ttl
        FROM unnest ($1);
        RETURN True;
    EXCEPTION
        WHEN others THEN
            RETURN False;
    END $$;

从Postgres填写值很容易：

SELECT queue_push(
    array[
        ('e62c7924-2cd1-4dd6-9b55-d4e612816ce0', 2, 0, 0),
        ('a7e864af-4c4c-452d-9df2-f9d4f70ac02e', 2, 0, 0),
    ]::arr_message_input[]
);

但是通过SQLAlchemy我无法弄清楚该怎么做。我将列表作为数组传递，但是应该有列表或类似的列表。而且我无法用Python做到这一点。

例如：

功能模型描述如下：

class QueuePush(GenericFunction):

    name = "queue_push"

    @staticmethod
    def mapped_objects(**kwargs):
        return select(
            [
                Column('queue_push', BOOLEAN),
            ]
        ).select_from(
            func.queue_push(
                kwargs['arr_message_input'],
            )
        ).alias(name="queue_push")

对功能的要求：

import QueuePush


messages = [
    ['027d6e96-84b7-4f10-8640-13dfa1b05fd8', 3, 0, 2],
]

queue = db.query(QueuePush.mapped_objects(arr_message_input=messages)).all()

但是创建的类型仍然是一种数据结构。我显然做错了。

IljaEverilä

可以使用生成元组语法tuple_()，并且使用postgresql.array()可以生成数组文字。结合这些，您将获得：

from sqlalchemy import tuple_
from sqlalchemy.dialects.postgresql import array

class QueuePush(GenericFunction):

    name = "queue_push"

    @staticmethod
    def mapped_objects(**kwargs):
        input = array([tuple_(*t) for t in kwargs['arr_message_input']]
        return select(
            [
                Column('queue_push', BOOLEAN),
            ]
        ).select_from(
            func.queue_push( 
                # Eurgh, this cast handling is horrid, and uses the postgresql
                # specific cast syntax, but it'd require creating a UserDefinedType
                # or such just to use the proper cast method.
                input.op("::")(text("arr_message_input[]"))
            )
        ).alias(name="queue_push")

对于高级用法，您可以定义UserDefinedType和bind_expression()：

from sqlalchemy import tuple_
from sqlalchemy.types import UserDefinedType

class ArrMessageInput(UserDefinedType):

    def get_col_spec(self, **kw):
        return "arr_message_input"

    def bind_expression(self, val):
        return tuple_(*val)

接着

from sqlalchemy import tuple_
from sqlalchemy.dialects.postgresql import array, ARRAY

class QueuePush(GenericFunction):

    name = "queue_push"

    @staticmethod
    def mapped_objects(**kwargs):
        # For some reason the `bind_expression` is not applied to array elements,
        # even if using `array(..., type_=ArrMessageInput)`
        input = array([tuple_(*t) for t in kwargs['arr_message_input']])
        return select(
            [
                Column('queue_push', BOOLEAN),
            ]
        ).select_from(
            func.queue_push( 
                input.cast(ARRAY(ArrMessageInput))
            )
        ).alias(name="queue_push")

另一方面，如果您使用psycopg2，则可以依靠它对Python列表和元组的适应，只需将结构作为文字传递，转换为正确的类型：

from sqlalchemy import literal
from sqlalchemy.dialects.postgresql import ARRAY

class QueuePush(GenericFunction):

    name = "queue_push"

    @staticmethod
    def mapped_objects(**kwargs):
        # `literal` is required, so that the cast doesn't apply SQLA type handling.
        # Note the use of Python's `tuple` instead of SQLAlchemy `tuple_`
        input = literal([tuple(t) for t in kwargs['arr_message_input']])
        return select(
            [
                Column('queue_push', BOOLEAN),
            ]
        ).select_from(
            func.queue_push( 
                input.cast(ARRAY(ArrMessageInput))
            )
        ).alias(name="queue_push")

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。