Streamlit in Snowflake でSQL作成の補助ツールを作ってみた

2024年7月30日 18:32

はじめに

こんにちは、ライフイズテックのデータサイエンティストのホンディーです。

少し前から各所で名前を聞くようになったStreamlitについて、若干出遅れてしまいましたが、僕もようやくその便利さと面白さに気づき色々試して見るようになりました。ローカル端末上でも動かしていますが、ライフイズテックのデータ基盤で採用しているSnowflakeでも、Streamlit in Snowflake という機能でStreamlitを動かせます。その Streamlit in Snowflake で日々のSQL作成作業に便利なツールを作ったのでそれを紹介します。

この記事のヘッダー画像は上記のプレスリリースから拝借しました。

SQLを作成するときにselectしたい列の列挙が面倒なので効率化したい

職種の特性上、ほとんど毎日何かしらのSQLを書いています。その中で面倒に感じているのが抽出したい列名の記述です。SnowflakeのワークシートやReDashでSQLを書く場合は補完も効きますが、それでも列数が多いと面倒です。そしてテキストエディタやjupyterでSQLを書いているときはその補完さえ無いので一層手間がかかります。

少しでも効率化しようと、describe文を使って列名の一覧を取得してきて、次のような列の一覧をPythonで生成してコピペして使うといったことをよくやっていました。(それぞれ、aはエイリアス、column{数字} は列名のサンプルです。)

a.column1,
a.column2,
a.column3,
・・・

この、「列の情報をDBから取得してきて、コピペ用にエイリアスをつけての列挙をStreamlitのGUIツールで手軽にできるようにした」というのが作成したツールの機能です。

作成したツールの動作イメージ

出来上がった画面を先にお見せします。セキュリティ上の理由でDB名などはマスクし列名はダミーに差し替えさせていただきます。

左側のプルダウンでDB/スキーマ/テーブルを選び、エイリアスを入力する(省略も可能)と、そのテーブルから全列をSELECTしてくるSQLを構築して、更に各列のデータ型情報を表示するものです。

テキストエディタやjupyterでSQLを書く時、一旦ここで作成したSQLの列の一覧を全部コピペして来てから不要な列を消す手順にすることでかなり作業が楽になりました。必要なデータの取得の手間が下がる分、集計ロジックの記述や、SQLで抽出した後の分析に頭のリソースを集中できるようになったのも良いですね。

実際のコード

実装したコードはこちらです。Streamlit in Snowflake で動かす前提の実装なので他の環境のStreamlitでは動作しないのでご注意ください。

import streamlit as st
import pandas as pd
from snowflake.snowpark.context import get_active_session


session = get_active_session()


# データベース名の一覧を小文字で取得する。
def get_databases():
    df = pd.DataFrame(session.sql("show databases").collect())
    return df["name"].str.lower()


# 指定したデータベースのスキーマ名の一覧を小文字で取得する。
def get_schemas(database):
    query = f"show schemas in database {database}"
    df = pd.DataFrame(session.sql(query).collect())
    df["name"] = df["name"].str.lower()
    return df[df.name != "information_schema"]["name"]


# 指定したスキーマのテーブル名の一覧を小文字で取得する。
def get_tables(database, schema):
    query = f"show tables in {database}.{schema}"
    df = pd.DataFrame(session.sql(query).collect())
    return df['name'].str.lower()


# 指定したテーブルの列名とデータ型の一覧を小文字で取得する。
def get_columns(database, schema, table):
    query = f"describe table {database}.{schema}.{table}"
    df = pd.DataFrame(session.sql(query).collect())
    df["name"] = df["name"].str.lower()
    df["type"] = df["type"].str.lower()
    return df[["name", "type"]]


# 初期化
if 'databases' not in st.session_state:
    st.session_state.databases = get_databases()
if 'schemas' not in st.session_state:
    st.session_state.schemas = {}
if 'tables' not in st.session_state:
    st.session_state.tables = {}
if 'columns' not in st.session_state:
    st.session_state.columns = {}

# データベースの選択
selected_database = st.sidebar.selectbox('データベース選択', st.session_state.databases)

if selected_database:
    # スキーマの取得とキャッシュ
    if selected_database not in st.session_state.schemas:
        st.session_state.schemas[selected_database] = get_schemas(selected_database)

    selected_schema = st.sidebar.selectbox('スキーマ選択', st.session_state.schemas[selected_database])

    if selected_schema:
        # テーブルの取得とキャッシュ
        if f"{selected_database}.{selected_schema}" not in st.session_state.tables:
            st.session_state.tables[f"{selected_database}.{selected_schema}"] = get_tables(selected_database, selected_schema)

        selected_table = st.sidebar.selectbox('テーブル選択', st.session_state.tables[f"{selected_database}.{selected_schema}"])

        if selected_table:
            # 列の取得とキャッシュ
            if f"{selected_database}.{selected_schema}.{selected_table}" not in st.session_state.columns:
                st.session_state.columns[f"{selected_database}.{selected_schema}.{selected_table}"] = get_columns(selected_database, selected_schema, selected_table)

            alias = st.sidebar.text_input("エイリアス")

            # SQLの構築
            sql = "select\n"
            for col in st.session_state.columns[f"{selected_database}.{selected_schema}.{selected_table}"]["name"]:
                if alias == "":
                    sql += f"    {col},\n"
                else:
                    sql += f"    {alias}.{col},\n"
            sql += "from\n"
            if alias == "":
                sql += f"    {selected_database}.{selected_schema}.{selected_table}"
            else:
                sql += f"    {selected_database}.{selected_schema}.{selected_table} as {alias}"

            st.write('## SQL')
            st.text(sql)
            st.write("## テーブル定義")
            st.table(st.session_state.columns[f"{selected_database}.{selected_schema}.{selected_table}"])

工夫したところ

工夫した点をいくつか紹介します。

1. キャッシュを残し、取得済のデータは使い回すようにした

Streamlitは、なにかウィジェットを操作すると原則すべてのコードが再実行されてしまうのですが、毎回DBの一覧を取得すると無駄なSQLが何度も発行されてしまいます。

そこで、 st.session_state というStreamlitのデータをセッション内で保持する機能を使い、そこにデータが有れば使い回すようにしました。

このあたりで、初期化しています。

if 'databases' not in st.session_state:
    st.session_state.databases = get_databases()
if 'schemas' not in st.session_state:
    st.session_state.schemas = {}
if 'tables' not in st.session_state:
    st.session_state.tables = {}
if 'columns' not in st.session_state:
    st.session_state.columns = {}

そして、このように、キャッシュにデータがない場合だけSQLを発行してデータを取っています。(下のコードはスキーマの一覧取得部分)

if selected_database:
    # スキーマの取得とキャッシュ
    if selected_database not in st.session_state.schemas:
        st.session_state.schemas[selected_database] = get_schemas(selected_database)

2. 操作に関する部分と結果の表示部分を分離

Streamlitは基本的にはコード内で作成作成されたコンポーネントを実行した順に上から下へ順番に並べていく配置になります。ただ、それだと列が多いテーブルの場合に画面が縦に長くなり操作性がいまいちだったので、 st.sidebar を利用してドロップダウンとエイリアスのテキストボックスを画面左に分離しました。

3. ドロップダウンの選択肢を動的に変更

{データベース}.{スキーマ}.{テーブル} という結合した選択肢を作って一つのドロップダウンで選択するのではなく、選択したデータベース、スキーマに応じて、そこに含まれるテーブルだけがテーブル選択のドロップダウン内に現れるような作りにしました。

ライフイズテックでは多くの事業を展開しているので必然的にDBの種類が多くあります。そしてその多くは教育に関する事業なので、生徒、教室、教科書、問題、といった内容のテーブルが多くのDBに存在し名前も似通っていたり同じだったりします。そのため、先にDBを指定してその中のテーブルだけ選択する作りにすると紛らわしさが減り便利になりました。

Snowflake in Streamlit の便利なところ

ローカル端末でつかうStreamlit と比較して、Snowflake in Streamlitで使う便利さが２つありました。

1. 起動しなくていい

自分の端末のStreamlit はいちいち streamlit run コマンドで起動しないと使えません。しかも別のを使いたくなったら起動中のを止めてもう一つを起動する必要があり少々不便です。

一方で、Snowflake in Streamlit であれば、Snowflakeにログインしてページにアクセスするだけで使えます。自分はSnowflake自体はだいたい常にログインしているのでこちらのほうが断然手軽です。

2. DBへの接続が簡単

これは実装が簡単になるという話なのですが、Snowflake にログインした状態で動かすため、追加の認証情報等がいりません。get_active_session() でセッションを取得したらそのままSQLが発行できます。

逆に、Snowflake上で動いているコードをそのままローカルに持ってくると動かないというデメリットにはなりえますが、今のところそのような場面はないので困っていません。

Snowflake in Streamlit のデメリット

良い点ばかりではなく、若干のデメリットも有りました。

1つはStreamlitのバージョンが少し古いことです。現在、Snowflakeで使えるStreamlitのバージョンは1.26 です。そのため、直近のStreamlitのバージョンで追加された機能は使えません。Streamlitの開発が非常に速いペースで進んでいるので、そのSnowflakeでは使えない機能の中に例えば、HTMLを表示する st.html()など、結構便利なものが多くあります。

2つめはStreamlitのいくつかの機能をサポートしていないことです。セキュリティ上の理由や技術的な成約などがあると思いますがそこそこの数の機能がサポートされていません。具体的にはこちらのページにまとまっています。

Streamlit in Snowflake の感想とまとめ

実は数か月前にもStreamlitを触ったことがありました。そのときはStreamlitの仕様を余り理解できておらず、プルダウンやスライドバーなどのウィジェットを少し動かすだけでコードがすべて再実行され、当然SQLも再発行されることでDBへの負担が大きい点や、レイアウトの不自由さなどを感じてすぐ使うのをやめてしまっていました。

しかし、改めて使い方を学んで、キャッシュ機能やレイアウトを整えるための各種メソッド、ウィジェットの使い方などを覚えていくと自分が作りたいWebアプリが非常に簡単に作れるライブラリだと感じるようになりました。

特に、僕はWeb開発の技術は殆ど持っていないので、こういうツールを作りたいなと思っても実際に作れるのはPythonバッチ等で実装できるものに限られていました。それがこれほど手軽にGUIを備えたアプリを作れるとなると、自分ができることの幅を大きく広げられそうだと感じています。

Streamlitはダッシュボード等のデータ分析系のアプリ開発に使うことを念頭において開発されているライブラリなので、今後はそちらの方面でも使っていきたいですね。

おしらせ

ライフイズテックサービス開発部では、気軽にご参加いただけるカジュアルなイベントを実施しています。開催予定のイベントは、 connpass のグループからご確認ください。興味のあるイベントがあったらぜひ参加登録をお願いいたします。皆さんのご参加をお待ちしています！

この記事が気に入ったらサポートをしてみませんか？