Google ColaboratoryからTiDB Serverlessに接続する

はじめに

先日、みんなのPython勉強会でTiDB Serverlessの紹介をする機会がありました。
今までに100回を超えて実施している勉強会で発表させて頂いてありがとうございました！

その中でGoogle ColaboratoryからTiDB Serverlessに接続するNotebookを作成したのですが、公開し忘れていたのでここで公開します。データの分析にはもちろんですが、SQLを勉強したり、バグの再現手順とか書くのにNotebook形式は便利だと思います。是非試してみてください。

Notebookはこちらです

ノートブックの解説

事前作業

• TiDB Serverlessクラスタを作成
• TiDB Cloudコンソールの右下の "?" マークから、Interactive Tutorial → S&P 500 Analysis を選択して、サンプルデータをロード(下図参照)

必要なライブラリのインストール

他のpythonのmysqlライブラリでもいいと思うのですが、いつも使っているのはpymysqlです。
NotebookのセルにSQLを書いて実行したいため、jupysql をインストールしています。
また jupysql の初期化がなぜかf文字列ではできなかったので、接続オブジェクトを作成するために sqlalchemy もインストールしています。

!pip install pymysql jupysql sqlalchemy --quiet
%load_ext sql

TiDB Serverlessへの接続

Notebookを公開することも考えて、ユーザー名とパスワードはGoogle Colaboratoryのシークレットに保管します。

シークレットはノートブックから userdata を使ってアクセスできます。

これで接続文字列を作り、sqlalchemyのengineオブジェクトを作成し、%sql に指定することで、%sqlや%%sql が利用できます。

セルでのSQLの利用

%%sql をセルの最初に書くことで、そのセルでSQLを書くことができます。Gistではシンタックスハイライトされていませんが、Google Colaboratoryではシンタックスハイライトされてキーワードや関数が分かりやすくなります。

結果をpandas DataFrameとして取得

%sql を使うと、SELECTの結果をDataFrameとして取得できます。
ただ %sqlは一行しか書けないので、長いSQLはあらかじめ %%SQL --save [name] を使って別名をつけておくと、CTEのように利用できるのでそれを利用しています。