Create CompositeArray of CompositeType without using sqlalchemy_utils

Question:

Within my FastAPI, SQLAlchemy, Alembic project, I was using postgresql engine, and now I’m trying to migrate to async with postgresql+asyncpg.
The issue here is that one of my DB schemas has this structure:

class MyTable(...):
    __tablename__ = 'my_table'

    name = Column(String(255), nullable=False, unique=True)
    tridimensional = Column(CompositeArray(
            CompositeType(
                'tridimensional_type', [
                    Column('x', Numeric(4, 0), nullable=False, default=0),
                    Column('y', Numeric(4, 0), nullable=False),
                    Column('z', Numeric(4, 0), nullable=False),
                ]
            ),
        ),
    )

Since this was relying entirely on sqlalchemy_utils.types.pg_composite (Both CompositeArray and CompositeType) and this does not have support to register_composites for postgresql+asyncpg, I was wondering (if possible) how to:

  • Create my own sqlalchemy.types.UserDefinedType that involves this tridimensional type
  • Add it to an alembic migration
  • Create a column using it within an array
Asked By: lvl

||
Source

Answers:

So, I managed to figure it out kinda using CompositeType:
I created this two classes:

  1. To access elements from a dict using the dot operator .
class Dotdict(dict):
    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

And the inherited class to process it from

class AsyncCompositeType(CompositeType):
    # async sessions return asyncpg.Record not the type define in sqlalchemy_utils
    # these wrappers returns a dict when the composite type is loaded using an async session
    # see https://docs.sqlalchemy.org/en/14/core/custom_types.html#typedecorator-recipes
    def result_processor(self, dialect, coltype):
        def process(record):
            if isinstance(record, asyncpg.Record):  # handle async sqlalchemy session
                return Dotdict(record)
            return record
        return process

Finally, your table would only be changed to:

class MyTable(...):
    __tablename__ = 'my_table'

    name = Column(String(255), nullable=False, unique=True)
    tridimensional = Column(CompositeArray(
            AsyncCompositeType(
                'tridimensional_type', [
                    Column('x', Numeric(4, 0), nullable=False, default=0),
                    Column('y', Numeric(4, 0), nullable=False),
                    Column('z', Numeric(4, 0), nullable=False),
                ]
            ),
        ),
    )

Hopefully this helps someone in the future.

Answered By: lvl

I’ve stumbled upon the same issue and came up with this (probably incomplete and not very thoroughly tested) implementation. It’s also based on ideas from sqlalchemy-utils but avoids hard psycopg2 dependency.

Mind that there is no CompositeArray equivalent. It seems that no custom type is actually required for such arrays. The same functionality can be achieved by explicitly declaring dimensions and item type of your ARRAY of composites and SQLAlchemy will properly pass each array entry to result_processor of item’s UserDefinedType (and hence to process_result_value of TypeDecorator wrapper).

All in all, the following code demonstrates intended usage (tested on Python 3.11 and SQLAlchemy 1.4):

import asyncio
import dataclasses
import pprint
from contextlib import AsyncExitStack
from decimal import Decimal

import sqlalchemy as sa
import sqlalchemy.ext.asyncio
from sqlalchemy.dialects import postgresql as sa_pg

import composite

if __name__ == "__main__":
    async def main():
        metadata = sa.MetaData()

        demo_composite_type = composite.define_composite(
            "demo_composite",
            metadata,
            ("field1", sa.TEXT),
            ("field2", sa.Numeric),
        )

        # NOTE: The following class might be omitted if `asyncpg.Record` as a return type is sufficient for your
        #       needs and there is no need to convert items to some custom (data)class.
        class demo_composite_type(sa.types.TypeDecorator):
            impl, cache_ok = demo_composite_type, True

            python_type = dataclasses.make_dataclass(
                "DemoComposite",
                [name for name in demo_composite_type.fields],
            )

            def process_result_value(self, value, dialect):
                if value is not None:
                    return self.python_type(**value)

        async with AsyncExitStack() as deffer:
            pg_engine = sa.ext.asyncio.create_async_engine(
                "postgresql+asyncpg://scott:tiger@localhost/test"
            )
            deffer.push_async_callback(pg_engine.dispose)

            pg_conn = await deffer.enter_async_context(
                pg_engine.connect()
            )

            await pg_conn.run_sync(metadata.create_all)
            deffer.push_async_callback(pg_conn.run_sync, metadata.drop_all)

            values_stmt = (
                sa.sql.Values(
                    sa.column("column1", sa.TEXT),
                    sa.column("column2", sa.Numeric),
                )
                .data([
                    ("1", Decimal(1)),
                    ("2", Decimal(2)),
                ])
                .alias("samples")
            )

            result = (await pg_conn.execute(
                sa.select([
                    sa.func.array_agg(
                        sa.cast(
                            sa.tuple_(
                                values_stmt.c.column1,
                                values_stmt.c.column2,
                            ),
                            demo_composite_type,
                        ),
                        type_=sa_pg.ARRAY(
                            demo_composite_type,
                            dimensions=1,
                        ),
                    )
                    .label("array_of_composites")
                ])
                .select_from(
                    values_stmt
                )
            )).scalar_one()

            pprint.pprint(result)


    asyncio.run(main())

Output of the sample above is the following:

[DemoComposite(field1='1', field2=Decimal('1')),
 DemoComposite(field1='2', field2=Decimal('2'))]
Answered By: zeronineseven