cloud-utils / sqlalchemy-aurora-data-api

An AWS Aurora Serverless Data API dialect for SQLAlchemy
https://cloud-utils.github.io/sqlalchemy-aurora-data-api/
Apache License 2.0
113 stars 21 forks source link

Support non-String ARRAYs in Postgresql dialect #21

Open jmullercuber opened 3 years ago

jmullercuber commented 3 years ago

Right now this project only supports string arrays. However, postgresql supports arrays of many other types [1].

A demonstrative example of this behavior is below:

import enum

from sqlalchemy import create_engine, Column, Enum, Integer, String
from sqlalchemy.dialects.postgresql import JSONB, ARRAY
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

Base = declarative_base()

class SockColor(enum.Enum):
    red = 1
    green = 2
    black = 3

class User(Base):
    __tablename__ = "user"

    id = Column(Integer, primary_key=True)
    name = Column(String)
    friends = Column(ARRAY(String))
    lucky_numbers = Column(ARRAY(Integer))  # this is not okay
    docs = Column(ARRAY(JSONB))  # this is not okay
    socks = Column(ARRAY(Enum(SockColor)))  # this is not okay

class TestPostgresDialect:
    engine = create_engine(
        "postgresql+auroradataapi://:@/TEST_DB_NAME",
        connect_args=dict(
            aurora_cluster_arn="TEST_CLUSTER_ARN", secret_arn="TEST_SECRET_ARN"
        ),
    )

    def test_orm(self):
        friends = ["Scarlett O'Hara", 'Ada "Hacker" Lovelace']
        lucky_numbers = [2, 3, 5, 7, 11, 13, 17, 19]
        docs = [
            {"fizz": [1, 2, 3]},
            {"buzz": [4, 5, 6]},
            {"bazz": [7, 8, 9]},
        ]
        socks = [SockColor.red, SockColor.red, SockColor.green]
        Base.metadata.create_all(self.engine)

        ed_user = User(
            name="ed",
            friends=friends,
            lucky_numbers=lucky_numbers,
            docs=docs,
            socks=socks,
        )

        Session = sessionmaker(bind=self.engine)
        session = Session()

        # clear data table
        session.query(User).delete()
        session.commit()

        # add record for test user
        session.add(ed_user)
        session.commit() # <--- test breaks here

        # query user record that we just inserted
        u = session.query(User).filter(User.name.like("%ed")).first()

        self.assertEqual(u.friends, friends)
        self.assertEqual(u.lucky_numbers, lucky_numbers)
        self.assertEqual(u.docs, docs)
        self.assertEqual(u.socks, socks)

Expected Result

Test should pass

Actual Result

value = [2, 3, 5, 7, 11, 13, ...]

    def process(value):
        # FIXME: escape strings properly here
>       return "\v".join(value) if isinstance(value, list) else value
E       sqlalchemy.exc.StatementError: (builtins.TypeError) sequence item 0: expected str instance, int found
E       [SQL: INSERT INTO "user" (name, friends, lucky_numbers, docs, socks) VALUES (:name, string_to_array(:friends, :string_to_array_1), string_to_array(:lucky_numbers, :string_to_array_2), string_to_array(:docs, :string_to_array_3), string_to_array(:socks, :string_to_array_4)) RETURNING "user".id]
E       [parameters: [{'lucky_numbers': [2, 3, 5, 7, 11, 13, 17, 19], 'docs': [{'fizz': [1, 2, 3]}, {'buzz': [4, 5, 6]}, {'bazz': [7, 8, 9]}], 'name': 'ed', 'friends': ["Scarlett O'Hara", 'Ada "Hacker" Lovelace'], 'socks': [<SockColor.red: 1>, <SockColor.red: 1>, <SockColor.green: 2>]}]]

Inserting just one of the un-supported array types produces similar error messages:

(builtins.TypeError) sequence item 0: expected str instance, SockColor found
(builtins.TypeError) sequence item 0: expected str instance, dict found
kislyuk commented 2 years ago

Thank you for bringing up this idea. At this time I lack the bandwidth to implement it. PRs (with documentation and tests) are welcome.