Jupyter notebook on VSCodeでAthenaに接続

本職はアナリストではないでのすが、たまにJupyter notebookを使う場合、notebook自体のエディターは使いにくいのでVisual studio Codeでなんとかならないかと思ったら出来たのでメモ。

ついでにAWS Athenaへクエリーを投げる。やり方はいろいろありましたが、awswranglerを使うと簡単だったのでこれを利用。
awslabs/aws-data-wrangler - https://github.com/awslabs/aws-data-wrangler

VSCodeにPython/JupyterのExtensionをインストール

Python (ms-python.python Microsoft)
Jupyter (ms-toolsai.jupyter Microsoft)
Python自体は3.8をインストールしています。もちろん、awswranglerもインストール。

notebookファイルを作成

VSCodeでファイルを作成・保存する時に拡張子を「.ipynb」にすると自動でJupyter notebookのUIと同じものがVSCode上に表示されます。

やってることは簡単で、Athenaへクエリーを投げてその結果をpandasでplotしてるだけ。
サンプルのコード：

import sys, os, array, time

import awswrangler as wr
import numpy as np
import pandas as pd

wg = 'primary'
db = 'test'
query = "SELECT ymd, count FROM test.tbl_test01"

df = wr.athena.read_sql_query(
    query,
    database=db,
    workgroup=wg)
df.plot(figsize=(5,2), kind='bar', legend=False, title="", x='ymd')

実行してみる

今回はAWSに接続しているので、$HOME/.aws/credentials を準備しておく。
後は通常のnotebookと使い方は全く同じで実行ボタンを押すグラフを表示してくれます。便利。