GoogleAppsEngineでDatastoreの更新をasyncにやる書き方

タイトルで言いたいことを書きました。Python2.7

from google.appengine.ext import ndb

class TestModel(ndb.Model):
    name = ndb.StringProperty()

@ndb.transactional_tasklet
def txn (key):
    entity = yield key.get_async()
    if not entity:
      entity = TestModel(key=key)
    entity.populate(key=key, name='foo')
    yield entity.put_async()
    raise ndb.Return(entity)

def put_multi_async(keys):
    @ndb.tasklet
    def _put_multi_async(keys):
        results = yield [txn(key) for key in keys]
        raise ndb.Return(results)

    return _put_multi_async(keys).get_result()

def main():
  key1 = ndb.Key(TestModel, 1)
  key2 = ndb.Key(TestModel, 2)
  key3 = ndb.Key(TestModel, 3)

  res = put_multi_async([key1, key2, key3])
  print res

main()

'''
=>
[TestModel(key=Key('TestModel', 1), name=u'foo'),
 TestModel(key=Key('TestModel', 2), name=u'foo'),
 TestModel(key=Key('TestModel', 3), name=u'foo')]
'''