たごもりすメモ

コードとかその他の話とか。

run_in_transaction内でトランザクション外データストアアクセスを行う方法

基本的なトランザクションの話をすこし

GAE/Pythonではトランザクション処理を行うのに db.run_in_transaction() を使う。典型的な使いかたは以下のような感じで、指定した関数内の処理がまるごとトランザクション内として扱われ、関数から出たところで自動的にCommitされる。

app1> x1 = X(key_name="x1")
app1> x1.put()
datastore_types.Key.from_path(u'X', u'x1', _app=u'tagomoris-test')
app1> def tx1():
...   x = X.get_by_key_name("x1")
...   y = Y(parent=x)
...   y.put()
...   
app1> db.run_in_transaction(tx1)
app1> 

GAEのトランザクションには強い制約がいくつかあるが、そのうち最大のものは「トランザクション内で操作するエンティティはすべて同じエンティティグループに属さなければならない」というものがある。エンティティグループとはエンティティのキーの親子関係によって作られるグループで、エンティティの種類などは関係ない。
以下のような操作はこのお約束を破っているので、例外が出る。

app1> class XE(db.Model):
...   pass
... 
app1> class XXE(db.Model):
...   xe = db.ReferenceProperty(XE)
... 
app1> xe1 = XE(key_name="xe1")
app1> xe1.put()
datastore_types.Key.from_path(u'XE', u'xe1', _app=u'tagomoris-test')
app1> def tx1():
...   xe1 = XE.get_by_key_name("xe1")
...   xxe1 = XXE(xe=xe1)
...   xxe1.put()
... 
app1> db.run_in_transaction(tx1)
Traceback (most recent call last):
  File "<console>", line 1, in <module>
  File "/home/sigh/google_appengine/google/appengine/api/datastore.py", line 2132, in RunInTransaction
    DEFAULT_TRANSACTION_RETRIES, function, *args, **kwargs)
  File "/home/sigh/google_appengine/google/appengine/api/datastore.py", line 2223, in RunInTransactionCustomRetries
    result = function(*args, **kwargs)
  File "<console>", line 4, in tx1
  File "/home/sigh/google_appengine/google/appengine/ext/db/__init__.py", line 805, in put
    return datastore.Put(self._entity, rpc=rpc)
  File "/home/sigh/google_appengine/google/appengine/api/datastore.py", line 278, in Put
    tx = _MaybeSetupTransaction(req, keys)
  File "/home/sigh/google_appengine/google/appengine/api/datastore.py", line 2315, in _MaybeSetupTransaction
    raise _DifferentEntityGroupError(expected_group, group)
  File "/home/sigh/google_appengine/google/appengine/api/datastore.py", line 2361, in _DifferentEntityGroupError
    b.kind(), id_or_name(b)))
BadRequestError: Cannot operate on different entity groups in a transaction: (kind=u'XE', name=u'xe1') and (kind=u'XXE', id=None).
app1> 

「えーなんで? XXEのプロパティxeでXEのエンティティを参照してるじゃん!」とか言ってもダメ。エンティティグループはあくまでキーの親子関係により規定されるので、プロパティで何をやろうとも関係はない。
エンティティグループの道は険しい。

GAE/PythonとGAE/Jの違い

GAE/Jの場合はデータストア操作を行うときに「どのトランザクション中の処理として実行するか」をコード上で指定するようになっており、複数トランザクションをひとつのプロセスから並行して走らせる、ということができる。(slim3のグローバルトランザクションなんかは完全にこの構造を利用して成立している。)

いっぽうGAE/Pythonトランザクション内で処理を実行するインターフェイスが db.run_in_transaction() しかなく、この内部で実行されるコードはすべて単一のトランザクションとして実行される。
少し注意すれば問題ない場合が多いのだが、しかし実は db.Model.get_or_insert() が内部的にトランザクションを使っていたりするため、迂闊にrun_in_transaction()内で実行したりして「Nested Transactionはダメよ」とか言われてがっくりきたりする。なんだよ関係ないエンティティひとつくらいget/putさせてよ、とか思ってもダメなものはダメなのだ。
GAE/Pythonトランザクションの道は険しい。

GAE/Pythonのrun_in_transaction内でトランザクション外のデータストア操作を行う

さて本題。GAE/Python環境でもちょっとくらいトランザクションから外れた処理をしたいこともある!(かもしれない!)という要求が自分にあったので、作ってみた。
まず以下のような関数を用意する。用途別にふたつ。(run_notin_transaction / run_notin_tx)

def _dummy_CurrentTransactionKey():
    return None

def run_notin_transaction(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        _current_tx_key_func = datastore._CurrentTransactionKey
        datastore._CurrentTransactionKey = _dummy_CurrentTransactionKey
        try:
            return func(*args, **kwargs)
        finally:
            datastore._CurrentTransactionKey = _current_tx_key_func
    return wrapper

def run_notin_tx(f):
    _current_tx_key_func = datastore._CurrentTransactionKey
    datastore._CurrentTransactionKey = _dummy_CurrentTransactionKey
    try:
        return f()
    finally:
        datastore._CurrentTransactionKey = _current_tx_key_func

run_notin_transactionは関数デコレータとして使うもので、run_notin_txはlambdaを受け取って即座に実行するもの。以下のようにして使う。

app1> class P(db.Model): pass
... 
app1> def tx2():
...   p1 = P()
...   run_notin_tx(lambda: P().put())
...   p1.put()
... 
app1> db.run_in_transaction(tx2)
app1> 

このコードを実行すると P というkindのエンティティがふたつputされる。ふたつのエンティティはそれぞれ別のエンティティグループに所属している。
このコードは動作する。ただしremote_apiのみで(もしくは同等のコードをデプロイしたproduction環境のみで)。datastore_file_stubやdatastore_sqlite_stubを使用する各環境では動作しない。(トランザクションの実装がなんちゃってなので。)

関数デコレータ版の方は以下のような感じ。

app1> @run_notin_transaction
... def isolated_put():
...   p2 = P()
...   p2.put()
... 
app1> def tx3():
...   p3 = P()
...   isolated_put()
...   p3.put()
... 
app1> db.run_in_transaction(tx3)
app1> 

このコードでも同じように、エンティティグループの異なるエンティティがふたつputされる。

何に使うの?

自分はDatastoreを執拗にいじりたかったのでこういうものが欲しかったけど、まあ要望はあるにはあるかもしれない。例えば前述の get_or_insert を実行したい!というケースとか。外部でやりゃいいじゃんって場合がほとんどなんだろうけど、その場でやらせてくれてもいいじゃんね、という。
他に考えつくのはログ的なエンティティをput()するのにトランザクション内外だのを注意していたくない、とかそういう用途だろうか。

実はあまりにニッチすぎて公開する気もなかったんだけど、用途はあるはず、というコメントをもらったので書いてみた次第。どうなのかなー。あるんですかね。使ってみたかった!という人は教えてください。(笑)

どうやってるの?

GAE/Pythonトランザクション機構は実に単純で、データストア操作が発生した瞬間に自身のプロセスの実行フレーム(スタックフレーム)を上にたどっていき、run_in_transaction関数*1がスタックフレームに積まれていれば、そのデータストア操作はトランザクション内で実行されているものとみなす。*2トランザクション複数同時に実行されることはない、という構造なので、この程度の機構でいいんですな。具体的なコードは以下。(google.appengine.api.datastore)

def _FindTransactionFrameInStack():
  """Walks the stack to find a RunInTransaction() call.

  Returns:
    # this is the RunInTransactionCustomRetries() frame record, if found
    frame record or None
  """
  frame = sys._getframe()
  filename = frame.f_code.co_filename

  frame = frame.f_back.f_back
  while frame:
    if (frame.f_code.co_filename == filename and
        frame.f_code.co_name == 'RunInTransactionCustomRetries'):
      return frame
    frame = frame.f_back

  return None

_CurrentTransactionKey = _FindTransactionFrameInStack

Pythonである以上、どのモジュールのどんな関数でも動的に入れ替えることが可能。トランザクション外でコードを実行したかったら、その間は datastore._CurrentTransactionKey() が常に None を返すようにしてやればいい、というのが前述の run_notin_transaction() がやっていたことだ。Noneを返すダミーに入れ替えてからトランザクション外の処理を実行し、完了したら _CurrentTransactionKey の定義を戻しておけばいい。*3

どうだこの自由っぷり! Pythonばんざい! みんなもっとAppEngineでPython使おうぜ!

*1:正確にはgoogle.appengine.api.datastore内でのみ使用されるRunInTransactionCustomRetries関数

*2:どのトランザクションか、というキーにはその実行フレームのオブジェクトそのものを使う!

*3:試すのも恐ろしいが、忘れると多分ひどいことになる。