add update compiler

Author: iurii iurii

src/drivers/base.py

@@ -3,7 +3,7 @@
 import six
 from sqlalchemy import schema, types as sqltypes, exc, util as sa_util
 from sqlalchemy.engine import default, reflection
-from sqlalchemy.sql import compiler, expression, type_api
+from sqlalchemy.sql import compiler, expression, type_api, crud
 from sqlalchemy.types import DATE, DATETIME, INTEGER, VARCHAR, FLOAT
 
 from .. import types
@@ -179,6 +179,26 @@ def visit_join(self, join, asfrom=False, **kwargs):
              "ON", self.process(join.onclause, asfrom=True, **kwargs)
              ))
 
+    def visit_update(self, update_stmt, asfrom=False, **kw):
+        text = 'ALTER TABLE '
+        table_text = self.update_tables_clause(update_stmt, update_stmt.table, [], **kw)
+        text += table_text
+        text += ' UPDATE '
+        crud_params = crud._setup_crud_params(
+            self, update_stmt, crud.ISUPDATE, include_table=False, **kw)
+
+        text += ', '.join(
+            c[0]._compiler_dispatch(self,
+                                    include_table=False) +
+            '=' + c[1] for c in crud_params
+        )
+
+        if update_stmt._whereclause is not None:
+            t = update_stmt._whereclause._compiler_dispatch(self, include_table=False)
+            if t:
+                text += " WHERE " + t
+        return text
+
 
 class ClickHouseDDLCompiler(compiler.DDLCompiler):
     def visit_create_column(self, create, **kw):

tests/orm/test_update.py

@@ -0,0 +1,36 @@
+from sqlalchemy import Column, update
+
+from src import types, Table
+from tests.testcase import BaseTestCase
+
+
+class UpdateTestCase(BaseTestCase):
+    def create_table(self):
+        metadata = self.metadata()
+
+        return Table(
+            't1', metadata,
+            Column('x', types.Int32, primary_key=True),
+            Column('y', types.Int32, primary_key=True)
+        )
+
+    def test_update(self):
+        table = self.create_table()
+        query = table.update().values(x=3).where(table.c.x.in_([1, 2]))
+        self.assertEqual(
+            self.compile(query),
+            'ALTER TABLE t1 UPDATE x=%(x)s WHERE x IN (%(x_1)s, %(x_2)s)'
+        )
+        query = update(table).values(x=3).where(table.c.x.in_([1, 2]))
+        self.assertEqual(
+            self.compile(query),
+            'ALTER TABLE t1 UPDATE x=%(x)s WHERE x IN (%(x_1)s, %(x_2)s)'
+        )
+
+    def test_update_from_own_field(self):
+        table = self.create_table()
+        query = table.update().values(x=table.c.y).where(table.c.x.in_([1, 2]))
+        self.assertEqual(
+            self.compile(query),
+            'ALTER TABLE t1 UPDATE x=y WHERE x IN (%(x_1)s, %(x_2)s)'
+        )