はじめに

これは Kyash Advent Calendar 2021 の7日目の記事です

Kyashで、経理・会計・資金決済法の担当エンジニアをしているPrinceチームの清水(@thimi0412)です。
Princeチームについてはこちら → https://foxtoy.hateblo.jp/entry/2021/07/19/100223

KyashではGoogle Cloud Composerを使用して会計業務に使うViewテーブルの作成などをおこなっています。
その中でBigQueryに対してクエリを実行し、結果をスプレッドシートに書き込み指定された共有ドライブのフォルダ内に入れる等の操作が必要になりCloudComposerでGoogle Workspace系のAPIを使う機会があったのでハマった点やどう使っているかを書きます。

Cloud Composerとは
Apache Airflow(ワークフロー管理ツール) で構築された、フルマネージドのワークフローオーケストレーションサービスです。
https://cloud.google.com/composer

Google Workspace系のOperatorは共有ドライブに対応していないので注意

スプレッドシートを作成したのちそのシートに対してBigQueryのクエリ結果を書き込む処理をしたいと思ったのでGoogleSheetsCreateSpreadsheetOperatorを使用してスプレッドシートを作成したのですがフォルダIDが見つからないというエラーが出ました。

Operatorの使用の例

SPREADSHEET = {
    "name": sheet_name, # シート名
    "parents": [parent_folder_id], # シートを作成するフォルダID
    "mimeType": "application/vnd.google-apps.spreadsheet",
}

create_spreadsheet = GoogleSheetsCreateSpreadsheetOperator(
    task_id="create_spreadsheet", spreadsheet=SPREADSHEET
)

ググっても情報が出てこなかったのでairflowのソースコードを調べました。
分かったことはcreate()の引数としてsupportsAllDrives=trueのパラメータがセットされていないことが問題でした。
共有ドライブに作成する際にはこのパラメータをセットする必要があります(他のGoogle Workspace系のOperatorも同様)。なんで対応してないんだ?
ref Implement shared drive support

def create_spreadsheet(self, spreadsheet: Dict[str, Any]) -> Dict[str, Any]:
    """
    Creates a spreadsheet, returning the newly created spreadsheet.
    :param spreadsheet: an instance of Spreadsheet
        https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets#Spreadsheet
    :type spreadsheet: Dict[str, Any]
    :return: An spreadsheet object.
    """
    self.log.info("Creating spreadsheet: %s", spreadsheet['properties']['title'])

    response = (
        self.get_conn().spreadsheets().create(body=spreadsheet).execute(num_retries=self.num_retries)
    )
    self.log.info("Spreadsheet: %s created", spreadsheet['properties']['title'])
    return response

GoogleSheetsCreateSpreadsheetOperatorは共有ドライブには対応していないのでスプレッドシート作成の関数を自作しました。
この関数をPythonOperator内から呼び出してスプレッドシートを作成しています。フォルダを作成する場合はmimeTypeをapplication/vnd.google-apps.folderにすればフォルダも作成できます。

from airflow.providers.google.suite.hooks.sheets import GSheetsHook

def create_spread_sheet(sheet_name, parent_folder_id):
    """スプレッドシートを作成する

    Args:
        sheet_name (string): シート名
        parent_folder_id (string): シートを作成する親フォルダのID

    Returns:
        [type]: [description]
    """
    hook = GoogleDriveHook()
    drive = hook.get_conn()
    file_metadata = {
        "name": sheet_name,
        "parents": [parent_folder_id],
        "mimeType": "application/vnd.google-apps.spreadsheet",
        # "mimeType": "application/vnd.google-apps.folder", フォルダの場合
    }
    # supportsAllDrives=Trueがないと共有ドライブに作成できない
    res = drive.files().create(body=file_metadata, supportsAllDrives=True).execute()

    return res

フォルダIDとシートIDはXComで受け取ると便利

XComは異なるタスク間で値をやりとりする手段です。
フォルダ作成→スプレッドシート作成→データの書きこみの依存関係はこのように作成していて、フォルダ、スプレッドシートの作成はそれぞれPythonOperator内でおこなっているのでPythonOperatorも戻り値としてレスポンスのid(フォルダやスプレッドシートのID)を返すことで別のタスクからXComを使用してその値を受け取れます。

def create_folder(**kwargs):
    """フォルダを作成する

    Returns:
        [str]: 作成されたフォルダのID
    """
    folder_name = '<yyyy-mm-dd>'
    folder_id = '<FOLDER_ID>'

    res = create_drive_folder(folder_name, folder_id)
    return res["id"]

def create_spreadsheet(**context):
    """スプレッドシートを作成する

    Returns:
        [str]: 作成したスプレッドシートのID
    """
    # 前taskのcreate_folderから作成したフォルダのIDを受け取る
    folder_id = context["task_instance"].xcom_pull(task_ids="create_folder")
    sheet_name = "sheet_A"

    res = create_spread_sheet(sheet_name, folder_id)
    return res["id"]

with DAG(
    DAG_NAME,
    description="spreadsheet_test",
    schedule_interval="0 23 * * *",  # 毎日8:00(JST)に実行
    catchup=False,
    default_args=default_args,
) as dag:

    task_create_folder = PythonOperator(
        task_id="create_folder",
        python_callable=create_folder,
        provide_context=True,
        dag=dag,
    )
   task_create_spreadsheet_A = PythonOperator(
        task_id="create_spreadsheet_A",
        python_callable=create_spreadsheet,
        provide_context=True,
        dag=dag,
    )
   task_create_spreadsheet_B = PythonOperator(
        task_id="create_spreadsheet_B",
        python_callable=create_spreadsheet,
        provide_context=True,
        dag=dag,
    )
   task_create_spreadsheet_C = PythonOperator(
        task_id="create_spreadsheet_C",
        python_callable=create_spreadsheet,
        provide_context=True,
        dag=dag,
    )
    task_create_folder >> task_create_spreadsheet_A
    task_create_folder >> task_create_spreadsheet_B
    task_create_folder >> task_create_spreadsheet_C

etc task失敗時のslack通知

各taskが失敗した際にon_failure_callbackに関数を登録しておくとslackの通知するようにしています。
logのurlの取得ができるのでslackの通知から直接logに飛べるのでとても便利。
↓のコードではslackwebを使用して通知していますがairflowのproviderにもslackのhookが用意されているのでそっちを使ってもよかったと作ってから思いました。

参考にしたもの https://www.mikulskibartosz.name/send-cusomized-slack-notification-when-airflow-task-fails/
slack provider https://github.com/apache/airflow/tree/main/airflow/providers/slack

def failure_notification(context):
    task_instance = context["task_instance"]
    dag_name = context["dag"]
    task_name = task_instance.task_id
    log_link = f"<{task_instance.log_url}|{task_name}>"
    error_message = str(context.get("exception") or context.get("reason"))

    slack = slackweb.Slack(url=SLACK_URL)
    attachments = [
        {
            "username": "Cloud Composer(Airflow)",
            "icon_emoji": ":airflow:",
            "color": "danger",
            "title": "エラーが発生しました",
            "fields": [
                {"title": "Project", "value": f"{PROJECT_ID}"},
                {"title": "<DAG name> Task name", "value": f"{dag_name} {task_name}"},
                {"title": "Log URL", "value": log_link},
                {"title": "Error Message", "value": error_message},
            ],
        }
    ]
    slack.notify(attachments=attachments)

最後に

CloudComposer(airflow)のprovider周りのOperatorはあんまりドキュメントが整理されていないので直接airflowのソースコードを読んでどういった仕様になっているのか調べたほうがいいなと思いました。
今回はGoogle WorkspaceのAPI周りを調べましたがairflowのソースコードとGoogle WorkspaceのAPIのドキュメントを繰り返し見て作業してました。
providerはgoogleやslack以外にも多くあるので他のサービスと連携する際にproviderを見てみて連携するサービスがあれば一から実装しなくても良くなるかもしれないのでチェックしておくといいかもと思いました。PR送れるくらいGoogleのproviderのコードを読んだので共有ドライブ対応のPRを送ろうと思っています。

Kyash Advent Calendar 2021の他の記事も是非ご覧ください〜!