Access composite data type using sqlalchemy in postgres

Question

I'm trying to pull out composite columns from tiger.geocode function in python using sqlalchemy.
In pure sql form this looks like this:

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 
;

This produces the following output:

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

The issue I'm faced with is how to reference composite columns when querying the object from sqlalchemy (rating, lon, lat, stno, street, styp, city, st, zip)?

Please and Thank you.

• univerio: Have you looked at sqlalchemy_utils.types.pg_composite?
• BrianEdwardHoover: I can see how that would work in the case when my composite type is a column on a static table, but in my case, it is a composite type returned by a function. In my setup, we haven't defined sqlalchemy classes for any functions (only tables). Any thoughts on this portion of the problem?

Answer 1

SQLAlchemy doesn't support set returning functions directly, but its FunctionElements are considered FromClauses, which means you can already treat them as tables; we just need to add the ability to select a particular column from the function. Fortunately, this is straightforward (although not obvious):

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),
 )

Subclassing from GenericFunction has the added benefit of registering the geocode function globally so that func.geocode will work as expected.

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)

Unfortunately this doesn't quite work. There seems to be a bug that makes g.c.addy.address syntax not work on recent versions of SQLAlchemy. We can fix it real quick (although this should really be fixed in 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}

Now it works:

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