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
(cherry picked from commit 780213bfefac7046ac889ffbc1c51e0d244dd678)

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 5fdb065cd6f61d3f84ba829f2282261f1143dfc7..d064ebc1d76f4009a6b0946f0f3158206bd4842c 100644 (file)
--- a/lib/sqlalchemy/dialects/postgresql/ext.py
+++ b/lib/sqlalchemy/dialects/postgresql/ext.py
@@ -52,6 +52,7 @@ class aggregate_order_by(expression.ColumnElement):
 
     def __init__(self, target, *order_by):
         self.target = elements._literal_as_binds(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 4aa85cc926a9def1bf468c95e62b6a42f17f02f9..d5921ece1d967ffcf84b636ea4a65a91759a699c 100644 (file)
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -1317,12 +1317,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):