社内の分析・開発作業を効率化するために、VSCodeからAmazon Athenaに接続するextension(拡張機能)を作った話

※ この記事はモチベーションクラウドシリーズ [Advent Calendar 2022](https://qiita.com/advent-calendar/2022/mcs) の7日目の記事になります。（遅刻しました）

# はじめに

どうも、リンクアンドモチベーション・データユニット所属の みく / Ryo Koizumi (@slpwalks) です。

突然ですが、弊社ではデータ分析基盤・機械学習モデル構築のデータレイクとして主に [Amazon Athena](https://aws.amazon.com/jp/athena/) を利用しています。 
時に長大なクエリを書くこともありますが、クエリをデバッグしたり結果を出力する際に利用する Athena コンソールがどうにも使いづらく(vimキーバインドが使えないetc.)日々の業務においてなかなかストレスを感じていました。
普段使っている VSCode でもっと手軽にクエリをプレビュー出来ないかなと思い、既存の拡張機能も試しましたが、問題（後述）が発生したため自分で開発することにしました。

# 作ったもの

- VSCode Athena Viewer

## 内容

- VSCode から Amazon Athena にクエリを投げて結果をCSVで取得するための拡張機能です。
- Assume Role、及び MFA Challenge に対応しています。

# リリース・インストール方法

1. 以下より最新版のダウンロードを行なってください。（インストールは自己責任でお願いします）
[https://github.com/Mayu-mic/vscode-athena-viewer/releases](https://github.com/Mayu-mic/vscode-athena-viewer/releases)

2. 拡張機能ペインを開き、右上のメニューから VSIXファイルからのインストール を選択し、ダウンロードしたファイルを指定してください。
![スクリーンショット_2022-12-08_20_43_26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64915/51168840-ef1d-eb5e-1aab-b5cd0a8a99aa.png)

## 使い方

### Configuration (初回のみ)

- 左のペインより VSCode Athena Viewer を開いてください。
- メニューより、プロファイル・リージョン・ワークグループを設定してください。
![スクリーンショット_2022-12-08_20_38_28.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64915/87202124-2b9b-99bb-90fd-5730920157ba.png)

### Run

- SQLファイルを開くか、エディタの言語モードを SQL にしてください。
- コマンドパレットより、Amazon Athena: Run SQL を実行してください。

# 機能

現在実装している機能になります。

1. エディタの言語モードがSQLの時、コマンドパレット(cmd+shift+p) から Amazon Athena: Run SQL を実行することで、現在開いているエディタの内容でクエリを投げ、CSVを表示する機能
![スクリーンショット_2022-12-08_20_26_43.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64915/8c3f3953-f0d2-941b-4d3f-cfed1a92557e.png)

2. 左のペインからデータベース・テーブルの構造にアクセスし、内容をプレビューする機能（すぐスキーマ忘れるので）
![Notification_Center.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64915/8cb124a4-9d01-a02d-323d-746a33214e9d.png)

3. 実行ログからクエリーを再実行する機能

4. [Parameterized Query](https://docs.aws.amazon.com/ja_jp/athena/latest/ug/querying-with-prepared-statements.html) を解釈し、SQL実行後にパラメータを付与する機能（実験的）
![スクリーンショット 2022-12-08 20.31.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/64915/584a85c7-ee6f-ab5e-d8a5-b643f9bc8bbb.png)

# 自分で作る以外の選択肢はなかったの？!

[Athena Driver](https://marketplace.visualstudio.com/items?itemName=brunoventura.sqltools-athena-driver) は既に拡張機能として存在していて、こちらは [SQLTools for VSCode](https://marketplace.visualstudio.com/items?itemName=mtxr.sqltools) のプラグインとして実装されています。SQLTools は VSCode 上で SQL を扱うのに便利なツールですが、Athena Driver プラグインに関しては以下の問題点を発見しました。

1. 取得結果のページネーションに対応していない（1000件までしか取得できない）
2. AssumeRole に対応していない

そのうち1.は修正して[PRまで投げ](https://github.com/kovihq/sqltools-athena-driver/pull/25)ましたが、2.の対応をしようとして **AssumeRole時のMFA Codeの入力が実現できずに諦めました**。
SQLToolsは接続処理を拡張機能の実行とは異なるバックエンドのnode.jsに投げているようですが、これにより VSCodeのAPI(InputBox) を用いた入力が出来ず、SQLTools側が出力コンソールに投げているためnodeの標準入力も使えません。
Athena には普通のAPIを用いて接続すれば良いため、SQLTools をゴニョゴニョするより **もう自前で書いちゃったほうがいいな** と思いスタンドアロンでの実装に切り替えました。

# 作ってみて

- 作り方については [公式](https://code.visualstudio.com/api/get-started/your-first-extension) などを参照してください。
- Command, Workspace, Codelens など、概念がたくさんありますが、解説が不十分な印象があります。特に Codelens は公式での説明では全然分からず、紹介されている [Regex Previewer](https://marketplace.visualstudio.com/items?itemName=chrmarti.regex) の実装を見てようやく理解しました。
- しかしフルTypeScriptなので開発体験自体は良いと思います。デバッグも別窓で立ち上げるだけのため、開発においてストレスは感じませんでした。
- アイコンを用意するのが何気にむずいです。どなたかアイコンを書いてくれないでしょうか？

# 大きなフレームで業務改善を行う

今回エディタの拡張機能を作るのは初めての試みでしたが、普段使い慣れているエディタの機能を向上させるというのは日々の生産性を向上させるチャンスに溢れているなと思いました。

現在は、この拡張機能を前提としたリファクタリングを行なっています。SQLをコードから分離し、単体でデバッグ実行できるようにすることで、エディタを行ったり来たりしなくても良くなり、より効率的な開発が出来ることを期待しています。

みなさんも気軽に extensions を書いてみて、個人やチームの生産性を上げてみてください、ってことで本記事は〆です。ありがとうございました！