Preserve ARRAY update expression arguments and binary values

Author: laughingman7743

pyathena/formatter.py

@@ -22,7 +22,7 @@
 class _ComplexParameter:
     """Typed complex value supplied by the SQLAlchemy dialect."""
 
-    constructor: Literal["ARRAY", "MAP", "ROW", "JSON_PARSE"]
+    constructor: Literal["ARRAY", "MAP", "ROW", "JSON_PARSE", "FROM_HEX"]
     values: tuple[Any, ...]
 
 

pyathena/sqlalchemy/_array_update.py

@@ -3,11 +3,12 @@
 from typing import Any
 
 from sqlalchemy import exc, types, util
-from sqlalchemy.sql import operators, visitors
+from sqlalchemy.sql import operators
 from sqlalchemy.sql.elements import BinaryExpression, BindParameter, ColumnElement, Null, Slice
 from sqlalchemy.sql.schema import Column
 from sqlalchemy.sql.visitors import InternalTraversal
 
+from pyathena.formatter import _ComplexParameter
 from pyathena.sqlalchemy.types import _array_item_type, _bind_complex, _literal_complex
 
 
@@ -20,7 +21,13 @@ def __init__(self, item_type):
         self.item_type = item_type
 
     def bind_processor(self, dialect):
-        return lambda value: _bind_complex(value, self.item_type, dialect)
+        def process(value):
+            value = _bind_complex(value, self.item_type, dialect)
+            if isinstance(value, (bytes, bytearray)):
+                return _ComplexParameter("FROM_HEX", (value.hex(),))
+            return value
+
+        return process
 
     def literal_processor(self, dialect):
         return lambda value: _literal_complex(value, self.item_type, dialect)
@@ -118,12 +125,8 @@ def _index_sql(compiler, index: ColumnElement[Any], **kw):
     if not isinstance(index.type, (types.Integer, types.NullType, types.ARRAY)):
         raise exc.CompileError("ARRAY write indices must be integers")
 
-    def type_bind(element: Any, **kwargs: Any) -> Any:
-        if isinstance(element, BindParameter):
-            return element._with_binary_element_type(_IndexType())
-        return None
-
-    index = visitors.replacement_traverse(index, {}, type_bind)
+    if isinstance(index, BindParameter):
+        index = index._with_binary_element_type(_IndexType())
     sql = compiler.process(index, **kw)
     failure = (
         f"CAST(concat('Invalid ARRAY index: ', coalesce(CAST({sql} AS VARCHAR), 'NULL')) AS BIGINT)"

tests/pyathena/sqlalchemy/test_array_update.py

@@ -1,5 +1,5 @@
 import pytest
-from sqlalchemy import Column, Integer, MetaData, Table, bindparam, types
+from sqlalchemy import Column, Integer, MetaData, Table, bindparam, func, types
 from sqlalchemy import exc as sa_exc
 
 from pyathena.formatter import DefaultParameterFormatter
@@ -88,3 +88,36 @@ def test_generic_array_partial_update():
     table = array_table(types.ARRAY(Integer))
     sql = str(table.update().values({table.c["items"][1]: 2}).compile(dialect=AthenaDialect()))
     assert "SET items=transform(" in sql
+
+
+def test_write_index_expression_keeps_its_argument_types():
+    table = array_table()
+    index = func.length("abc")
+    statement = table.update().values({table.c["items"][index]: 9})
+    compiled = statement.compile(dialect=AthenaDialect())
+    params = {
+        name: compiled._bind_processors.get(name, lambda value: value)(value)
+        for name, value in compiled.params.items()
+    }
+    assert "length('abc')" in DefaultParameterFormatter().format(str(compiled), params)
+
+
+def test_ordered_partial_update_with_sql_expression():
+    table = array_table()
+    items = table.c["items"]
+    statement = table.update().ordered_values((items[2], items[1] + 1), (table.c.id, 2))
+    compiled = str(statement.compile(dialect=AthenaDialect()))
+    assert compiled.index("SET items=") < compiled.index(", id=")
+    assert "element_at(arrays.items" in compiled
+
+
+def test_binary_element_assignment_uses_native_hex_parameter():
+    table = array_table(AthenaArray(types.BINARY))
+    compiled = (
+        table.update().values({table.c["items"][1]: b"\x00\xff"}).compile(dialect=AthenaDialect())
+    )
+    params = {
+        name: compiled._bind_processors.get(name, lambda value: value)(value)
+        for name, value in compiled.params.items()
+    }
+    assert "FROM_HEX('00ff')" in DefaultParameterFormatter().format(str(compiled), params)

tests/sqlalchemy/test_suite.py

@@ -168,6 +168,27 @@ def test_nested_zero_indexed_and_cached_bindings(self, connection, metadata):
         with pytest.raises(sa_exc.DBAPIError):
             connection.execute(statement, {"row_id": 1, "outer": -1, "inner": 0, "value": 9})
 
+    def test_expression_values_and_indices(self, connection, metadata):
+        table = Table(
+            "array_expression_updates",
+            metadata,
+            Column("id", Integer),
+            Column("items", AthenaArray(Integer)),
+            Column("binary_items", AthenaArray(types.BINARY)),
+        )
+        table.create(connection)
+        connection.execute(table.insert().values(id=1, items=[1, 2, 3], binary_items=[b"abc"]))
+        items = table.c["items"]
+        connection.execute(
+            table.update().ordered_values(
+                (items[func.length("abc")], items[1] + 8),
+                (table.c.binary_items[1], b"\x00\xff"),
+                (table.c.id, 2),
+            )
+        )
+        connection.execute(table.update().values({items[1:2]: items[2:3].concat([4])}))
+        eq_(connection.execute(select(table)).one(), (2, [2, 9, 4, 9], [b"\x00\xff"]))
+
     def test_null_slice_binding_rejected(self, connection, metadata):
         table = Table("array_null_slice", metadata, Column("items", types.ARRAY(Integer)))
         table.create(connection)