Airflowの単体テストを書きましょう

データ基盤は下流の分析・可視化・モデリングの「基盤」となるので、品質の担保は言うまでもなく重要ですね。品質を確保するには、ワークフローの監視・検証、ワークフローのテスト、そして加工用クエリのテストがいずれも欠かせません。この記事では、ワークフロー(Airflow)の単体テスト方法について紹介します。また、ワークフローの監視・検証に関しては、過去の記事も合わせてご覧いただけると幸いです。 ワークフローの監視 ワークフローの検証 DAGの単体テスト まずは、DAGの単体テストについて説明します。厳密に言えば、DAGの実行ではなく、DAGが正確に構築されたかどうかのテストを行います。 https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#unit-tests Airflowの公式ベストプラクティスでは簡潔に紹介されていますが、具体例を挙げてさらに詳しく説明しましょう。 importのテスト importが正常にできることを確認する(importが失敗するとWeb UIからも確認できるが、単体テストする時点で確認するともっと便利) import時間を制限する。(冗長なDAGがあると解析するのに時間がかかるので、import時間を制限することで、事前に冗長なDAGを発見できる) 最低でも1つのタスクが含まれていることを確認する。 import unittest from datetime import timedelta from airflow.models import DagBag class TestImportDags(unittest.TestCase): IMPORT_TIMEOUT = 120 @classmethod def setUpClass(cls) -> None: cls.dagbag = DagBag() cls.stats = cls.dagbag.dagbag_stats def test_import_dags(self): self.assertFalse( len(self.dagbag.import_errors), f"DAG import failures. Errors: {self.dagbag.import_errors}", ) def test_import_dags_time(self): duration = sum((o.duration for o in self.stats), timedelta()).total_seconds() self.assertLess(duration, self.IMPORT_TIMEOUT) def test_dags_have_at_least_one_task(self): for key, dag in self.dagbag.dags.items(): self.assertTrue(dag, f"DAG {key} not exsit") self....

April 6, 2023 · Me

AirflowからAirbyteをトリッガーする際にハマるポイント

https://docs.airbyte.com/operator-guides/using-the-airflow-airbyte-operator/ AirflowからAirbyte Operatorを利用するための設定について、Airbyte公式の記事は既にわかりやすくまとめています。実際に試してみて、少しハマったところがあったので、その知見を共有したいと思います。 1. Airflowを2.3.0以上にアップグレードする必要がある apache-airflow-providers-airbyte[http]を利用するのにAirflowを2.3.0以上に上げないといけません。(apache-airflow-providers-airbyte[http]をdocker-composer.ymlの_PIP_ADDITIONAL_REQUIREMENTSに追加することも忘れずに) Cloud Composerなどを利用している場合、GUIからアップグレード可能です。 https://airflow.apache.org/docs/apache-airflow-providers-airbyte/stable/index.html version: '3' x-airflow-common: &airflow-common image: apache/airflow:2.3.4-python3.8 environment: &airflow-common-env PYTHONPATH: /opt/airflow/dags AIRFLOW__CORE__EXECUTOR: CeleryExecutor AIRFLOW__CORE__SQL_ALCHEMY_CONN: postgresql+psycopg2://airflow:password@postgres/airflow AIRFLOW__CELERY__RESULT_BACKEND: db+postgresql://airflow:password@postgres/airflow AIRFLOW__CELERY__BROKER_URL: redis://:@redis:6379/0 AIRFLOW__API__AUTH_BACKEND: 'airflow.api.auth.backend.basic_auth' # 追加 _PIP_ADDITIONAL_REQUIREMENTS: apache-airflow-providers-airbyte[http]==3.2.0 2. Airflowの古いバージョンから2.3.4上げるとdocker-composeがバグる airflow 2.2.xでは問題なく環境構築できていましたが、イメージをapache/airflow:2.3.4-python3.8に変更してdocker compose up airflow-initを実行したら怒られます。 You are running pip as root. Please use 'airflow' user to run pip! Airflowの古いdocker-composer.ymlのバグのようなので、 https://github.com/apache/airflow/pull/23517/files services -> airflow-init -> environmentに_PIP_ADDITIONAL_REQUIREMENTS: ''を追加すれば解決できます。 ... environment: <<: *airflow-common-env _AIRFLOW_DB_UPGRADE: 'true' _AIRFLOW_WWW_USER_CREATE: 'true' _AIRFLOW_WWW_USER_USERNAME: ${_AIRFLOW_WWW_USER_USERNAME:-airflow} _AIRFLOW_WWW_USER_PASSWORD: ${_AIRFLOW_WWW_USER_PASSWORD:-password} # 追加 _PIP_ADDITIONAL_REQUIREMENTS: '' ....

March 6, 2023 · Me

Cloud SQLにある大量なテーブルをBigQueryに入れる話

経緯 こんにちは、データエンジニアのjcです。 昨年度から大規模なデータ分析基盤の構築に携わっています。最近Cloud SQLにある6つのDBの数百個のテーブルを日次洗い替えでBigQueryにあるデータ基盤に入れるタスクを取り組んでいます。 Cloud SQLとBigQuery両方ともGCPのサービスのため、federated queriesを利用すると簡単にできそうに見えますが、 https://cloud.google.com/bigquery/docs/federated-queries-intro 実際に行ってみると、以下の3つの課題を気づきました。 BigQuery側でスキーマ情報を含めたテーブルを一々作成するのは現実的ではない プロダクトの進化とともにテーブル・カラムが頻繁に作成・変更されるため、BigQuery側でも対応しないといけない Cloud SQLにあるテーブルの定義をそのまま取ってきてもBigQueryではMySQLとPostgreSQLの一部の型が対応されていない https://cloud.google.com/bigquery/docs/tables#sql_1 少し苦労していましたが、幸い解決方法を見つけました。 今後躓く方もいるかもしれないので、知見を共有したいと思います。 BigQuery側でスキーマ情報を含めたテーブルを一々作成するのは現実的ではない BigQueryはクエリの結果によってテーブルを作成できるので、事前にテーブルを作っておく必要がありません。 https://cloud.google.com/bigquery/docs/tables#sql_1 大量なテーブルを一々作成するのは現実的ではない課題の解決法としてはDBのメタ情報(descriptionを含めて)をそのまま生かしてテーブル作成用クエリを生成し、テーブルを作成します。 例えばPostgreSQLの場合、まずはテーブルのメタ情報 SELECT schemaname, relname AS table_name, obj_description(relid) AS description FROM pg_catalog.pg_statio_all_tables WHERE schemaname = '{YOUR_SCHEMA}' とカラムのメタ情報を取得します。 SELECT c.table_schema, c.table_name, c.column_name, c.data_type, pgd.description FROM pg_catalog.pg_statio_all_tables AS st INNER JOIN pg_catalog.pg_description pgd on ( pgd.objoid = st.relid ) RIGHT JOIN information_schema.columns c ON ( pgd.objsubid = c.ordinal_position and c.table_schema = st.schemaname and c....

February 2, 2023 · Me

dbtでBigQuery上に構築したデータ基盤のメタデータ管理

データ基盤におけるETL/ELT開発のT(Transform)を担うツールdbtは最近注目を浴びています。dbtでデータモデリングする方法既に多く紹介されたので、この記事では手を動かしながらdbtでBigQuery上に構築したデータ基盤のメタデータを管理する方法を紹介します。 環境構築 dbt公式はHomebrewを推していますが、ローカル環境が汚染されるのをなるべく避けたいので、Dockerで環境構築します。 dbtのプロジェクトとプロファイルの設定ファイルを用意しておかないと、公式のドキュメントそのまま実行したらコケます。しかし設定ファイルの生成は環境を構築する必要があるので無限ループになっています。 https://docs.getdbt.com/docs/get-started/docker-install そのため、公式のサンプルプロジェクトをforkし、事前にローカル環境で生成した設定ファイルを追加しました。 https://github.com/aibazhang/dbt-metadata-management profiles.ymlを編集 {YOUR_DATASET_NAME}と{YOUR_PROJECT_ID}を置き換えます。 複数のデータセットのメタデータも作成可能ですが、一旦任意のデータセット名を指定する必要があります。dbtの問題点でもありますが、後ほど説明します。 Dockerイメージをプル docker pull ghcr.io/dbt-labs/dbt-bigquery:1.2.0 コンテナを立ち上げる git clone https://github.com/aibazhang/dbt-metadata-management cd dbt-metadata-management gcloud認証 認証済みの場合、このステップは不要です。 gcloud auth login --no-launch-browser gcloud auth application-default login --no-launch-browser コンテナを立ち上げる docker run --rm \ --network=host \ --platform linux/amd64 \ --mount type=bind,source=`PWD`,target=/usr/app \ --mount type=bind,source=`PWD`/profiles.yml,target=/root/.dbt/profiles.yml \ --mount type=bind,source=$HOME/.config/gcloud/application_default_credentials.json,target=/root/.config/gcloud/application_default_credentials.json \ ghcr.io/dbt-labs/dbt-bigquery:1.2.0 \ ls データモデルのリストが表示されたら、環境構築が無事終了です。 ドキュメントを生成する docs generate 以下のコマンドを実行すれば、models/配下のクエリとメタデータ(yamlファイル)を参照して、target/配下にメタデータのドキュメントが生成されます。 docker run --rm \ --network=host \ --platform linux/amd64 \ --mount type=bind,source=`PWD`,target=/usr/app \ --mount type=bind,source=`PWD`/profiles....

December 11, 2022 · Me

アラートを出す際にAirflowのContextから誤ったtask idが取得されてしまうバグの対処法

先日投稿した記事はAirflow DAGのon_failure_callbackとdagrun_timeoutを組み合わせることでDAGの遅延を監視する方法を紹介しました。 Contextから誤ったtask idが取得されてしまう contextからdag_runの情報を取得してチャットツールやメールにアラートを出すのは一般的です。Slackにアラートを出す際の例ですが、dag_id, run_id, task_id, reason, log_urlを取得して、webhookでSlackの特定なチャンネルに投稿し、log_urlをクリックするだけですぐローカルあるいはクラウド環境(例えばCloud Composer)で失敗したtaskのログを確認できるので、アラート解消の効率化に繋がります。 ソースは以下となります。 from slack_sdk.webhook import WebhookClient from airflow.models import Variable from textwrap import dedent def notify_error(workflow: str, context: dict) -> None: webhook = WebhookClient(Variable.get("slack_webhook_access_token")) log_url = context.get("task_instance").log_url message = dedent( f""" :x: Task has failed. *Workflow*: {workflow} *DAG*: {context.get('task_instance').dag_id} *Run ID* {context.get('dag_run').run_id} *Task*: {context.get('task_instance').task_id} *Reason*: {context.get('reason')} <{log_url}| *Log URL*> """ ) webhook.send( text="alert", blocks=[ { "type": "section", "text": {"type": "mrkdwn", "text": message}, } ], ) しかし、数回検証してみた結果、実行が失敗したタスクtask_idではなく、誤ったtask_idが取得されてしまう事象がしばしば発生します。Airflowの既知バグで、現時点(2022....

November 21, 2022 · Me

Airflowのon_failure_callbackとdagrun_timeoutを組み合わせることでDAGの遅延を監視する

https://buildersbox.corp-sansan.com/entry/2022/08/18/110000 この記事では新しくDAGを作成してデータの転送処理が遅れているかを監視する方法が紹介されました。監視と実行用のDAGを分離することでスッキリにはなりますが、DAGが増えることによって管理の手間が生じます(特にDAGが大量にある場合)。 https://stackoverflow.com/questions/69312630/airflow-triggering-the-on-failure-callback-when-the-dagrun-timeout-is-exceed を参考にして DAGの引数on_failure_callbackとdagrun_timeoutを組み合わせることでDAGの遅延を監視する方法を試してみました。 from datetime import datetime, timedelta from airflow import DAG from airflow.models import TaskInstance from airflow.operators.bash import BashOperator def _on_dag_run_fail(context): print("***DAG failed!! do something***") print(f"The DAG failed because: {context['reason']}") print(context) default_args = { "owner": "mi_empresa" } with DAG( dag_id="failure_callback_example", start_date=datetime(2021, 9, 7), schedule_interval=None, default_args=default_args, catchup=False, on_failure_callback=_on_dag_run_fail, dagrun_timeout=timedelta(seconds=45), ) as dag: delayed = BashOperator( task_id="delayed", bash_command='echo "waiting..";sleep 60; echo "Done!!"', ) will_fail = BashOperator( task_id="will_fail", bash_command="exit 1" ) delayed >> will_fail 実行時間が45秒超えると、_on_dag_run_failが実行され失敗したDAGのコンテキスト情報がプリントされます。誰かの役に立つかもしれないので、試しながら気づいた2つのポイントを共有します...

November 17, 2022 · Me

Airflowで構築したワークフローを検証する

データ基盤の監視 データ基盤は下流の分析・可視化・モデリングの「基盤」となるので、監視は言うまでもなく品質を担保するため重要な存在です。データ基盤監視の考え方についてこの2つの記事が紹介しています。 https://tech-blog.monotaro.com/entry/2021/08/24/100000 https://buildersbox.corp-sansan.com/entry/2022/08/18/110000 同じくSQLによるデータ基盤を監視しており、最も大きな違いは自作ツールかAirflowで検証することだけです。本文はAirflowで構築したワークフローの検証についてもう少し紹介したいと思います。 まず、Data Pipelines Pocket Referenceではデータ基盤検証の原則が紹介されました。 Validate Early, Validate Often 要はできるだけ早く、できるだけ頻繁に検証するとのことです。ELTあるいはETL処理においては、Extract, Load, Transformそれぞれのステップが終了した直後に監視するのは最も理想的だと思います。 Transformはデータセットのコンテキストを把握しておかないと検証できないため、データセットごとに対応していく必要があります。ExtractとLoadはnon-contextで汎用的な検証ができるため、データ基盤構築の序盤からやっておいた方が安心だと思います。 Extractの検証 ストレージサービス(例えばGCS)をデータレイクにする場合、データソースからデータレイクにデータがちゃんとレプリケートされたかを検証するためにAirflowのairflow.providers.google.cloud.sensors.gcsを利用すると簡単にできます。 https://airflow.apache.org/docs/apache-airflow-providers-google/stable/_api/airflow/providers/google/cloud/sensors/gcs/index.html 一つのファイルをチェックする場合はほとんどないと思うので、GCSObjectExistenceSensorよりGCSObjectsWithPrefixExistenceSensorの方がもっと実用的でしょう。下記のタスクをExtractと次の処理の間に挟むと、障害の早期発見が期待できます。なお、Extract時点で既に障害が起きている場合はほとんどデータソース側の処理が失敗しているので、アプリケーション側と連携して作業する必要があります。 check_extract = GCSObjectsWithPrefixExistenceSensor( task_id="check_extract", bucket="{YOUR_BUCKET}", prefix="{YOUR_PREFIX}", ) Loadの検証 ELTとETL処理、Loadするタイミングは異なりますが、検証の方法(データがデータウェアハウスあるいはデータマートにロードされたか)は同じです。よく使われているデータウェアハウスサービスBigQueryだと、Airflowのairflow.contrib.operators.bigquery_check_operator.BigQueryCheckOperatorを利用できます。 https://airflow.apache.org/docs/apache-airflow/1.10.10/_api/airflow/contrib/operators/bigquery_check_operator/index.html#airflow.contrib.operators.bigquery_check_operator.BigQueryCheckOperator 下記のタスクをLoad処理の後ろに追加すれば良いです。 check_load = BigQueryCheckOperator( task_id="check_load", sql={YOUR_SQL}, use_legacy_sql=False, params={ "bq_suffix": Variable.get("bq_suffix"), "dataset": setting.bq.dataset, "table": setting.bq.table_name, }, location="asia-northeast1", ) SQLは検証用のクエリとなっており、BigQueryのメタデータテーブルがよく用いられます。例えばこの記事が紹介されたクエリでテーブルが空になっているかを確認できます。 SELECT total_rows FROM ${dataset_id}.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '${table_name}' ORDER BY last_modified_time DESC LIMIT 1; GCPについて簡単に説明しましたが、AWSとAzureも似たようなことはできるはずです。 皆さんのワークフロー設計にご参考になれば幸いです。

November 1, 2022 · Me

AirflowのBigQueryInsertJobOperatorのオプションについて

先日「AirflowでGCS外部パーティションテーブルをBigQueryに取り込む方法」についての記事を書きましたが、BigQueryInsertJobOperatorにフォーマットのオプションを追加する際に、実際の挙動は公式ドキュメントの記述と若干違うので、追記しておきます。 なぜ気づいたのか 外部テーブルを取り組む際に、閉じるカッコが見つからないというエラーが出て、 Error detected while parsing row starting at position: 4204. Error: Missing close double quote (") character. 調査してみた結果、改行記号が含まれているのが原因のようでした。stack overflowによるとallowQuotedNewlinesといったCSVのオプションを追加すると解決できるらしいです。 https://stackoverflow.com/questions/65782572/how-to-avoid-missing-close-double-quote-character-error-in-google-bigquer https://qiita.com/sawanoboly/items/410b3346518e569da581 ただし、configurationにcsvOptionsを追加しても解決されず(同じエラーが出た)、困っていました。 configuration={ "load": { "destinationTable": { "projectId": PROJECT_ID, "datasetId": DATASET_NAME, "tableId": TABLE, }, "sourceUris": ["gs://my-bucket/biostats/*.csv"], "autodetect": True, "sourceFormat": "CSV", "hivePartitioningOptions": { "mode": "AUTO", "sourceUriPrefix": "gs://my-bucket/biostats/", }, "csvOptions": { "allowQuotedNewlines": True }, } }, 解決法 CLI bq loadを利用す際にどうなっているかを確認したら、どうやらcsvOptionsを明示的に書く必要がなく--allow_quoted_newlinesだけ追加すとでうまくいきました。AirflowのBigQueryInsertJobOperatorで同じことをすると、問題なくloadできました。 configuration={ "load": { "destinationTable": { "projectId": PROJECT_ID, "datasetId": DATASET_NAME, "tableId": TABLE, }, "sourceUris": ["gs://my-bucket/biostats/*....

August 30, 2022 · Me