トップ > 技術的な取り組み > Snowflakeで複数のSQLをまとめて実行するストアドプロシージャを作った話


Snowflakeで複数のSQLをまとめて実行するストアドプロシージャを作った話

技術的な取り組み

Loading...

KADOKAWA Connected / ドワンゴの @saka1 です。 最近は、ドワンゴ・KADOKAWA向けに提供しているデータプラットフォームの移行作業を行っています。

engineering.kdx.co.jp

記事執筆時点（2022年3月初旬）では、移行作業はそろそろ佳境あるいは完了というところで、なんとかかんとか年度末には間に合うんじゃないかという進行です。

移行作業の進展に伴って、事前検証では見つけられなかった細かな問題や運用上の不便さなどが見つかるようになってきました。 新しく発生する課題に対して、データプラットフォーム運用チームは色々な方法で問題の解決・軽減をしようとしています。今回紹介するストアドプロシージャもそんな工夫の1つです。

直面した課題

Snowflakeのタスクの制限

Snowflakeでは、定期的に実行したい処理を書く場合、タスクというオブジェクトを定義することになります。

create task example_task
    schedule = 'USING CRON 0 15 * * * Asia/Tokyo' -- 毎日15:00にスケジュール
    warehouse = default_wh
    -- as以降には定期実行させたいSQLを書く
    as
        select 1;

タスクの制限として、SQL文をひとつしか登録できません。バッチ処理では複数のSQLを連ねて一つの処理とすることも多いので、複数SQLをスケジュールできないのは困ります。

1つの逃げ方としてタスク同士の依存関係を定義することはできるので、複数のタスクを連ねれば複数SQLの実行はできます。

-- このタスクはexample_taskが完了後に発火する
create task after_example_task
    warehouse = default_wh
    after example_task  -- ここで先行タスクとの依存関係を定義する
    as
        select 2;

じゃあこれで万事解決なのかというと、これはこれで扱いが難しいところがあります。

  • タスクが複数個必要になるので、記述としてはかなり冗長になってしまう。
  • 異なるタスクはセッションが分離されるので、例えば一時データ用の仮テーブル(temporary table)の共有や、途中のSQLが失敗した時に全体のロールバック、みたいなことができない。

要するに、あまり使い勝手がよくありません。「複数SQLを実行したいのであって複数タスクを作りたいわけじゃないのに」ってことですね。

解決策はストアドプロシージャだけど

Snowflakeでは、ストアドプロシージャを定義することで、発行するSQLを動的に切り替えたりと、非常に柔軟な処理を実現できます。「最悪でもストアドプロシージャあるから書けないことはないだろう」みたいな面があるので重宝します。

ストアドプロシージャの呼び出し自体はcall文ひとつなので「複数SQLを呼び出すストアドプロシージャを定義し、それをcallするタスクを作る」とすれば、1つのタスクで複数SQLの実行ができます。

参考リンク: ストアドプロシージャの利点

前述の単純な例をストアドプロシージャに書き換えてみます。

create or replace procedure procedure_one_two()
    returns string
    language javascript
    strict
    execute as caller
    as
    $$
        snowflake.execute({sqlText: 'select 1'});
        snowflake.execute({sqlText: 'select 2'});
        return "success";
    $$
;

create task example_task
    schedule = 'USING CRON 0 15 * * * Asia/Tokyo'
    warehouse = default_wh
    as
        call procedure_one_two();

こんな感じで、ストアドプロシージャではSQLの中に文字列リテラル$$ string $$で表されたJavaScript(JS)ソースコードを埋め込みます。JSの中では snowflake などの独自オブジェクトが触れるので、これらに対して操作することで、SQLの実行や結果取得などができます。

というわけで、確かに技術的にはできはするのですが……まあその、書いてみるとわかるのですが、お世辞にも使いやすくはないです。

  • SQLが書けるアナリストであっても、JSを書けるとは限らないです。
    • 典型的なアナリストのスキルとSnowflakeの機能がマッチしていない状況が見られました。
  • ストアドプロシージャは記述難易度が高いです。
    • SQL上にJSのソースコードを埋め込むので、例えばエディタの支援がかなり受けづらいです。
    • SQLの文字列として評価された後にJSとして実行されるので、エスケープ周りの書き方など、細かいハマりポイントがあります。

私たちのチームやアナリストの人達で使ってみたところ、案の定というか、利用にハードルがあることがわかりました。数百人の利用者にJSを覚えてくれと言うのも現実的ではなく、何かしらの工夫が必要そうでした。

課題のまとめ

  • 例えばタスクにおいて、複数のSQLをまとめて実行したいというニーズはあります。
  • 複数タスクの定義やストアドプロシージャの利用で機能的な実現はできます。
  • しかし複数タスクをこの用途に用いるのは使いづらさがあり、かといってストアドプロシージャは利用難易度が高いようでした。

緩和策

この課題をなんとかしようとチームでアイデア出ししたところ「複数SQL文字列を順次実行する高階関数みたいなものがあればいいのでは？」という発案が出てきました。こんな風に書けたらいいよねってことです。

run_queries([
    select 1,
    select 2,
])

もちろんSQLにこんな構文はないので、ある程度妥協することになります。SQLは文字列で受け取るしかありません。 幸いにもSnowflakeには半構造化データのサポートが入っているので、配列をネイティブで扱うことはできます。

run_queries(array_construct(
    'select 1',
    'select 2'
));

このような書き方を実現する run_queries は、それほど複雑ではありません。

create or replace procedure run_queries(query_array array)
    returns string
    language javascript
    strict
    execute as caller
    as
    $$
    try {
        for (const query of QUERY_ARRAY){
            snowflake.execute ({
                sqlText: query,
            });
        }
        return "Success";
    } catch (err) {
        throw err;
    }
$$

これで最低限の用途はカバーできるんじゃないかということで、データプラットフォーム利用者にこのrun_queriesを提供することにしました。

得られた効果

アナリストの方などにrun_queriesを使ってもらったところ、評判はまずまずでした。

  • 「JavaScriptを書かなくても、それなりに簡潔に複数SQLの実行ができる」のは期待通り便利なようです。
  • 順次実行するSQLを文字列の形で書くしかないため、利便性に限界があるのは否めません。

筆者の視点では、単純なアイディアでそれなりに便利になったのでコストパフォーマンスは良かったんじゃないか、と感じています。

将来展望

Snowflakeスクリプトで状況は改善される……はず？

Snowflake社も私たちのようなユースケースをなんとかしたいと思っているのかもしれませんが、記事執筆時点でSnowflakeスクリプトという機能がプレビュー（GAではないが利用者が触ってみることができる状態）になっています。

SnowfalkeスクリプトはストアドプロシージャをSQL風（SnowflakeのSQL拡張記法）に書けるというものです。例えばbegin〜endによるブロック構文がサポートされています。今回まさに欲しかったものです。

create or replace procedure procedure_one_two()
  returns string
  language sql
  as
  $$
    begin
      select 1;
      select 2;
      return 'success';
    end;
  $$;

create task example_task
    schedule = 'USING CRON 0 15 * * * Asia/Tokyo'
    warehouse = default_wh
    as
        call procedure_one_two();

さらにexecute immediate文を使うと、ブロックを直接埋め込めるようです。

create task example_task
    schedule = 'USING CRON 0 15 * * * Asia/Tokyo'
    warehouse = default_wh
    as
        execute immediate $$
            begin
              select 1;
              select 2;
              return 'success';
            end;
        $$;

JSが書けないとストアドプロシージャが書けない問題は、これで改善されるのかもしれません。

ワークフローエンジンは別で整備するのがよい？

もともとSnowflakeのタスク機能は素朴です。全てがSnowflake上で完結する点や、自分たちでインフラを管理しなくても良い点などは優れているのですが、複雑なパイプラインを組むにはそれほど向いていないように感じます。

実際、私達データエンジニアの部署や連携先システムでは、タスクを使わずにワークフローエンジン（例えばApache Airflow）を使ってSQLの実行を管理している事例も多いです。

データプラットフォームの主要な利用者にはアナリストの方々もいます。アナリストにAirflowを使ってくださいとするのは、少なくとも私たちの状況ではあまり現実味がありません。明確な答えはまだなく、今後要検討だなあと考えている次第です。

まとめ

Snowflake上の操作はおおむね全てSQLで書くのですが、どうしても基本的な表現力の限界みたいな所にぶつかることがあります。 そういう場合にストアドプロシージャはJSが使えるので表現力を相当補えるのですが、使用にJSスキルが必要になってきます。

対策の1つとして、run_queriesというストアドプロシージャを定義する方法を試行しています。 これが満点の解決策なのかは怪しいところだと思いつつ、1つの事例として紹介しました。

©KADOKAWA Connected Inc.