Доступ к составному типу данных с использованием sqlalchemy в postgres

Question

Я пытаюсь извлечь составные столбцы из функции tiger.geocode в python, используя sqlalchemy.
В чистом виде sql это выглядит так:

SELECT   
    g.rating  
    ,ST_X(g.geomout) As lon  
    ,ST_Y(g.geomout) As lat  
    ,(addy).address As stno  
    ,(addy).streetname As street  
    ,(addy).streettypeabbrev As styp  
    ,(addy).location As city  
    ,(addy).stateabbrev As st  
    ,(addy).zip  
FROM geocode(pagc_normalize_address('1 Capitol Square Columbus OH 43215')) As g  
;

Это дает следующий вывод:

#   rating  lon lat stno    street  styp    city    st  zip
1   17  -82.99782603089086  39.96172588526335   1   Capital St  Columbus    OH  43215

Проблема, с которой я сталкиваюсь, заключается в том, как ссылаться на составные столбцы при запросе объекта из sqlalchemy (рейтинг, lon, lat, stno, улица, styp, город, st, zip)?

Пожалуйста и спасибо.

Answer 1

SQLAlchemy не поддерживает функции, напрямую возвращающие множество, но его FunctionElementс считаются FromClauses, что означает, что вы уже можете рассматривать их как таблицы; нам просто нужно добавить возможность выбора определенного столбца из функции. К счастью, это просто (хотя и не очевидно):

from sqlalchemy.sql.base import ColumnCollection
from sqlalchemy.sql.expression import column
from sqlalchemy.sql.functions import FunctionElement

NormAddy = CompositeType(
    "norm_addy",
    [
        Column("address", Integer),
        Column("predirAbbrev", String),
        Column("streetName", String),
        Column("streetTypeAbbrev", String),
        Column("postdirAbbrev", String),
        Column("internal", String),
        Column("location", String),
        Column("stateAbbrev", String),
        Column("zip", String),
        Column("parsed", Boolean),
    ],
)

class geocode(GenericFunction):
    columns = ColumnCollection(
        Column("rating", Integer),
        column("geomout"),  # lowercase column because we don't have the `geometry` type
        Column("addy", NormAddy),
    )

Подклассы от GenericFunction имеет дополнительное преимущество регистрации geocode функционировать глобально, так что func.geocode будет работать как положено.

g = func.geocode(func.pagc_normalize_address("1 Capitol Square Columbus OH 43215")).alias("g")
query = session.query(
    g.c.rating,
    func.ST_X(g.c.geomout).label("lon"),
    func.ST_Y(g.c.geomout).label("lat"),
    g.c.addy.address.label("stno"),
    g.c.addy.streetName.label("street"),
    g.c.addy.streetTypeAbbrev.label("styp"),
    g.c.addy.location.label("city"),
    g.c.addy.stateAbbrev.label("st"),
    g.c.addy.zip,
).select_from(g)

К сожалению, это не совсем работает. Кажется, есть ошибка, которая делает g.c.addy.address синтаксис не работает на последних версиях SQLAlchemy. Мы можем исправить это очень быстро (хотя это должно быть исправлено в sqlalchemy_utils):

from sqlalchemy_utils.types.pg_composite import CompositeElement
import sqlalchemy_utils

class CompositeType(sqlalchemy_utils.CompositeType):
    class comparator_factory(_CompositeType.comparator_factory):
        def __getattr__(self, key):
            try:
                type_ = self.type.typemap[key]
            except KeyError:
                raise AttributeError(key)
            return CompositeElement(self.expr, key, type_)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.typemap = {c.name: c.type for c in self.columns}

Теперь это работает:

print(query.statement.compile(engine))
# SELECT g.rating, ST_X(g.geomout) AS lon, ST_Y(g.geomout) AS lat, (g.addy).address AS stno, (g.addy).streetName AS street, (g.addy).streetTypeAbbrev AS styp, (g.addy).location AS city, (g.addy).stateAbbrev AS st, (g.addy).zip AS zip_1 
# FROM geocode(pagc_normalize_address(%(pagc_normalize_address_1)s)) AS g