Kindless Ancestor Queryをローカル環境で使用する #appengine

AppEngineで実装されているQueryで、Kind指定なしで親キーのみ指定して、親キーに属するEntityGroupをKindをまたいでゴッソリ持ってくるという便利なクエリがあるのですが、開発環境で動かないという問題を抱えています。開発環境で動いてくれないって事はデプロイ前にテストできないって事で、結局プロダクトコードに含めることもできずまだ使えないカンジで勿体無いです。
makeSyncCallという仕組みで開発環境とProduction環境のDatastoreを直結するとテストできない事も無いですが、無理矢理過ぎるのでそれは考慮してません

ApiProxy.Delegateの出番

このブログを読んでくれている方はもう聞き飽きておられる事でしょうから詳しく説明しませんが、アレです、フックします。んでProtocolBufferをゴリゴリと触ります。

今回は、以下のようにいじってやる作戦です。もちろん、結局は環境を弄っているワケなので、必ずProduction環境と同じというワケでもないですし、多少の違いが出てしまうのは確かですが、それを言い始めると開発環境はProduction環境とはそもそも違うという割り切りはできて触っているワケで、今回のこのエントリの内容に関してもその割り切りができればそこそこ便利かもしれません。

1. フックして、datastore_v3#RunQueryと関係がないものは保存しておいたApiProxy#getDelegate()に処理を委譲する。
2. サービス側へリクエストされるbyte配列を奪って、ProtocolBufferオブジェクトとして組立て直す。それにより、AncestorKeyが指定されているか？Kindが指定されていないか？というチェックを行うことができる。これらの条件を満たさない場合はkindless ancestor queryでは無いので、これまた保存しておいたApiProxy#getDelegate()に処理を委譲する。
3. ここまですり抜けて来た処理は、kindless ancestor queryであるとわかるので、違うクエリに組み立てる。

違うクエリと書いていますが、実際に何をするかと言うとKind指定アリのAncestorクエリをKindの数だけ実行するという処理を行います。毎回クエリ結果のProtocolBufferオブジェクトを組み立てなおして、最後に全てのクエリの結果をマージして返します。かなり昔に書いた、開発環境でのみ実行可能なgetSchemaというデータストア内の全てのKindを取得できる機能を利用しています。

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Future;

import com.google.appengine.api.datastore.dev.LocalDatastoreService;
import com.google.appengine.tools.development.ApiProxyLocalImpl;
import com.google.apphosting.api.ApiProxy;
import com.google.apphosting.api.DatastorePb;
import com.google.apphosting.api.ApiProxy.ApiConfig;
import com.google.apphosting.api.ApiProxy.ApiProxyException;
import com.google.apphosting.api.ApiProxy.Environment;
import com.google.apphosting.api.ApiProxy.LogRecord;
import com.google.apphosting.api.DatastorePb.GetSchemaRequest;
import com.google.apphosting.api.DatastorePb.QueryResult;
import com.google.apphosting.api.DatastorePb.Schema;
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
import com.google.storage.onestore.v3.OnestoreEntity.Reference;
import com.google.storage.onestore.v3.OnestoreEntity.Path.Element;

/**
 * kindless ancestor queryをローカルでも動作させる{@link ApiProxy.Delegate}の実装.
 * @author shin1ogawa
 */
public class EnableAncestorQueryDelegate implements ApiProxy.Delegate<Environment> {

  @SuppressWarnings("unchecked")
  ApiProxy.Delegate<Environment> before = ApiProxy.getDelegate();

  public byte[] makeSyncCall(Environment env, String service, String method, byte[] request)
      throws ApiProxyException {
    if (service.equals("datastore_v3") == false || method.equals("RunQuery") == false) {
      return before.makeSyncCall(env, service, method, request);
    }
    // kindless ancestor queryかどうかをチェックするために、pbを組立て直す。
    DatastorePb.Query requestPb = new DatastorePb.Query();
    requestPb.mergeFrom(request);
    if (requestPb.getAncestor() == null || requestPb.getKind().isEmpty() == false) {
      return before.makeSyncCall(env, service, method, request);
    }
    
    // kindless ancestor queryの時は全てのKindに対してancestor queryを実行する。
    String[] kinds = getKinds();
    // ancestor keyのpbからancestorのkindを取得し、最初はancestor kindから実行する。
    Reference ancestor = requestPb.getAncestor();
    String ancestorKind = ancestor.getPath().getElement(0).getType();
    QueryResult resultOfAncestor =
        runAncestorQuery(env, service, method, ancestorKind, requestPb);
    resultOfAncestor.setMoreResults(true);
    for (String kind : kinds) {
      if (kind.equals(ancestorKind) == false) {
        // ancestor kind以外の全てのkindに対してancestorクエリを実行する
        QueryResult result = runAncestorQuery(env, service, method, kind, requestPb);
        Iterator<EntityProto> i = result.resultIterator();
        while (i.hasNext()) {
          EntityProto next = i.next();
          resultOfAncestor.addResult(next);
        }
      }
    }
    return resultOfAncestor.toByteArray();
  }

  DatastorePb.QueryResult runAncestorQuery(Environment env, String service, String method,
      String kind, DatastorePb.Query requestPb) {
    requestPb.setKind(kind);
    byte[] requestbytes = requestPb.toByteArray();
    DatastorePb.QueryResult result = new DatastorePb.QueryResult();
    result.mergeFrom(before.makeSyncCall(env, service, method, requestbytes));
    return result;
  }

  String[] getKinds() {
    LocalDatastoreService datastoreService =
        (LocalDatastoreService) ((ApiProxyLocalImpl) before).getService("datastore_v3");
    Schema schema =
        datastoreService.getSchema(null, new GetSchemaRequest().setApp(ApiProxy
          .getCurrentEnvironment().getAppId()));
    List<EntityProto> entityProtoList = schema.kinds();
    List<String> kindList = new ArrayList<String>(entityProtoList.size());
    for (EntityProto entityProto : entityProtoList) {
      List<?> path = entityProto.getKey().getPath().elements();
      Element element = (Element) path.get(path.size() - 1);
      kindList.add(element.getType());
    }
    return kindList.toArray(new String[0]);
  }

  public void log(Environment env, LogRecord logRecord) {
    before.log(env, logRecord);
  }

  public Future<byte[]> makeAsyncCall(Environment env, String service, String method,
      byte[] request, ApiConfig config) {
    return before.makeAsyncCall(env, service, method, request, config);
  }
}

やっつけで書いたので、getSchemaを実行するためにLocalDatastoreServiceが必要になり、ApiProxy#getDelegate()が必ずApiProxyLocalImplじゃないとダメ（つまり、他のDelegateが設定されているとダメ）という問題がありますが、そのあたりは外部からLocalDatastoreServiceを渡すなどの工夫をすればもっと便利になりますね。なかなか便利そうな気がします。