TiDB の SQL Parser を使用して CSV ファイルの不備を検知する仕組み

はじめに

こんにちは、ATOM 事業本部のエンジニアの岸田 (@mwudo) です。

集計基盤の機能開発や保守、API サーバー、バッチ処理などを担当しており、ATOM のバックエンド周りを見守っています。

趣味はボルダリングで、毎週、そびえ立つ壁に挑戦しています。

バックエンドの開発のテストで使用している CSV ファイルの不備に気づくために SQL Parser の使った取り組みについてご紹介します。

CSVの不備

バックエンドの開発でダミーデータとして DB 定義に則したダミーデータを CSV ファイルとしてリポジトリで用意してます。

このリポジトリをサブモジュールとして API やバッチ処理などのリポジトリで参照し、テスト時にあらかじめ投入してテストを行うことがあります。

CSV ファイルを DB へ投入するときのコマンドは LOAD DATA LOCAL INFILE を使用しています。

このコマンドは、投入先のテーブルの列数より CSV ファイルの列数が少なければ、値が入らない列にはデフォルト値が入るようになっています。

dev.mysql.com

これによって、バックエンドの開発において既存の DB のテーブル定義に対して列の追加・削除などの対応を行った時に気付きづらく、CSV ファイルを使用している各リポジトリのテストが失敗して気付くケースが少なからずありました。

このケースを踏んでしまって悲しい気持ちにならないために、DB 定義変更時に CSV ファイルの列とテーブル定義の列を比較する仕組みを作成しました。

SQL Parser の導入

SQL Parser を導入する前にテーブル定義の DDL を正規表現を使用して解析すること検討しました。

ChatGPT などを使えば正規表現を用意することは可能ですが、正規表現自体が複雑になりがちで、保守性が低いと考えました。

そこで、SQL を解析できる Go のライブラリを探したところ、github.com/pingcap/tidb/pkg/parser を見つけました。

github.com

こちらのライブラリは、PingCAP が開発している TiDB で使用されている SQL Parser です。

このライブラリを使用することで SQL を抽象構文木（以下 AST）に変換し、AST のノードを参照することで SQL の情報を取得することができます。

README によると、以下の特徴があるとのことです。

• MySQL のほとんどの構文をサポートしている
• 数行の実装を追加するだけで、独自の構文をサポートできる
• パフォーマンスがよい

Quickstart もあり、これを読むだけである程度の使い方がわかるようになっています。

実装しているテストを一部紹介

導入した SQL Parser を使用して、どのようにテストを実装しているか一部を紹介します。

AST のノードを取得するためには、以下の interface を実装する必要があります。

type Visitor interface {
    Enter(n Node) (node Node, skipChildren bool)
    Leave(n Node) (node Node, ok bool)
}

DDL の情報を取得するために構造体は以下のように定義しています。

// テーブル定義の情報を管理
type TableDef struct {
    Name    string      // テーブル名
    Columns []columnDef // カラム定義
}

// カラム定義の情報を管理
type ColumnDef struct {
    Field   string // カラム名
    Type    string // カラムの型
    NotNull bool   // NULL 許容かどうか
}

この構造体に以下のメソッドを実装して　Visitor として扱えられるようにしています。

func (t *TableDef) Enter(in ast.Node) (ast.Node, bool) {
    switch v := in.(type) {
    case *ast.CreateTableStmt:
        t.Name = v.Table.Name.String()
    case *ast.ColumnDef:
        t.Columns = append(t.Columns, columnDef{
            Field:   v.Name.String(),
            Type:    v.Tp.InfoSchemaStr(),
            NotNull: slices.ContainsFunc(v.Options, func(v *ast.ColumnOption) bool { return v.Tp == ast.ColumnOptionNotNull }),
        })
    }
    return in, false
}

func (t *TableDef) Leave(in ast.Node) (ast.Node, bool) { return in, true }

今回のテストでは DDL の SQL のためテーブル情報やカラム情報を持つ型に絞っていますが、SELECT や INSERT などの DML も取得することができます。

tidb/pkg/parser/ast at master · pingcap/tidb · GitHub に各型の実装があります。

最後に SQL を引数にして parser.New().ParseSQL を呼び出し Accept を実行すると、AST のノードが取得され def に情報が格納されます。

func extractDDL(t *testing.T, ddl string) TableDef {
    astNodes, _, err := parser.New().ParseSQL(ddl)
    assert.NoError(t, err) // github.com/stretchr/testify/assert
    def := TableDef{}
    astNodes[0].Accept(&def)
    return def
}

TableDef と ColumnDef で定義しているフィールドを使って行っているテストの概要はそれぞれ以下のようになっています。

フィールド	テスト概要
TableDef.Name	{Name}.csv が存在するかどうか
ColumnDef.Field	{Name}.csv のヘッダーの列に存在するかどうか
ColumnDef.Type	{Name}.csv のヘッダーの列のデータが型に合致しているかどうか
ColumnDef.NotNull	{Name}.csv のヘッダーの列に NULL が存在してもよいかどうか

まとめ

今回の仕組みを導入したことで今まで気づかず放置されていた CSV ファイルの不備を検知することができるようになりました。

人が手作業で確認をする必要がないものは仕組み化をすることで、手間を減らし、開発の生産性を向上させてプロダクト作りに集中できるようにしていきたいと思います！