Add extra_returning_fields in apply

Author: kigawas

src/gino/crud.py

@@ -102,7 +102,7 @@ def _set_prop(self, prop, value):
             self._literal = False
         self._props[prop] = value
 
-    async def apply(self, bind=None, timeout=DEFAULT):
+    async def apply(self, bind=None, timeout=DEFAULT, extra_returning_fields=tuple()):
         """
         Apply pending updates into database by executing an ``UPDATE`` SQL.
 
@@ -113,6 +113,9 @@ async def apply(self, bind=None, timeout=DEFAULT):
           ``None`` for wait forever. By default it will use the ``timeout``
           execution option value if unspecified.
 
+        :param extra_returning_fields: A `tuple` of returning fields besides
+          fields to create/update, e.g. (`updated_at`, `created_at`)
+
         :return: ``self`` for chaining calls.
 
         """
@@ -174,9 +177,9 @@ async def apply(self, bind=None, timeout=DEFAULT):
             )
             .execution_options(**opts)
         )
-        await _query_and_update(
-            bind, self._instance, clause, [getattr(cls, key) for key in values], opts
-        )
+        cols = tuple(getattr(cls, key) for key in values)
+        extra_cols = tuple(getattr(cls, key) for key in extra_returning_fields)
+        await _query_and_update(bind, self._instance, clause, cols + extra_cols, opts)
         for prop in self._props:
             prop.reload(self._instance)
         return self

tests/models.py

@@ -8,6 +8,7 @@
 
 from gino import Gino
 from gino.dialects.asyncpg import JSONB
+from sqlalchemy import func
 
 DB_ARGS = dict(
     host=os.getenv("DB_HOST", "localhost"),
@@ -49,6 +50,12 @@ class User(db.Model):
     weight = db.IntegerProperty(prop_name='parameter')
     height = db.IntegerProperty(default=170, prop_name='parameter')
     bio = db.StringProperty(prop_name='parameter')
+    updated_at = db.Column(
+        db.DateTime(timezone=True),
+        nullable=False,
+        onupdate=func.now(),
+        server_default=func.now(),
+    )
 
     @balance.after_get
     def balance(self, val):

tests/test_crud.py

@@ -209,6 +209,20 @@ async def test_update_multiple_primary_key(engine):
     assert f2
 
 
+async def test_update_extra_returing_fields(engine):
+    u1 = await test_create(engine)
+    updated1 = u1.updated_at
+    await u1.update(nickname="new_nickname").apply(bind=engine)
+    assert updated1 == u1.updated_at
+
+    u2 = await test_create(engine)
+    updated2 = u2.updated_at
+    await u2.update(nickname="new_nickname").apply(
+        bind=engine, extra_returning_fields=("updated_at",)
+    )
+    assert updated2 < u2.updated_at
+
+
 async def test_delete(engine):
     u1 = await test_create(engine)
     await u1.delete(bind=engine, timeout=10)