diff options
| -rw-r--r-- | doc/build/changelog/unreleased_13/4756.rst | 11 | ||||
| -rw-r--r-- | lib/sqlalchemy/dialects/postgresql/array.py | 30 | ||||
| -rw-r--r-- | test/dialect/postgresql/test_types.py | 112 |
3 files changed, 150 insertions, 3 deletions
diff --git a/doc/build/changelog/unreleased_13/4756.rst b/doc/build/changelog/unreleased_13/4756.rst new file mode 100644 index 000000000..d27c39235 --- /dev/null +++ b/doc/build/changelog/unreleased_13/4756.rst @@ -0,0 +1,11 @@ +.. change:: + :tags: usecase, postgresql + :tickets: 4756 + + Added support for multidimensional Postgresql array literals via nesting + the :class:`.postgresql.array` object within another one. The + multidimensional array type is detected automatically. + + .. seealso:: + + :class:`.postgresql.array` diff --git a/lib/sqlalchemy/dialects/postgresql/array.py b/lib/sqlalchemy/dialects/postgresql/array.py index 594cd3a0c..81bde2a02 100644 --- a/lib/sqlalchemy/dialects/postgresql/array.py +++ b/lib/sqlalchemy/dialects/postgresql/array.py @@ -60,7 +60,7 @@ class array(expression.Tuple): array([1,2]) + array([3,4,5]) ]) - print stmt.compile(dialect=postgresql.dialect()) + print(stmt.compile(dialect=postgresql.dialect())) Produces the SQL:: @@ -73,6 +73,24 @@ class array(expression.Tuple): array(['foo', 'bar'], type_=CHAR) + Multidimensional arrays are produced by nesting :class:`.array` constructs. + The dimensionality of the final :class:`.ARRAY` type is calculated by + recursively adding the dimensions of the inner :class:`.ARRAY` type:: + + stmt = select([ + array([ + array([1, 2]), array([3, 4]), array([column('q'), column('x')]) + ]) + ]) + print(stmt.compile(dialect=postgresql.dialect())) + + Produces:: + + SELECT ARRAY[ARRAY[%(param_1)s, %(param_2)s], + ARRAY[%(param_3)s, %(param_4)s], ARRAY[q, x]] AS anon_1 + + .. versionadded:: 1.3.6 added support for multidimensional array literals + .. seealso:: :class:`.postgresql.ARRAY` @@ -83,7 +101,15 @@ class array(expression.Tuple): def __init__(self, clauses, **kw): super(array, self).__init__(*clauses, **kw) - self.type = ARRAY(self.type) + if isinstance(self.type, ARRAY): + self.type = ARRAY( + self.type.item_type, + dimensions=self.type.dimensions + 1 + if self.type.dimensions is not None + else 2, + ) + else: + self.type = ARRAY(self.type) def _bind_param(self, operator, obj, _assume_scalar=False, type_=None): if _assume_scalar or operator is operators.getitem: diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py index 72335ebe3..557b91622 100644 --- a/test/dialect/postgresql/test_types.py +++ b/test/dialect/postgresql/test_types.py @@ -950,6 +950,55 @@ class TimePrecisionTest(fixtures.TestBase, AssertsCompiledSQL): class ArrayTest(AssertsCompiledSQL, fixtures.TestBase): __dialect__ = "postgresql" + def test_array_literal(self): + obj = postgresql.array([1, 2]) + postgresql.array([3, 4, 5]) + + self.assert_compile( + obj, + "ARRAY[%(param_1)s, %(param_2)s] || " + "ARRAY[%(param_3)s, %(param_4)s, %(param_5)s]", + params={ + "param_1": 1, + "param_2": 2, + "param_3": 3, + "param_4": 4, + "param_5": 5, + }, + ) + self.assert_compile( + obj[1], + "(ARRAY[%(param_1)s, %(param_2)s] || ARRAY[%(param_3)s, " + "%(param_4)s, %(param_5)s])[%(param_6)s]", + params={ + "param_1": 1, + "param_2": 2, + "param_3": 3, + "param_4": 4, + "param_5": 5, + }, + ) + + def test_array_literal_getitem_multidim(self): + obj = postgresql.array( + [postgresql.array([1, 2]), postgresql.array([3, 4])] + ) + + self.assert_compile( + obj, + "ARRAY[ARRAY[%(param_1)s, %(param_2)s], " + "ARRAY[%(param_3)s, %(param_4)s]]", + ) + self.assert_compile( + obj[1], + "(ARRAY[ARRAY[%(param_1)s, %(param_2)s], " + "ARRAY[%(param_3)s, %(param_4)s]])[%(param_5)s]", + ) + self.assert_compile( + obj[1][0], + "(ARRAY[ARRAY[%(param_1)s, %(param_2)s], " + "ARRAY[%(param_3)s, %(param_4)s]])[%(param_5)s][%(param_6)s]", + ) + def test_array_type_render_str(self): self.assert_compile(postgresql.ARRAY(Unicode(30)), "VARCHAR(30)[]") @@ -1350,7 +1399,7 @@ class ArrayRoundTripTest(object): [[util.ue("m\xe4\xe4")], [util.ue("m\xf6\xf6")]], ) - def test_array_literal(self): + def test_array_literal_roundtrip(self): eq_( testing.db.scalar( select( @@ -1360,6 +1409,67 @@ class ArrayRoundTripTest(object): [1, 2, 3, 4, 5], ) + eq_( + testing.db.scalar( + select( + [ + ( + postgresql.array([1, 2]) + + postgresql.array([3, 4, 5]) + )[3] + ] + ) + ), + 3, + ) + + eq_( + testing.db.scalar( + select( + [ + ( + postgresql.array([1, 2]) + + postgresql.array([3, 4, 5]) + )[2:4] + ] + ) + ), + [2, 3, 4], + ) + + def test_array_literal_multidimensional_roundtrip(self): + eq_( + testing.db.scalar( + select( + [ + postgresql.array( + [ + postgresql.array([1, 2]), + postgresql.array([3, 4]), + ] + ) + ] + ) + ), + [[1, 2], [3, 4]], + ) + + eq_( + testing.db.scalar( + select( + [ + postgresql.array( + [ + postgresql.array([1, 2]), + postgresql.array([3, 4]), + ] + )[2][1] + ] + ) + ), + 3, + ) + def test_array_literal_compare(self): eq_( testing.db.scalar(select([postgresql.array([1, 2]) < [3, 4, 5]])), |