Org-modeでSnowflake SQLを実行する

Org-modeのコードブロックからSQLクエリを実行できるとドキュメントの作成や、ちょっとしたデータ分析に便利です。

Org-babelはSQLクエリ実行に標準で対応していますが、標準でサポートしているSQLエンジンはMicrosoft SQL Server, MySQL, PostgreSQLなど限られています。 今回はこの一覧には存在しないSnowflakeのSQLクエリをコードブロックから実行し、実行結果をorg-mode内に表示させます。

Org-modeはDBI (dbish) などの汎用インターフェイスにも対応しているので、それらを経由してもよいのですが、Snowflakeは公式からsnowsqlというCLIツールが出ているので、 今回はこれを利用して情報の受け渡しを行うようにします。なお、認証情報は ~/.snowsql/config 等で設定されていることを前提とします。

設定

leaf.elの形式で書くと以下のようになります。

SQLクエリ実行時に大量に結果を出力しすぎてバッファを溢れさせることを防ぐために、 1000行より多く出力するようなときは確認ダイアログをミニバッファに出すようにしています。

(leaf ob-snowflake
  :if (executable-find "snowsql")
  :custom
  ;; CSV -> Tableに変換する行の上限をあげておく （デフォルト999）
  (org-table-convert-region-max-lines . 5000)
  :config
  ;; snowsql の実行結果が大きい場合に確認ダイアログを出すための閾値
  (defcustom ob-snowflake-max-rows 1000
    "Rows above which Org-Babel asks before inserting the result."
    :group 'org
    :type  'integer)

  ;; org-babel-execute:sql を Snowflake 対応でラップ
  (defun my/org-babel-execute:sql--snowflake (orig-fn body params)
    (let ((engine (cdr (assq :engine params))))
      (if (equal engine "snowflake")
          (let* ((cmdline       (cdr (assq :cmdline params)))
                 (result-params (cdr (assq :result-params params)))
                 (colnames-p    (not (equal (cdr (assq :colnames params)) "no")))
                 (out-file      (or (cdr (assq :out-file params))
                                    (org-babel-temp-file "sql-out-")))
                 (in-file       (org-babel-temp-file "sql-in-")))
            ;; クエリを一時ファイルへ
            (with-temp-file in-file
              (insert (org-babel-expand-body:sql body params)))
            ;; snowsql 実行（バナー抑止＋CSV＋ヘッダー可）
            (let ((cmd (format
                        "snowsql -f %s -o friendly=false -o output_format=csv -o header=%s %s > %s"
                        (org-babel-process-file-name in-file)
                        (if colnames-p "true" "false")
                        (or cmdline "")
                        (org-babel-process-file-name out-file))))
              (org-babel-eval cmd ""))
            ;; 行数を数えて 1000 行超なら確認ダイアログ
            (let* ((row-count (with-temp-buffer
                                (insert-file-contents out-file)
                                (count-lines (point-min) (point-max))))
                   (skip      (and (> row-count ob-snowflake-max-rows)
                                   (not (yes-or-no-p
                                         (format "Query returned %d rows. Insert into Org buffer? "
                                                 row-count))))))
              (if skip
                  "[Snowflake query skipped by user]"
                ;; 結果をテーブル／文字列として返す
                (if (member "table" result-params)
                    (let* ((table (with-temp-buffer
                                    (org-table-import out-file ",")
                                    (org-table-to-lisp))))
                      (if colnames-p
                          (append (list (car table) 'hline) (cdr table))
                        table))
                  (with-temp-buffer
                    (insert-file-contents out-file)
                    (buffer-string))))))
        ;; Snowflake 以外はデフォルト実装へ
        (funcall orig-fn body params))))
  ;; advice を登録
  (advice-add 'org-babel-execute:sql :around
              #'my/org-babel-execute:sql--snowflake))

上記を init.el で設定した後は、下記のようにヘッダに :engine snowflake としてSQLエンジンを指定してあげることで、Snowflakeでのコードブロックが実行可能になります。

#+begin_src sql :engine snowflake :results table
select *
from (values
        (1, 2, 3),
        (4, 5, 6),
        (7, 8, 9)
     ) as t(col1, col2, col3);
#+end_src

また 同ヘッダーに :results table を指定すればテーブル形式で表示され、 :results output を指定すればCSV形式で出力されます。テーブル形式での出力時には最大行数が org-table-convert-region-max-lines で制約されるので、表示可能な行数を適宜設定しておきましょう。