SA20: Add CrateDB-specific patches to CrateCompilerSA20

This effectively concludes the support for SQLAlchemy 2.0, by
implementing the same strategy as for the previous versions:
After vendoring the vanilla dialect's compiler's `visit_update` and
`_get_crud_params` methods, they are patched at a few spots to
accommodate CrateDB's features.

The changed behavior is mostly for running update statements on nested
`OBJECT` or `ARRAY` data types.

Author: amotl

src/crate/client/sqlalchemy/compat/core20.py

@@ -21,8 +21,9 @@
 
 from typing import Any, Dict, List, MutableMapping, Optional, Tuple, Union
 
+import sqlalchemy as sa
 from sqlalchemy import ColumnClause, ValuesBase, cast, exc
-from sqlalchemy.sql import crud, dml
+from sqlalchemy.sql import dml
 from sqlalchemy.sql.base import _from_objects
 from sqlalchemy.sql.compiler import SQLCompiler
 from sqlalchemy.sql.crud import (REQUIRED, _as_dml_column, _create_bind_param,
@@ -47,6 +48,11 @@ def visit_update(self, update_stmt, **kw):
         )
         update_stmt = compile_state.statement
 
+        # [20] CrateDB patch.
+        if not compile_state._dict_parameters and \
+                not hasattr(update_stmt, '_crate_specific'):
+            return super().visit_update(update_stmt, **kw)
+
         toplevel = not self.stack
         if toplevel:
             self.isupdate = True
@@ -87,7 +93,8 @@ def visit_update(self, update_stmt, **kw):
         table_text = self.update_tables_clause(
             update_stmt, update_stmt.table, render_extra_froms, **kw
         )
-        crud_params_struct = crud._get_crud_params(
+        # [20] CrateDB patch.
+        crud_params_struct = _get_crud_params(
             self, update_stmt, compile_state, toplevel, **kw
         )
         crud_params = crud_params_struct.single_params
@@ -105,12 +112,38 @@ def visit_update(self, update_stmt, **kw):
         text += table_text
 
         text += " SET "
+
+        # [20] CrateDB patch begin.
+        include_table = extra_froms and \
+            self.render_table_with_column_in_update_from
+
+        set_clauses = []
+
+        for c, expr, value, _ in crud_params:
+            key = c._compiler_dispatch(self, include_table=include_table)
+            clause = key + ' = ' + value
+            set_clauses.append(clause)
+
+        for k, v in compile_state._dict_parameters.items():
+            if isinstance(k, str) and '[' in k:
+                bindparam = sa.sql.bindparam(k, v)
+                clause = k + ' = ' + self.process(bindparam)
+                set_clauses.append(clause)
+
+        text += ', '.join(set_clauses)
+        # [20] CrateDB patch end.
+
+        """
+        # TODO: Complete SA20 migration.
+        # This is the column name/value joining code from SA20.
+        # It may be sensible to use this procedure instead of the old one.
         text += ", ".join(
             expr + "=" + value
             for _, expr, value, _ in cast(
                 "List[Tuple[Any, str, str, Any]]", crud_params
             )
         )
+        """
 
         if self.implicit_returning or update_stmt._returning:
             if self.returning_precedes_values:
@@ -356,6 +389,24 @@ def _get_crud_params(
             kw,
         )
 
+    # [20] CrateDB patch.
+    #
+    # This sanity check performed by SQLAlchemy currently needs to be
+    # deactivated in order to satisfy the rewriting logic of the CrateDB
+    # dialect in `rewrite_update` and `visit_update`.
+    #
+    # It can be quickly reproduced by activating this section and running the
+    # test cases::
+    #
+    #   ./bin/test -vvvv -t dict_test
+    #
+    # That croaks like::
+    #
+    #   sqlalchemy.exc.CompileError: Unconsumed column names: characters_name
+    #
+    # TODO: Investigate why this is actually happening and eventually mitigate
+    #       the root cause.
+    """
     if parameters and stmt_parameter_tuples:
         check = (
             set(parameters)
@@ -367,6 +418,7 @@ def _get_crud_params(
                 "Unconsumed column names: %s"
                 % (", ".join("%s" % (c,) for c in check))
             )
+    """
 
     if (
         _compile_state_isinsert(compile_state)

src/crate/client/sqlalchemy/dialect.py

@@ -32,7 +32,7 @@
     CrateDDLCompiler
 )
 from crate.client.exceptions import TimezoneUnawareException
-from .sa_version import SA_VERSION, SA_1_4
+from .sa_version import SA_VERSION, SA_1_4, SA_2_0
 from .types import Object, ObjectArray
 
 TYPES_MAP = {
@@ -155,7 +155,10 @@ def process(value):
 }
 
 
-if SA_VERSION >= SA_1_4:
+if SA_VERSION >= SA_2_0:
+    from .compat.core20 import CrateCompilerSA20
+    statement_compiler = CrateCompilerSA20
+elif SA_VERSION >= SA_1_4:
     from .compat.core14 import CrateCompilerSA14
     statement_compiler = CrateCompilerSA14
 else:

src/crate/client/sqlalchemy/sa_version.py

@@ -25,3 +25,4 @@
 SA_VERSION = V(sa.__version__)
 
 SA_1_4 = V('1.4.0b1')
+SA_2_0 = V('2.0.0b1')

src/crate/client/sqlalchemy/tests/bulk_test.py

@@ -78,4 +78,4 @@ def test_bulk_save(self):
             ('Banshee', 26),
             ('Callisto', 37)
         )
-        self.assertEqual(expected_bulk_args, bulk_args)
+        self.assertSequenceEqual(expected_bulk_args, bulk_args)