Repository: marshmallow-code/marshmallow-sqlalchemy

Issue:
Bugs with enums, (PG) arrays of enums, and (PG) multidimensional arrays
Hello all,

Would like to just say thank you for your work on this project - I've been using it for a few years now and it's been great.

Recently, as was migrating away from `marshmallow_enum` and playing around with some others things, I noticed some issues/bugs in enum definitions se/de.

Here's a quick summary:
1. Creating an array of enums with the `sqlalchemy.dialects.postgresql.ARRAY` instead of the `sqlalchemy.ARRAY` type immediately breaks
2. When using an AutoSchema to infer the se/de schema, Enums are granted Length validators which breaks things when the enum class only extends `enum.Enum`
3. AutoSchema for multidimensional Postgres arrays break

My fix for the first issue was a bit complex, so I'm wondering if there's a better direction to approach it. Here's a reproduction of our initial issue-

```python
import enum
from marshmallow_sqlalchemy import SQLAlchemyAutoSchema
import sqlalchemy as sa
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.dialects.postgresql import ARRAY

# config
db_url = "postgresql+psycopg://***"
engine = create_engine(db_url)
Session = sessionmaker(bind=engine)
session = Session()

Base = declarative_base()


#
class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3


class Data(Base):
    __tablename__ = "data"

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.Enum(MyEnum))
    value_array = sa.Column(sa.ARRAY(sa.Enum(MyEnum)))
    value_array_pg = sa.Column(ARRAY(sa.Enum(MyEnum)))  # pg dialect


class DataSchema(SQLAlchemyAutoSchema):
    class Meta:
        model = Data
        load_instance = True


#


def test_enum():
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)

    new_data = Data(value=MyEnum.one)
    session.add(new_data)
    session.commit()

    results = session.query(Data).all()
    for row in results:
        it = DataSchema().dump(row)
        print(it)


if __name__ == "__main__":
    test_enum()
```

Which gives `TypeError: Enum.__init__() missing 1 required positional argument: 'enum'`

```python
import enum
import functools

from marshmallow_sqlalchemy import ModelConverter, SQLAlchemyAutoSchema
from marshmallow import fields
from marshmallow.base import FieldABC

import sqlalchemy as sa
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.dialects.postgresql import ARRAY as pgARRAY

# config
db_url = "postgresql+psycopg://***"
engine = create_engine(db_url)
Session = sessionmaker(bind=engine)
session = Session()

Base = declarative_base()
#

# problems
# 1. w/ no modification pgARRAY of enums gets `TypeError: Enum.__init__() missing 1 required positional argument: 'enum'`
#   a. add _improved_enum_factory factory to fix this, but then (2)
# 2. pgARRAY of enums breaks w/ error `ValueError: The list elements must be a subclass or instance of marshmallow.base.FieldABC.`
#   a. add `instance = resolved if isinstance(resolved, FieldABC) else resolved()`


# fix/patch
def _improved_array_factory(converter: ModelConverter, data_type):
    resolved = converter._get_field_class_for_data_type(data_type.item_type)
    instance = resolved if isinstance(resolved, FieldABC) else resolved()

    return functools.partial(fields.List, instance)


def _improved_enum_factory(_: ModelConverter, data_type):
    return (
        functools.partial(fields.Enum, enum=data_type.enum_class)
        if data_type.enum_class
        else fields.Field
    )


ModelConverter.SQLA_TYPE_MAPPING[pgARRAY] = _improved_array_factory
# ModelConverter.SQLA_TYPE_MAPPING[sa.ARRAY] = _improved_array_factory
ModelConverter.SQLA_TYPE_MAPPING[sa.Enum] = _improved_enum_factory

# Note that the corresponding se/de schema for an sa.ARRAY gets interpreted as fields.Raw
# so the enums get returned as enum objects (that's why it works natively)


#


#
class MyEnumFix(enum.Enum):
    one = 1
    two = 2
    three = 3


class Data(Base):
    __tablename__ = "data_fix"

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.Enum(MyEnumFix))
    value_array = sa.Column(sa.ARRAY(sa.Enum(MyEnumFix)))
    value_array_pg = sa.Column(pgARRAY(sa.Enum(MyEnumFix)))


class DataSchema(SQLAlchemyAutoSchema):
    class Meta:
        model = Data
        load_instance = True
        sqla_session = session
        include_fk = False


#


def test_enum():
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)

    # test orm
    new_data = Data(
        value=MyEnumFix.one,
        value_array=[MyEnumFix.one, MyEnumFix.two],
        value_array_pg=[MyEnumFix.one, MyEnumFix.two],
    )
    session.add(new_data)
    session.commit()

    results = session.query(Data).all()
    for row in results:
        it = DataSchema().dump(row)
        print(it)


if __name__ == "__main__":
    test_enum()
```

This details the entire fix, and you can see why I needed to make modification to both factories from the comments. I thought it a little weird that `sa.ARRAY(sa.ENUM)` gets se/de as `fields.RAW` which is why I also added an array factory for `sa.ARRAY`, but that doesn't influence the bug/reproduction.

---
Next, I was having trouble with the `TypeError: object of type 'MyEnum' has no len()` error. I noticed in another issue [here](https://github.com/marshmallow-code/marshmallow-sqlalchemy/issues/615) , you can simply subclass enum with its type so below `class MyEnum(str, enum.Enum):` produces no error (this reproduction was mostly taken from the above issue).

```python
import enum
from marshmallow_sqlalchemy import SQLAlchemyAutoSchema
import sqlalchemy as sa
from sqlalchemy.orm import sessionmaker, declarative_base, scoped_session

# config
engine = sa.create_engine("sqlite:///:memory:")
session = scoped_session(sessionmaker(bind=engine))
Base = declarative_base()


#
class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3


class Data(Base):
    __tablename__ = "data"

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(sa.Enum(MyEnum))


class DataSchema(SQLAlchemyAutoSchema):
    class Meta:
        model = Data
        load_instance = True


#


def test_enum():
    data = {"value": "one"}
    DataSchema().load(data, session=session)


if __name__ == "__main__":
    test_enum()

```

Is that the intended way to create enums? We have a bunch of enums and I wouldn't really want to subclass each one. I also created a different fix (before knowing about the subclassing) in which I extended `fields.Enum` and created an `_improved_enum_factory` that applies it (vs the regular `fields.Enum`). Really I just removed the Length validation

```
class EnumFieldForFactory(fields.Enum):

    def __init__(self, enum, *args, **kwargs):
        kwargs["validate"] = []

        # ignoring Length validation for now
        if not kwargs.get("dump_only"):
            if kwargs.get("by_value", False):
                choices = list(enum._member_map_.values())
            else:
                choices = enum._member_names_
            kwargs["validate"].append(validate.OneOf(choices=choices))

        super().__init__(enum, *args, **kwargs)
```

I feel like the behavior is a bit buggy because `_add_column_kwargs` adds a Length validator by default, but would like to know what other people think.

And quickly - I noticed in a recent version (the PR to the issue I referenced above) that most Enum schemas generated from AutoSchema no longer has the OneOf validation, which I thought was odd. I didn't completely dive into the issue / find a fix yet, but that feels weird to me.

---
Lastly, in a slight extension to 1, I had another issue with multidimensional PG arrays and AutoSchema. 

```python
import enum
import functools

from marshmallow_sqlalchemy import ModelConverter, SQLAlchemyAutoSchema
from marshmallow import fields
from marshmallow.base import FieldABC

import sqlalchemy as sa
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.dialects.postgresql import ARRAY as pgARRAY

# config
db_url = "postgresql+psycopg://***"
engine = create_engine(db_url)
Session = sessionmaker(bind=engine)
session = Session()

Base = declarative_base()
#

class Data(Base):
    __tablename__ = "data_multi"

    id = sa.Column(sa.Integer, primary_key=True)
    value = sa.Column(pgARRAY(sa.Float, dimensions=2))
    value2 = sa.Column(sa.ARRAY(sa.Float, dimensions=2))


class DataSchema(SQLAlchemyAutoSchema):
    class Meta:
        model = Data
        load_instance = True


#


def test_enum():
    Base.metadata.drop_all(engine)
    Base.metadata.create_all(engine)

    # test orm
    new_data = Data(
        value=[[1.1, 2.2], [3.3, 4.4]],
        value2=[[1.1, 2.2], [3.3, 4.4]],
    )
    session.add(new_data)
    session.commit()

    results = session.query(Data).all()
    for row in results:
        it = DataSchema().dump(row)
        print(it)


if __name__ == "__main__":
    test_enum()
```

Multidimensional `sa.ARRAY`s work fine because it's deserialized as `fields.RAW`, but not for Postgres, which gets a `fields.List` with an `inner=fields.Float`, not capturing the multidimensionality of the array. I fixed this by modifying the `_improved_array_factory` from above to be the following

```python
def _improved_array_factory(converter: ModelConverter, data_type):
    resolved = converter._get_field_class_for_data_type(data_type.item_type)
    instance = resolved if isinstance(resolved, FieldABC) else resolved()

    dimensions = data_type.dimensions
    for _ in range(dimensions - 1 if dimensions else 0):
        instance = fields.List(instance)

    return functools.partial(fields.List, instance)
```

So while these patches worked for me, I'm not sure if they're right for the codebase, though I do see them potentially fitting in. Would love to discuss this more

Discussion:
@sloria: thanks for the detailed report! will investigate when i get some time@carterjc: No problem! Would be happy to work on a PR if you think that's the right move@sloria: @carterjc feel free to open a PR, even if it's a WIP or just a test reproducing the issue. it'd be greatly appreciated

Solution patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1aee458..5a0234e 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,6 +6,8 @@ Changelog
 
 Bug fixes:
 
+* Fix handling of arrays of enums and multidimensional arrays (:issue:`653`).
+  Thanks :user:`carterjc` for reporting and investigating the fix.
 * Fix handling of `sqlalchemy.PickleType` columns (:issue:`394`)
   Thanks :user:`Eyon42` for reporting.
 
diff --git a/pyproject.toml b/pyproject.toml
index c7cd35d..9a4948a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,11 @@ classifiers = [
   "Programming Language :: Python :: 3.13",
 ]
 requires-python = ">=3.9"
-dependencies = ["marshmallow>=3.18.0", "SQLAlchemy>=1.4.40,<3.0"]
+dependencies = [
+  "marshmallow>=3.18.0",
+  "SQLAlchemy>=1.4.40,<3.0",
+  "typing-extensions; python_version < '3.10'",
+]
 
 [project.urls]
 Changelog = "https://marshmallow-sqlalchemy.readthedocs.io/en/latest/changelog.html"
diff --git a/src/marshmallow_sqlalchemy/convert.py b/src/marshmallow_sqlalchemy/convert.py
index 8c83ae8..e73bce1 100644
--- a/src/marshmallow_sqlalchemy/convert.py
+++ b/src/marshmallow_sqlalchemy/convert.py
@@ -4,13 +4,28 @@ import functools
 import inspect
 import uuid
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any, Callable, Literal, cast, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    Union,
+    cast,
+    overload,
+)
+
+# Remove when dropping Python 3.9
+try:
+    from typing import TypeAlias, TypeGuard
+except ImportError:
+    from typing_extensions import TypeAlias, TypeGuard
 
 import marshmallow as ma
 import sqlalchemy as sa
 from marshmallow import fields, validate
 from sqlalchemy.dialects import mssql, mysql, postgresql
 from sqlalchemy.orm import SynonymProperty
+from sqlalchemy.types import TypeEngine
 
 from .exceptions import ModelConversionError
 from .fields import Related, RelatedList
@@ -18,14 +33,17 @@ from .fields import Related, RelatedList
 if TYPE_CHECKING:
     from sqlalchemy.ext.declarative import DeclarativeMeta
     from sqlalchemy.orm import MapperProperty
-    from sqlalchemy.types import TypeEngine
 
     PropertyOrColumn = MapperProperty | sa.Column
 
-_FieldClassFactory = Callable[[Any, Any], type[fields.Field]]
+_FieldPartial: TypeAlias = Callable[[], fields.Field]
+# TODO: Use more specific type for second argument
+_FieldClassFactory: TypeAlias = Callable[
+    ["ModelConverter", Any], Union[type[fields.Field], _FieldPartial]
+]
 
 
-def _is_field(value) -> bool:
+def _is_field(value: Any) -> TypeGuard[type[fields.Field]]:
     return isinstance(value, type) and issubclass(value, fields.Field)
 
 
@@ -49,17 +67,30 @@ def _is_auto_increment(column) -> bool:
     return column.table is not None and column is column.table._autoincrement_column
 
 
-def _postgres_array_factory(converter: ModelConverter, data_type: postgresql.ARRAY):
-    return functools.partial(
-        fields.List,
-        converter._get_field_class_for_data_type(data_type.item_type),
-    )
+def _list_field_factory(
+    converter: ModelConverter, data_type: postgresql.ARRAY
+) -> Callable[[], fields.List]:
+    FieldClass = converter._get_field_class_for_data_type(data_type.item_type)
+    inner = FieldClass()
+    if not data_type.dimensions or data_type.dimensions == 1:
+        return functools.partial(fields.List, inner)
+
+    # For multi-dimensional arrays, nest the Lists
+    dimensions = data_type.dimensions
+    for _ in range(dimensions - 1):
+        inner = fields.List(inner)
+
+    return functools.partial(fields.List, inner)
 
 
 def _enum_field_factory(
     converter: ModelConverter, data_type: sa.Enum
-) -> type[fields.Field]:
-    return fields.Enum if data_type.enum_class else fields.Raw
+) -> Callable[[], fields.Field]:
+    return (
+        functools.partial(fields.Enum, enum=data_type.enum_class)
+        if data_type.enum_class
+        else fields.Raw
+    )
 
 
 class ModelConverter:
@@ -72,6 +103,7 @@ class ModelConverter:
     ] = {
         sa.Enum: _enum_field_factory,
         sa.JSON: fields.Raw,
+        sa.ARRAY: _list_field_factory,
         sa.PickleType: fields.Raw,
         postgresql.BIT: fields.Integer,
         postgresql.OID: fields.Integer,
@@ -82,7 +114,7 @@ class ModelConverter:
         postgresql.JSON: fields.Raw,
         postgresql.JSONB: fields.Raw,
         postgresql.HSTORE: fields.Raw,
-        postgresql.ARRAY: _postgres_array_factory,
+        postgresql.ARRAY: _list_field_factory,
         postgresql.MONEY: fields.Decimal,
         postgresql.DATE: fields.Date,
         postgresql.TIME: fields.Time,
@@ -335,14 +367,18 @@ class ModelConverter:
     def _get_field_class_for_data_type(
         self, data_type: TypeEngine
     ) -> type[fields.Field]:
-        field_cls = None
+        field_cls: type[fields.Field] | _FieldPartial | None = None
         types = inspect.getmro(type(data_type))
         # First search for a field class from self.SQLA_TYPE_MAPPING
         for col_type in types:
             if col_type in self.SQLA_TYPE_MAPPING:
-                field_cls = self.SQLA_TYPE_MAPPING[col_type]
-                if callable(field_cls) and not _is_field(field_cls):
-                    field_cls = cast(_FieldClassFactory, field_cls)(self, data_type)
+                field_or_factory = self.SQLA_TYPE_MAPPING[col_type]
+                if _is_field(field_or_factory):
+                    field_cls = field_or_factory
+                else:
+                    field_cls = cast(_FieldClassFactory, field_or_factory)(
+                        self, data_type
+                    )
                 break
         else:
             # Try to find a field class based on the column's python_type


Test patch:
diff --git a/tests/test_conversion.py b/tests/test_conversion.py
index 1180e90..01ec7ef 100644
--- a/tests/test_conversion.py
+++ b/tests/test_conversion.py
@@ -226,6 +226,46 @@ class TestPropertyFieldConversion:
         inner_field = getattr(field, "inner", getattr(field, "container", None))
         assert type(inner_field) is fields.Int
 
+    @pytest.mark.parametrize(
+        "array_property",
+        (
+            pytest.param(make_property(sa.ARRAY(sa.Enum(CourseLevel))), id="sa.ARRAY"),
+            pytest.param(
+                make_property(postgresql.ARRAY(sa.Enum(CourseLevel))),
+                id="postgresql.ARRAY",
+            ),
+        ),
+    )
+    def test_convert_ARRAY_Enum(self, converter, array_property):
+        field = converter.property2field(array_property)
+        assert type(field) is fields.List
+        inner_field = field.inner
+        assert type(inner_field) is fields.Enum
+
+    @pytest.mark.parametrize(
+        "array_property",
+        (
+            pytest.param(
+                make_property(sa.ARRAY(sa.Float, dimensions=2)), id="sa.ARRAY"
+            ),
+            pytest.param(
+                make_property(postgresql.ARRAY(sa.Float, dimensions=2)),
+                id="postgresql.ARRAY",
+            ),
+        ),
+    )
+    def test_convert_multidimensional_ARRAY(self, converter, array_property):
+        field = converter.property2field(array_property)
+        assert type(field) is fields.List
+        assert type(field.inner) is fields.List
+        assert type(field.inner.inner) is fields.Float
+
+    def test_convert_one_dimensional_ARRAY(self, converter):
+        prop = make_property(postgresql.ARRAY(sa.Float, dimensions=1))
+        field = converter.property2field(prop)
+        assert type(field) is fields.List
+        assert type(field.inner) is fields.Float
+
     def test_convert_TSVECTOR(self, converter):
         prop = make_property(postgresql.TSVECTOR)
         with pytest.raises(ModelConversionError):