Fix array_agg to accommodate ARRAY arguments

Fixed bug in :func:`.array_agg` function where passing an argument
that is already of type :class:`.ARRAY`, such as a Postgresql
:obj:`.postgresql.array` construct, would produce a ``ValueError``, due
to the function attempting to nest the arrays.

Change-Id: Ibe5f6275d90e4868e6ef8a733de05acd44c05d78
Fixes: #4107
(cherry picked from commit 4bb8397ae3a9d65bd18eb1d7c951bf5121ea280a)

diff --git a/doc/build/changelog/unreleased_11/4107.rst b/doc/build/changelog/unreleased_11/4107.rst
new file mode 100644 (file)
index 0000000..013a8e8
--- /dev/null
+++ b/doc/build/changelog/unreleased_11/4107.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: bug, postgresql
+    :tickets: 4107
+    :versions: 1.2.0b3
+
+    Fixed bug in :obj:`.array_agg` function where passing an argument
+    that is already of type :class:`.ARRAY`, such as a Postgresql
+    :obj:`.postgresql.array` construct, would produce a ``ValueError``, due
+    to the function attempting to nest the arrays.
\ No newline at end of file

diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
index 08f1d32a504798be7dd9506df998107de39b48b8..42c5df464109dd6b84166268047b55c665475d9b 100644 (file)
--- a/lib/sqlalchemy/sql/functions.py
+++ b/lib/sqlalchemy/sql/functions.py
@@ -689,7 +689,14 @@ class array_agg(GenericFunction):
 
     def __init__(self, *args, **kwargs):
         args = [_literal_as_binds(c) for c in args]
-        kwargs.setdefault('type_', self.type(_type_from_args(args)))
+
+        if 'type_' not in kwargs:
+
+            type_from_args = _type_from_args(args)
+            if isinstance(type_from_args, sqltypes.ARRAY):
+                kwargs['type_'] = type_from_args
+            else:
+                kwargs['type_'] = sqltypes.ARRAY(type_from_args)
         kwargs['_parsed_args'] = args
         super(array_agg, self).__init__(*args, **kwargs)
 

diff --git a/test/sql/test_functions.py b/test/sql/test_functions.py
index 51ce6b2386bef1bbfc8df23adc1374cbf7decfe8..9fa3089856fbbbf8818ff4bae752f3d88775f3ad 100644 (file)
--- a/test/sql/test_functions.py
+++ b/test/sql/test_functions.py
@@ -554,13 +554,49 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
         )
 
 
-class ReturnTypeTest(fixtures.TestBase):
+class ReturnTypeTest(AssertsCompiledSQL, fixtures.TestBase):
 
     def test_array_agg(self):
         expr = func.array_agg(column('data', Integer))
         is_(expr.type._type_affinity, ARRAY)
         is_(expr.type.item_type._type_affinity, Integer)
 
+    def test_array_agg_array_datatype(self):
+        expr = func.array_agg(column('data', ARRAY(Integer)))
+        is_(expr.type._type_affinity, ARRAY)
+        is_(expr.type.item_type._type_affinity, Integer)
+
+    def test_array_agg_array_literal_implicit_type(self):
+        from sqlalchemy.dialects.postgresql import array, ARRAY as PG_ARRAY
+        expr = array([column('data', Integer), column('d2', Integer)])
+
+        assert isinstance(expr.type, PG_ARRAY)
+
+        agg_expr = func.array_agg(expr)
+        assert isinstance(agg_expr.type, PG_ARRAY)
+        is_(agg_expr.type._type_affinity, ARRAY)
+        is_(agg_expr.type.item_type._type_affinity, Integer)
+
+        self.assert_compile(
+            agg_expr,
+            "array_agg(ARRAY[data, d2])",
+            dialect="postgresql"
+        )
+
+    def test_array_agg_array_literal_explicit_type(self):
+        from sqlalchemy.dialects.postgresql import array
+        expr = array([column('data', Integer), column('d2', Integer)])
+
+        agg_expr = func.array_agg(expr, type_=ARRAY(Integer))
+        is_(agg_expr.type._type_affinity, ARRAY)
+        is_(agg_expr.type.item_type._type_affinity, Integer)
+
+        self.assert_compile(
+            agg_expr,
+            "array_agg(ARRAY[data, d2])",
+            dialect="postgresql"
+        )
+
     def test_mode(self):
         expr = func.mode(0.5).within_group(
             column('data', Integer).desc())