Sansan Tech Blog

Sansanのものづくりを支えるメンバーの技術やデザイン、プロダクトマネジメントの情報を発信

Dataformの単体テストにおけるチームの取り組み

技術本部 データ戦略部の松本と申します。

最近は気温の変化が激しく、体調を崩しやすい日々が続いております。皆さんもお元気でお過ごしでしょうか。

そんな中ですが、今日はGoogle Cloud Platformで提供されている マネージドサービスのDataformについての紹介と、チームにおいての単体テストの取り組みを紹介させてください。

なお、本記事は Sansan Advent Calendar 2023 3日目の記事です。 adventar.org

この記事の対象者は次のような方です

  • Dataformを少しでも触ったことがある人
  • Dataformの単体テストに困っている人

Dataformとは?

Dataformは、BigQueryなどのデータ分析基盤上で複雑なSQLを実行する際に利用するプロジェクト管理プラットフォームです。 次のような機能が特徴で、大規模なデータの加工処理を細かいパイプラインに分割して管理することが可能になっています。

  • Git連携
  • テーブルの依存関係をコード上で表現できる
  • 依存関係のツリーの可視化
  • JavaScriptを用いたSQLの生成及びコードの再利用
  • 単体テストのサポート

Dataformの採用理由

データ戦略部で構築する企業データベースのデータ統合部分で利用しています。 データ戦略部では様々な経路を用いてデータを収集していますが、チームでは収集結果に対してELT処理を行い、統一された企業データベースに変換するロジックを開発しています。

多数のデータをBigQuery上で効率よく処理できること、Google Cloud Platformのマネージドサービスであり導入が容易であること、SQLベースであり習得が容易であることを鑑み、今回 Dataform を利用して企業データベースを作成するロジックを構築しました。

Dataformの単体テスト

Dataformの単体テストの仕組み

最初に、Dataformの単体テストの仕組みを紹介します。Dataformでは次に示す手順で単体テストを行い、SQLXが仕様通り実装されているかを確認しています。

例として次のSQLXのファイルおよびテストファイルを想定します

config {
  type: "table",
}

SELECT
  MAX(input) AS output
FROM
  ${ref("input_table")}
config {
  type: "test",
  dataset: "xxx"
}

-- expectエリア
SELECT
  3 AS output

input "input_table" {
  -- inputエリア
  SELECT
    1 AS input
  UNION ALL
  SELECT
    2 AS input
  UNION ALL
  SELECT
    3 AS input
}

① テスト対象のテーブルを生成するSQLXファイルのうち、ref句を除く部分をコンパイルします

refの部分をinputエリアで指定したSELECT文の内容で置き換えます(サブクエリとして利用)

手順 2 を経た後の例示のクエリのref部分がサブクエリとして展開され、次のようになります

SELECT
  MAX(input) AS output
FROM
  (SELECT
    1 AS input
  UNION ALL
  SELECT
    2 AS input
  UNION ALL
  SELECT
    3 AS input)

③ 2 で生成したクエリをBigQuery上で実行し、結果を取得します

expectエリアに指定したクエリをBigQuery上で実行し、結果を取得します

⑤ 3と4とを出力された順番で1行ずつ比較し、結果が異なっていれば不一致とみなします

このような流れでDataformは単体テストを行なっています。refinputで差し替え、実際のデータと期待されるデータを比較することで、クエリや変換の正確性を確認できます。

チームにおける取り組み

単体テストはJavaScriptで記述する

現在、チームでは、次のような理由で単体テストを全てJavaScriptで記述する方針を採用しています。

複数のテストを単一のファイルで管理

SQLXでは単一のファイル内に単一のテストしか書けず、複数パターンのテストを行う場合は単一のテストに複数パターンを埋め込むか、複数のファイルを用意する必要があります。

JavaScriptを使用することで、複数のテストを単一のファイルに記述できます。これにより、関連するテストケースをまとめて管理でき、全体の可読性が向上します。

(test の命名規則をデータセット名 + テストケース名にするという運用を行なっています)

テストの効率的な実施

同じテーブルに対するテストケースは似ていることが多いですが、ファイルが分割されていると再利用することが難しいです。これを回避し、テストの共通パターンをJavaScriptでまとめて記述することで、テストの効率化を図っています。

SQL内部のシンタックスハイライトが可能

JavaScriptのファイル中にSQL文を記載する必要があるため、その観点では可読性が低下します。 一部エディタでは、JavaScript内部でSQLを記述する際も、シンタックスハイライトが適用されます。この機能を用いて、JavaScriptファイル内部のSQL文に対しての可読性を担保しています。

test(
  "データセット名_テストケース名"
)
  .dataset("dataset")
  .input(
    "input1",
    /*sql*/`SELECT 1 as input_1, "a" as input_2, null as input_3`
  )
  .expect(`SELECT 2 as output`)

JavaScript内部でSQLのハイライトを行う

このようなアプローチにより、テストコードの保守性を高めつつ、開発プロセスをスムーズに進めることができます。

SELECT文を作成する関数を用意して、SELECT文の記載を効率化する

先述の例では入力が完全に同じ場合には定数に切り出すことで共通化できますが、入力がわずかでも異なる場合はSELECT文を再度記述する必要がありました。この課題に対し、toSelectQueryという簡単な変換関数を作成し、この関数にJavaScriptの配列を渡すことでSELECT文に変換するユーティリティを開発しました。

const serializeSqlValue = (value) => {
  // nullの場合はnullに変換
  if (value === null) {
    return "NULL";
  }

  // nullValueヘルパー関数で指定された場合はココで強制的に変換する
  if (value.__type) {
    return value.toJSON();
  }

  // 日付の場合はTIMESTAMP型に変換
  if (value instanceof Date) {
    return `TIMESTAMP '${value.toISOString()}'`;
  }

  // objectの場合はJSON型に変換
  if (typeof value === "object") {
    return `JSON '${JSON.stringify(value)}'`;
  }

  // それ以外はJSON文字列に変換
  return JSON.stringify(value);
};

const toSelectQuery = (objs) => {
  return objs
    .map((obj) => {
      return `SELECT ${Object.entries(obj)
        .map(([key, value]) => `${serializeSqlValue(value)} as \`${key}\``)
        .join(",")}`;
    })
    .join(" UNION ALL ");
};

// JSON形式でobject以外のデータを表現する場合に使用する
const json = (text) => {
  return {
    toJSON: () => text,
  };
};

// nullに形を与える場合に使用する
const nullValue = (type) => {
  return {
    __type: type,
    toJSON: () => `CAST(NULL AS ${type})`,
  };
};

module.exports = {
  serializeSqlValue,
  toSelectQuery,
  json,
  nullValue,
};

上述のテストを、toSelectQueryを用いて書き直すと、次のようになります。

test(
  "データセット名_テストケース名"
)
  .dataset("dataset")
  .input(
    "input1",
    toSelectQuery([{
      input_1: 1,
      input_2: "a",
      input_3: null
    }])
  )
  .expect(toSelectQuery([{output: 2}]))

これによりテストデータを単純なJavaScriptオブジェクトとして保持することが可能になり、定数として再利用することや、map関数などの配列操作関数を用いて一部データを書き換えるなど、類似するテストケースを効果的に管理できるようになりました。

UDF関数をoperationとして一つのファイルに切り出し、 ref を使用せずに利用する

チームではUDF自体の管理もDataformで一元管理する運用を取っています。しかし、現時点のDataformではUDFの扱いが不十分でした。

公式のGitHub上のIssue にある通り、UDFをDataformで管理するためには type: operations でUDF文を実装すること、およびそれを ref で参照することで実現されます。

(下記クエリはGitHubのIssueより引用)

// my_udf.sqlx
config {
  type: "operations",
  hasOutput: true,
}

create function ${self()}(x INT64) ...
// my_table.sqlx
config { type: "table" }
select ${ref("my_udf")}(column) from ...

しかし、関数の利用箇所にrefを使用すると、先述のテストの仕組みに引っかかり、うまくテストできません。(具体的には、ref部分をinputの内容でサブクエリ置換するため、CREATE FUNCTION文がサブクエリとして現れてしまうという問題が発生します。)

上記問題を避けて依存関係を解決するためにチームでは、UDF関数を利用箇所はテスタビリティのために、 dependencyに依存関係記述しかつrefを使用せずに関数を利用するという運用にしています。

-- UDF関数本体 (test.sqlx)
config {
  type: "operations",
  hasOutput: true,
  tags: ["ci_required"],
}

CREATE OR REPLACE FUNCTION test
-- 関数の利用箇所 (table.sqlx)

config {
  type: "table",
  hasOutput: true,
  dependencies: [
    "test" -- 依存関係を定義して実行時に関数を適切に更新する
  ]
}

SELECT `defaultDataset.test`(input_1) from ${ref("some_table")} -- 関数名はrefで解決させない

このアプローチにより、UDF関数の一元管理と依存関係の解決を効果的に行い、テストの容易性を確保しました。 *1

UDFのテストはGoogleが提供する関数を利用する

UDF関数をDataformで管理しているため、UDFに対してもテストを実施しています。次のブログにある関数を利用してUDFのテストを行いました。

cloud.google.com

ランダムなIDが各実行ごとに付与されるため、複数のテストに分割するとデバッグ性がとても下がってしまうのが難点です。 現状我々のチームでは、UDF関数のテストは単一ファイル内で配列で複数テストケースを表記した上で、コメントでテストケース名を補完するという運用を行なっています。

const { generate_udf_test } = unit_test_utils;

generate_udf_test("dataset.function_name", [
  // テスト補足コメントをここに書く
  { inputs: [`""`], expected_output: `""` },

  // テスト補足コメントをここに書く
  { inputs: [`NULL`], expected_output: `NULL` },
]);

UDF関数には ci_required のタグを付与

上2つで紹介した、UDF関数のテストおよびUDF関数を利用しているクエリのテストを行うためには、UDF関数が事前にテストを実行するプロジェクト上に作成されている必要があります。

そのためチームでは、UDF関数には ci_required というタグを付与し、以下のコマンドを使用してテスト前に関数の再生成を行なった上で単体テストを実行する運用を構築しました。

config {
  type: "operations",
  hasOutput: true,
  tags: ["ci_required"],
}

CREATE OR REPLACE FUNCTION ...
$ dataform run --tags ci_required
$ dataform test

上記対応はCI上で行なっていますが、これによりCIの実行時に必ず最新の関数が存在していることが保証され、UDF関数を効率的にテストできるようになりました。

最後に

今回の記事が皆様のご参考になれば幸いです。Dataformはまだまだ発展途上なので、コントリビュートもしていきたいなと思いました。 今年のアドベントカレンダーでは類似ツールである dbt の紹介も予定されています。ぜひ楽しみにしていてください。

*1:CREATE TEMP FUNCTIONを利用する案もチームでは出ましたが、単体テストが(執筆当時) preOperation句を無視するという仕様であるため、CREATE TEMP FUNCTIONを用いたファイルはテスト実行不可能になります。

© Sansan, Inc.