Amazon CodeWhispererでLambdaのコード（Python）を実装してみた

# 背景・目的
先日、[AWS Glue データカタログへのクロスアカウントアクセスを試してみた](https://qiita.com/zumax/items/cbe75c5dee757c7720c1)では、Athenaから別アカウントのData Catalogにアクセスしましたがこれを、Lambdaを使用して取得したいと思います。

このLambdaを、先日、一般提供が発表された「[Amazon CodeWhisperer](https://aws.amazon.com/jp/blogs/news/announcing-new-tools-for-building-with-generative-ai-on-aws/)」を使用して、実装してみようと思います。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/dd5c26d4-7baf-c436-c237-645723c4a159.png)


# まとめ
- CodeWhisperでは、コードの補完、コメントからコードを生成できます。
- LambdaコンソールとIDEでは、下記の違いがあります。
    - Lambdaコンソールの場合は、単一のSuggest、手動で取得
    - IDEの場合は、複数のSuggest、自動で取得されるようです。


# 概要
## What is CodeWhisperer?
- リアルタイムでコードの推奨事項を提供する、機械学習を利用した汎用コードジェネレーター
- コードを書くと、CodeWhisperer は既存のコードとコメントに基づいて自動的に提案を生成する。 
- パーソナライズされたレコメンデーションは、1 行のコメントから完全に形成された関数まで、サイズと範囲がさまざま。
- CodeWhisperer はコードをスキャンして、セキュリティの問題を強調して定義することもできる。
- CodeWhisperer がサポートするプログラミング言語については、[言語サポート](https://docs.aws.amazon.com/ja_jp/codewhisperer/latest/userguide/language-ide-support.html#language-support)を参照。

## Amazon CodeWhisperer Features
### Support for popular programming languages and IDEs
一般的なプログラミング言語と IDE をサポートしています。
下記の言語に対して、コードSuggestを提供
- Python
- Java
- JavaScript
- TypeScript
- C#
- Go
- Rust
- PHP
- Ruby
- Kotlin
- C、C++
- シェル スクリプト
- SQL
- Scala

下記のIDEで利用できます。
- JetBrains IDE (IntelliJ IDEA、PyCharm、WebStorm、Rider)
- Visual Studio (VS) Code
- AWS Cloud9
- AWS Lambda コンソールなど


### Built-in security scans
- 上位 10 の Open Worldwide Application Security Project (OWASP) に含まれる脆弱性や、暗号ライブラリのベスト プラクティスを満たさない脆弱性など、見つけにくい脆弱性を検出できる。

### Enterprise administration
SSO と AWS IAM Identity Center の統合を使用して、ユーザーまたはグループに CodeWhisperer へのアクセスを提供する。 

## Amazon CodeWhisperer Pricing
下記の2つのTierがあります。
- CodeWhisperer Individual Tier
    - 無料
    - 一ヶ月あたり最大50回のコードスキャン/ユーザが無料で利用できる
- CodeWhisperer Professional Tier
    - 「ユーザーごと、月ごと」に計算され、暦月の請求期間中に CodeWhisperer にアクセスできるユーザーの最大数に基づいて、組織に毎月請求される。

### Security scans
- セキュリティ スキャンの実行は、IDE エディタで現在アクティブなファイルと、プロジェクトの依存ファイルに対してセキュリティ スキャンを実行。 スキャンが終了すると、スキャンされたファイルのセキュリティの問題が VSC の問題パネルで強調表示される。
- CodeWhisperer のセキュリティ スキャンは、Amazon CodeGuru と統合されている。

# 実装
:::note info
バージニアリージョンで実行しています。
:::
## 事前準備
### Lambda関数を作成
まずは、Lambda関数を作成します。ここではcross-account-by-lambdaという関数名で作成しました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/c534047c-fa72-d6c1-2f26-32375b27d04e.png)


### IAMポリシーの設定
1. 上記のLambda関数作成時に、自動でIAMロールを作成しました。このIAMロールのポリシーに下記を追加します。
    ```
    {
        "Sid": "CodeWhispererPermissions",
        "Effect": "Allow",
        "Action": [
            "codewhisperer:GenerateRecommendations"
        ],
        "Resource": "*"
    }
    ```
1. GlueカタログにアクセスするGlueのIAMポリシーを追加します。
    ```
    {
        "Effect": "Allow",
        "Action": [
            "glue:CreateDatabase",
            "glue:DeleteDatabase",
            "glue:GetDatabase",
            "glue:GetDatabases",
            "glue:UpdateDatabase",
            "glue:CreateTable",
            "glue:DeleteTable",
            "glue:BatchDeleteTable",
            "glue:UpdateTable",
            "glue:GetTable",
            "glue:GetTables",
            "glue:BatchCreatePartition",
            "glue:CreatePartition",
            "glue:DeletePartition",
            "glue:BatchDeletePartition",
            "glue:UpdatePartition",
            "glue:GetPartition",
            "glue:GetPartitions",
            "glue:BatchGetPartition"
        ],
        "Resource": [
            "*"
        ]
    }
    ```

1. Athenaの出力結果を格納するS3バケットのIAMポリシーを追加します。
    ```
    {
        "Effect": "Allow",
        "Action": [
            "s3:GetObject",
            "s3:PutObject",
            "s3:ListBucket"
        ],
        "Resource":[
            "arn:aws:s3:::{利用者のAthenaの結果を格納するバケット}",
            "arn:aws:s3:::{利用者のAthenaの結果を格納するバケット}/*"
        ]
    }
    ```

1. 提供側のS3バケットから取得する権限を付与します。
    ```
    {
        "Effect": "Allow",
        "Action": [
            "s3:GetObject",
            "s3:PutObject",
            "s3:ListBucket"
        ],
        "Resource":[
            "arn:aws:s3:::{提供側のデータを格納するバケット}",
            "arn:aws:s3:::{提供側のデータを格納するバケット}/*"
        ]
    }
    ```



1. 後ほどAthenaを実行するので、IAMポリシーをアタッチします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/19602e17-cc45-431a-0c5b-fdb0284565a3.png)


### CodeWhisperを有効化
1. Lambdaのコードソースで、Tools>「Amazon CodeWhisperer Code Suggestions」を選択し、有効にします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/3e65bde1-1645-c9b9-8449-b1d576ac5d43.png)

1. Preferencesで「Amazon CodeWhisperer Code Suggestions」が有効になっていることを確認します。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/6d6bf185-9fe7-e9ae-ee06-1872d48e1a2c.png)

### その他の設定
1. この後、Athenaを使用してクエリを実行します。デフォルトの３秒だと短いので、タイムアウト時間を3分に延ばしておきます。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/cf5e8d67-d57f-e5d7-15f1-12eac9e24678.png)


## デモ
[Introducing Amazon CodeWhisperer in the AWS Lambda console (In preview)](https://aws.amazon.com/blogs/compute/introducing-amazon-codewhisperer-in-the-aws-lambda-console-in-preview/)を参考に、操作感を試してみます。

:::note info
Lambdaコンソールの場合は、単一のSuggest、手動で取得、IDEの場合は、複数のSuggest、自動で取得されるようです。
:::
:::note info
以降から、Macを前提として記載しています。
:::
- Suggestを使うには、ショートカットキーを使用します。
    - Suggestを手動で取得する
        - Macの場合、Option+C
        - Windowsの場合、Alt+C
    - Suggestを受け入れる場合
        - Tab
    - Suggestを拒否
        - ESCまたは、Backspace

### とりあえず試す
1. import json直下でOption+C
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/1fbc19e7-0bf1-31b9-0d57-cd9cdafe1cb3.png)

1. Tabで決定
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/94b9116a-bfe7-f0a3-26e4-cdcca6b67152.png)


### Single-line code completion
行単位でコードの補完を行ってくれるか確認します。
:::note info
変数名によりSuggestが変わりました。
:::

1. `athena`と途中まで書いてOption＋CとTabを実行。boto3の初期化をSuggestしてくれました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/aec63c76-2937-27e9-f0c1-470a9fee62ea.png)

    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/e26dfad0-2fc2-796d-688b-7add79f0f960.png)

1. 次に`s3`と書いてみました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/abb3e91d-4c73-4eef-1bb8-367d9f9d3c07.png)
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/913f49c1-5e02-5f14-01dc-209b9940dd7f.png)


### Full function generation
関数単位で、コードの生成を行ってくれるか確認します。
1. `def read_from_s3`と書いて、Option+CとTabを実行。テンプレートを書いてくれました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/1b9c76ff-b319-5217-86a0-1056b933b440.png)
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/383db57e-72d4-8f75-bf4d-ed2c400e6905.png)

### Generate code from comments
コメントからコードの生成を行ってくれるか確認します。

1. `Function to upload a file to an S3 bucket`とコメントに書いて、Option+C＋Tabを実行。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/f85bec6c-c6fa-8a35-241b-84de9d57f1ba.png)

1. そのまま連続でOptopn+Cを実行すると続けてSuggestしてくれました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/c400ec29-842d-e7f3-86ba-35312b547b4e.png)


## Athenaのクライアントプログラムを書いてみる
### 2つの関数を作成する
Athenaに対してクエリを実行する関数と、クエリが正常に終了したか確認する関数を実装します。

1. 初期状態は下記の通りです。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/68996777-b5cd-0f3d-b4ef-23032da41d8b.png)

1. `Function to query using Athena`とコメントを入力します。query_athenaという関数が生成されました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/7b6b75e0-4bcb-6166-a665-f03db0096b71.png)

1. そのままOptions+Cを打ち込みます。雛形ができました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/034bb4b0-2d95-9849-c613-a7c757adf443.png)
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/9af3f930-a01d-7b67-3694-cd163bf1ef33.png)


1. 次に、Athenaはクエリの実行結果をチェックする必要があるので、`Function to check the result of "query_athena" function (input is query ID, output is true or false)`と入力します。check_query_statusという関数が生成されました。
1. そのままOptions+Cをクリックします。コードが生成されました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/64dfbfab-990e-faa7-79ea-580307b89db8.png)
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/17737bb6-d886-8f5a-f878-9b5e69b02a91.png)

### 仕様変更（戻り値の変更）
上記で作成したquery_athena関数の戻り値がresponseだったのでクエリIDを戻したいと思います。

1. `# return to queryId`とコメントし、Option+C & Tabで追加されました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/bb492568-40e0-5d39-f141-dcc81e69d4d9.png)

1. 元のコード（return response）は削除しておきます。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/49e933db-a333-ed6f-9c4f-5aa19a15f8ca.png)


### 仕様変更（仕様の変更）
Athenaのクエリ実行とチェックは非同期のため、実行後すぐに結果をチェックしても終了していない可能性があります。そのため最大値を設けて、定期的に確認するように仕様を変更します。（20秒ごとに最大3回確認します。）

1. コメントを追加し、Options+C＆Tabを実行します。生成されました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/dc5946e2-7384-149c-9a64-82d465e3fbdd.png)
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/745a5b4d-cb07-5365-58d3-44d467fc2c5a.png)

### 作成した関数を呼び出す
1. 関数（query_athenaとcheck_result）を呼び出すコードを生成します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/67c41c65-3adc-768a-298d-afcdc6f15e50.png)

### 残りのクエリを修正します。
下記については、手動で直します。
1. Athenaの出力結果であるS3バケット名を指定します。
1. 接続先を
1. クエリを変更します
    ```
    query_id = query_athena("SELECT * FROM account_a_catalog.crossaccount_access_db.test_table")
    ```

1. check_result関数のロジックが、このままでは一回目のチェックでいきなり抜けてしまうので、SUCCEEDEDではないときにSleepするように修正します。
    ```
    # Function to check the result of "query_athena" function (input :query ID, output :true or false) . 
    # Check the result of "query_athena" function every 20 seconds for max 3 times.
    def check_result(queryId):
        state = "RUNNING"
        client = boto3.client('athena')
        max_execution = 3
        
        while max_execution > 0 and state in ["RUNNING", "QUEUED"]:
            max_execution -= 1
            response = client.get_query_execution(QueryExecutionId=queryId)
            if "QueryExecution" in response and "Status" in response["QueryExecution"] and "State" in response["QueryExecution"]["Status"]:
                state = response["QueryExecution"]["Status"]["State"]
                if state == 'SUCCEEDED':
                    return True
    
            time.sleep(20)
    
        return False    
    ```

1. sleepを使用するために`import time`を追加します。

## 動作確認
1. Lambdaのコンソールから「テスト」をクリックします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/274dc2bf-ccc2-45c6-1667-c4e812548466.png)

1. Athenaの結果を確認します。ステータスが「完了済」になっています。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/20ec726c-31f2-ade6-897d-2f5e3d9cb913.png)

1. 結果をダウンロードします。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/78784aec-972d-2552-2c75-2e9228027974.png)

1. 確認できました。
    ![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/206276/df749464-70aa-ee3d-a995-8206eb7ead90.png)


# 考察
今回、CodeWhisperer を使用してLambdaをコード（Python）を実装してみました。
今までは、boto3のAPIなど、マニュアルを読んでコーディングするところから開始していましたが、CodeWhispererをすることで、ある程度実装された状態から始められる点は良いと感じました。

なお、今回はLambdaのコンソールを使用しましたが、IDEは複数のSuggestとキー入力が不要ということなので、今後はIDEで試してみたいと思います。

# 参考
https://aws.amazon.com/jp/blogs/news/announcing-new-tools-for-building-with-generative-ai-on-aws/

https://docs.aws.amazon.com/ja_jp/codewhisperer/latest/userguide/what-is-cwspr.html

https://docs.aws.amazon.com/ja_jp/codewhisperer/latest/userguide/lambda-setup.html

https://aws.amazon.com/blogs/compute/introducing-amazon-codewhisperer-in-the-aws-lambda-console-in-preview/

https://aws.amazon.com/jp/codewhisperer/features/