🐍

gcloud-aio-bigqueryでasync/awaitを使ってBigQueryのクエリを叩く

2024/01/09に公開

Python

BigQuery

tech

PythonからBigQueryを叩くには公式の google-cloud-bigquery があるのですが、asyncioに対応していません。

一部のライブラリではv2側でasyncioに対応しているのですが、BigQueryはv2でも対応していないように見えます。ちなみにv1はほとんどの実装でREST APIを叩いて、v2はほとんどの実装でRPC APIを叩いています（どちらも全部確認したわけではありません）。一部のライブラリでv2のほうを叩くとサーバーエラーが多発したこともあり、今のところなるべくv1でasyncioを使用したいです。

GCPのPythonライブラリをasyncioで動くように書き換えたものが gcloud-aio で用意されており、その中に gcloud-aio-bigquery があります。あまり使われているのを見かけませんが、AirflowのGCP周りの一部で使われているようです。

ただしREST APIそのままの実装になっているため、REST APIの知識が必要です。手軽に使えるとは言い難いです。

REST APIでの流れ

まずクエリを実行するには jobs.query を呼びます。

代表的なパラメーターは以下の通りです。

query : 必須。SQLクエリを指定する。
maxResults : １レスポンスに含まれる行数。指定しなくてもバイト数でページングされるが、指定しておいたほうが安全。
timeoutMs : タイムアウトの秒数
useLegacySql : BigQueryのレガシーSQLを使用するかどうか。デフォルトtrueなのでfalseを指定しておいたほうがいい。

実行するとジョブが作成されます。

レスポンスを確認してデータが完結しているか確認します。 jobComplete が true の場合、ジョブが完了して結果の行を取得できる状態です。この場合、 rows に結果の行が含まれています。また、 schema にスキーマが含まれています。 rows だけ見てもどの列の何のデータがの判別が難しいので、 schema と合わせて確認します。

jobComplete が true で、 pageToken が含まれている場合、まだ読み込むべきデータが残っているので続きを読み込みます。 pageToken が無ければ終了です。

データが完結していない場合は、レスポンスの jobReference の jobId を拾って、（ pageToken があればそれも一緒に） jobs.getQueryResults を実行します。結果は jobs.query と同じです。

全体の流れをまとめると以下のようになります。

jobs.query でジョブを作成
jobComplete が true なら
- rows がない（totalRows=0）→終了
- rows がある→そのデータを使う
- pageToken がない→終了
jobComplete が false 、または pageToken がある→ jobs.getQueryResults を実行し、上の判定ロジックに戻る

実際に使ってみる

実際に動くものが以下になります。

import asyncio
from collections.abc import AsyncGenerator

import aiohttp
from gcloud.aio.bigquery import Job
from google.cloud import bigquery
from google.cloud.bigquery._helpers import _rows_from_json
from google.cloud.bigquery.table import _parse_schema_resource  # type:ignore


async def query(
    sql: str, max_results: int = 100
) -> AsyncGenerator[list[bigquery.Row], None]:
    loop = asyncio.get_running_loop()
    async with aiohttp.ClientSession(loop=loop) as session:
        job = Job(session=session)
        query_request = {
            "query": sql,
            "maxResults": max_results,
            "useLegacySql": "false",
        }
        response = await job.query(query_request=query_request)
        while True:
            errors = response.get("errors")
            page_token = response.get("pageToken")
            if errors is not None:
                raise RuntimeError(errors)
            if response.get("jobComplete"):
                if int(response.get("totalRows", 0)) == 0:
                    return
                schema = _parse_schema_resource(response.get("schema", {}))
                yield _rows_from_json(response.get("rows", ()), schema)
                if page_token is None:
                    return
            else:
                await asyncio.sleep(1.0)  # 終わってなさそうなので少し待つ
            # ジョブの結果を確認
            job = Job(
                job_id=response["jobReference"]["jobId"],
                project=response["jobReference"]["projectId"],
                session=session,
            )
            params = {
                "location": response["jobReference"]["location"],
                "maxResults": max_results,
                "pageToken": page_token,
            }
            response = await job.get_query_results(params=params)


async def main() -> None:
    async for rows in query("SELECT * FROM `project_id.dataset_id.table`"):
        for row in rows:
            # row["key"] で値を取得可能
            for key, value in row.items():
                print(f"{key}={value}")


if __name__ == "__main__":
    asyncio.run(main())

雑に書いているので必要に応じて書き換えてください。タイムアウトの設定や例外の処理などが別途必要になります。

結果行の整形に公式の google-cloud-bigquery も使っています。また、 aiohttp パッケージも必要です。

プロジェクトIDを指定したい場合は Job のコンストラクタで指定してもいいですし、環境変数 GOOGLE_CLOUD_PROJECT などで指定しても反映されます。

まとめ

公式で対応してほしいです。

REST APIでの流れ

実際に使ってみる

まとめ

Discussion