Fix aggregate_order_by issue

Fixes: #5989
Using array_agg() with aggregate_order_by() in
postgresql would sometimes return ARRAY(NullType)
instead of ARRAY(actual_type).

Change-Id: I05a0b2b7ea59291e3c04575578adcc337296e5a8

diff --git a/doc/build/changelog/unreleased_13/5989.rst b/doc/build/changelog/unreleased_13/5989.rst
new file mode 100644 (file)
index 0000000..cccf227
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/5989.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: bug, orm, postgresql
+    :tickets: 5989
+
+    Fixed issue where using :class:`_ext.aggregate_order_by` would return
+    ARRAY(NullType) under certain conditions.

diff --git a/lib/sqlalchemy/dialects/postgresql/ext.py b/lib/sqlalchemy/dialects/postgresql/ext.py
index 71a0aa5a6894f0e3b89dab0999de1f8671005c88..959e6597b5f402b4758b5904cc3a571e9c989228 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/ext.py
+++ b/lib/sqlalchemy/dialects/postgresql/ext.py
@@ -57,6 +57,7 @@ class aggregate_order_by(expression.ColumnElement):
 
     def __init__(self, target, *order_by):
         self.target = coercions.expect(roles.ExpressionElementRole, target)
+        self.type = self.target.type
 
         _lob = len(order_by)
         if _lob == 0:

diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py
index b36f50408d5f9df11bf3d9a18cf7af06ce1fa580..b24132f699b9b06e8d949852aeac2f20ed63da79 100644 (file)
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -1280,12 +1280,31 @@ class ArrayTest(AssertsCompiledSQL, fixtures.TestBase):
         is_(expr.type.__class__, types.ARRAY)
         is_(expr.type.item_type.__class__, Integer)
 
-    def test_array_agg_specific(self):
+    @testing.combinations(
+        ("original", False, False),
+        ("just_enum", True, False),
+        ("just_order_by", False, True),
+        ("issue_5989", True, True),
+        id_="iaa",
+        argnames="with_enum, using_aggregate_order_by",
+    )
+    def test_array_agg_specific(self, with_enum, using_aggregate_order_by):
+        from sqlalchemy.dialects.postgresql import aggregate_order_by
         from sqlalchemy.dialects.postgresql import array_agg
+        from sqlalchemy.dialects.postgresql import ENUM
 
-        expr = array_agg(column("q", Integer))
+        element_type = ENUM if with_enum else Integer
+        expr = (
+            array_agg(
+                aggregate_order_by(
+                    column("q", element_type), column("idx", Integer)
+                )
+            )
+            if using_aggregate_order_by
+            else array_agg(column("q", element_type))
+        )
         is_(expr.type.__class__, postgresql.ARRAY)
-        is_(expr.type.item_type.__class__, Integer)
+        is_(expr.type.item_type.__class__, element_type)
 
 
 class ArrayRoundTripTest(object):