[Python]PyMongoで動的クエリ生成

2019年8月4日

PyMongoとは

PythonからMongoDBを扱えるライブラリです。

CLIのmongoとは書き方が微妙に違うため、そのまま転記すると色々とハマりかねません。
具体的には、あくまでPythonで動いているので、各種クエリやオペレータをクオテーションで囲む必要があります。

ですが、これは要するに「辞書を引数にしている」ということなので、動的にクエリを生成することがわりと簡単にできます。
特に、ベタ書きすると発狂しかねないaggregateのpipelineを簡潔に書くことができます。

具体例 : メンテしやすいswitchクエリ生成

ユーザーごとのscoreが格納されたコレクションから、得点ごとにランク分けした結果を取得するようなクエリを発行するケースを考えます。
score 80以上ならS, 60以上ならA、40以上ならB、それ以外ならCというランクに振り分けるとします。

ベタ書き

ベタに書くと、クエリは以下のようになります。

pipe = [
    {"$project":
        {
            "_id": 0,
            "uid": 1,
            "rank":
            {
                "$switch":
                {
                    "branches": 
                    [
                        {
                            "case": {"$gte": ["$scores", 80]},
                            "then": "S"
                        },
                        {
                            "case": {"$gte": ["$scores", 60]},
                            "then": "A"
                        },
                        {
                            "case": {"$gte": ["$scores", 40]},
                            "then": "B"
                        },
                    ],
                    "default": "C"
                }
            }
        }
    }
]
col.aggregate(pipeline=pipe)

これはこれでいいかもしれませんが、
– ランクの文字列を変更したい
– 得点XX以上を変更したい
– ランク帯を増やしたい
などの変更がかかると、メンテが途端に面倒になります。

辞書型変数に切り出す

そこで、以下のように書き換えてみます。

# 設定項目
rank_params = {
    "S": 80,
    "A": 60,
    "B": 40
}
default_rank = "C"

# キーをリストで取得
ranks = rank_params.key()

# switchオペレータの引数(branches, default)の値を生成
branches = []
for rank in ranks:
    cond = {
        "case": {$gte: ["$scores", rank_params[rank]]},
        "then": rank
    }
    branches.append(cond)

switch = {
    "branches": branches,
    "default" : default_rank
}

pipe = [
    {"$project":
        {
            "_id": 0,
            "uid": 1,
            "rank":
            {
                "$switch": switch
            }
        }
    }
]
col.aggregate(pipeline=pipe)

値を調整したいときは、rank_paramsとdefault_rankの値を書き換えるだけでよくなります。

PyMongoにおけるクエリは、辞書やリストの集まりなので、dict.keys()やforなどを使うことで、メンテしやすい書き方ができます。
“$gte”などのオペレータも文字列なので、条件に合わせて切り替える、といったこともできそうです。

当初仕組みがよくわからないまま書いてしまった非効率なクエリたちをリファクタリングしていきたいところ。

参考

2019/08/04修正

一部コードが誤っていたため修正しました。申し訳ございません

修正点 : $switch句における$gteの書き方

# 修正前
...
"case": {"score": {"$gte": 80}}, "then": "S" }
...

# 修正後
...
"case": {"$gte": ["$scores", 80]}, "then": "S"}
...

開発aggregate, PyMongo, python, switch