refetch after insert and update

MySQL doesn't support returning after insert and update

Author: wwwjfy

mysql_tests/test_bind.py

@@ -40,11 +40,9 @@ async def test_unbind(aiomysql_pool):
 
 
 async def test_db_api(bind, random_name):
-    lastrowid = await db.all(
-        User.insert().values(name=random_name).execution_options(
-            return_lastrowid=True)
-    )
-    r = await db.scalar(User.select('nickname').where(User.id == lastrowid))
+    result = await db.first(User.insert().values(name=random_name))
+    assert result is None
+    r = await db.scalar(User.select('nickname').where(User.nickname == random_name))
     assert r == random_name
     assert (
         await db.first(User.query.where(User.nickname == random_name))

mysql_tests/test_json.py

@@ -130,7 +130,7 @@ class PropsTest(db.Model):
         assert t.obj["x"] == 1
         assert t.arr[-1] == 6
         assert await db.select(
-            [PropsTest.profile, PropsTest.raw, PropsTest.bool,]
+            [PropsTest.profile, PropsTest.raw, PropsTest.bool]
         ).gino.first() == (
             {
                 "arr": [3, 4, 5, 6],

src/gino/crud.py

@@ -789,49 +789,69 @@ async def _query_and_update(bind, item, query, cols, execution_opts):
     if bind._dialect.support_returning:
         # noinspection PyArgumentList
         query = query.returning(*cols)
-        row = await bind.first(query)
+
+    async def _execute_and_fetch(conn, query):
+        context, row = await conn._first_with_context(query)
+        if not bind._dialect.support_returning:
+            if context.isinsert:
+                table = context.compiled.statement.table
+                key_getter = context.compiled._key_getters_for_crud_column[2]
+                compiled_params = context.compiled_parameters[0]
+                last_row_id = context.get_lastrowid()
+                if last_row_id is not None:
+                    lookup_conds = [
+                        c == last_row_id
+                        if c is table._autoincrement_column
+                        else c == _cast_json(
+                            c, compiled_params.get(key_getter(c), None))
+                        for c in table.primary_key
+                    ]
+                else:
+                    lookup_conds = [
+                        c == _cast_json(
+                            c, compiled_params.get(key_getter(c), None))
+                        for c in table.columns
+                    ]
+                query = sa.select(table.columns).where(
+                    sa.and_(*lookup_conds)).execution_options(**execution_opts)
+                row = await conn.first(query)
+            elif context.isupdate:
+                if context.get_affected_rows() == 0:
+                    raise NoSuchRowError()
+                table = context.compiled.statement.table
+                if len(table.primary_key) > 0:
+                    lookup_conds = [
+                        c == _cast_json(
+                            c, item.__values__[
+                                item._column_name_map.invert_get(c.name)])
+                        for c in table.primary_key
+                    ]
+                else:
+                    lookup_conds = [
+                        c == _cast_json(
+                            c, item.__values__[
+                                item._column_name_map.invert_get(c.name)])
+                        for c in table.columns
+                    ]
+                query = sa.select(table.columns).where(
+                    sa.and_(*lookup_conds)).execution_options(**execution_opts)
+                row = await conn.first(query)
+        return row
+
+    if isinstance(bind, GinoConnection):
+        row = await _execute_and_fetch(bind, query)
     else:
-        # CAVEAT: MySQL doesn't support RETURNING. The workaround here is
-        # to get lastrowid and load it after insertion.
-        # Note that this only works for tables with AUTO_INCREMENT column
-        # For update queries, update using its primary key
-
-        # make insertion and select in one transaction to get the might-be
-        # "dirty" row
-        release_conn = False
-        if not isinstance(bind, GinoConnection):
-            conn = await bind.acquire(reuse=True)
-            release_conn = True
-        else:
-            conn = bind
-        try:
-            lastrowid, affected_rows = await conn.all(
-                query.execution_options(return_affected_rows=True)
-            )
-            if not lastrowid and not affected_rows:
-                raise NoSuchRowError()
-            # It's insertion and primary key is AUTO_INCREMENT
-            if lastrowid:
-                pkey = cls.__table__.primary_key
-                query = (
-                    sa.select(cols)
-                    .where(pkey.columns.values()[0] == lastrowid)
-                    .execution_options(**execution_opts)
-                )
-            else:
-                try:
-                    query = (
-                        sa.select(cols)
-                        .where(item.lookup())
-                        .execution_options(**execution_opts)
-                    )
-                except LookupError:  # no primary key
-                    return None
-            row = await conn.first(query)
-        finally:
-            if release_conn:
-                await conn.release()
+        async with bind.acquire(reuse=True) as conn:
+            row = await _execute_and_fetch(conn, query)
     if not row:
         raise NoSuchRowError()
     for k, v in row.items():
         item.__values__[item._column_name_map.invert_get(k)] = v
+
+
+def _cast_json(column, value):
+    # FIXME: for MySQL, json string in WHERE clause needs to be cast to JSON type
+    if (isinstance(column.type, sa.JSON) or
+            isinstance(getattr(column.type, 'impl', None), sa.JSON)):
+        return sa.cast(value, sa.JSON)
+    return value

src/gino/dialects/aiomysql.py

@@ -48,7 +48,8 @@ class AiomysqlDBAPI(base.BaseDBAPI):
 # noinspection PyAbstractClass
 class AiomysqlExecutionContext(base.ExecutionContextOverride, MySQLExecutionContext):
     def get_lastrowid(self):
-        return self.cursor.last_row_id
+        lastrowid = self.cursor.last_row_id
+        return None if lastrowid == 0 else lastrowid
 
     def get_affected_rows(self):
         return self.cursor.affected_rows

src/gino/dialects/base.py

@@ -196,7 +196,8 @@ def __init__(self, context):
     def context(self):
         return self._context
 
-    async def execute(self, one=False, return_model=True, status=False):
+    async def execute(self, one=False, return_model=True, status=False,
+                      return_context=False):
         context = self._context
 
         param_groups = []
@@ -213,31 +214,26 @@ async def execute(self, one=False, return_model=True, status=False):
             return await cursor.async_execute(
                 context.statement, context.timeout, param_groups, many=True
             )
+        args = param_groups[0]
+        if context.baked_query:
+            rows = await cursor.execute_baked(
+                context.baked_query, context.timeout, args, one
+            )
         else:
-            args = param_groups[0]
-            if context.baked_query:
-                rows = await cursor.execute_baked(
-                    context.baked_query, context.timeout, args, one
-                )
+            rows = await cursor.async_execute(
+                context.statement, context.timeout, args, 1 if one else 0
+            )
+        item = context.process_rows(rows, return_model=return_model)
+        if one:
+            if item:
+                item = item[0]
             else:
-                rows = await cursor.async_execute(
-                    context.statement, context.timeout, args, 1 if one else 0
-                )
-            if not self.context.dialect.support_returning and (
-                self.context.isinsert or self.context.isupdate
-            ):
-                if self.context.execution_options.get("return_affected_rows", False):
-                    return context.get_lastrowid(), context.get_affected_rows()
-                return context.get_lastrowid()
-            item = context.process_rows(rows, return_model=return_model)
-            if one:
-                if item:
-                    item = item[0]
-                else:
-                    item = None
-            if status:
-                item = cursor.get_statusmsg(), item
-            return item
+                item = None
+        if status:
+            return cursor.get_statusmsg(), item
+        if return_context:
+            return context, item
+        return item
 
     def iterate(self):
         if self._context.executemany:
@@ -296,10 +292,6 @@ def timeout(self):
     def loader(self):
         return self._compiled_first_opt("loader", None)
 
-    @util.memoized_property
-    def return_affected_rows(self):
-        return self._compiled_first_opt("return_affected_rows", False)
-
     def process_rows(self, rows, return_model=True):
         if not rows:
             return []
@@ -417,6 +409,12 @@ def _init_baked_query(cls, dialect, connection, dbapi_connection, bq, parameters
         self.baked_query = bq
         return self
 
+    def get_lastrowid(self):
+        raise NotImplementedError
+
+    def get_affected_rows(self):
+        raise NotImplementedError
+
 
 class AsyncDialectMixin:
     cursor_cls = DBAPICursor

src/gino/engine.py

@@ -348,6 +348,10 @@ async def first(self, clause, *multiparams, **params):
         result = self._execute(clause, multiparams, params)
         return await result.execute(one=True)
 
+    async def _first_with_context(self, clause, *multiparams, **params):
+        result = self._execute(clause, multiparams, params)
+        return await result.execute(one=True, return_context=True)
+
     async def one_or_none(self, clause, *multiparams, **params):
         """
         Runs the given query in database, returns at most one result.