MWAAでRedshiftにIAM認証で接続してクエリを実行する（クロスアカウント）

# やりたいこと
Amazon Managed Workflows for Apache Airflow (MWAA) で**別のアカウント**にあるRedshiftクラスターに対してクエリを実行します。
同一アカウントに対してのクエリ実行はこちらの記事を参照してください。

https://qiita.com/yust0724/items/8841e4284cd5113e30f8

クロスアカウントの際のポイントは以下です。

- Redshiftアカウント側にIAMロールを作成して、MWAA実行ロールからAssumeRoleする
- Postgres Connectionを使う（理由は後述）

# 前提条件
- RedshiftのRA3クラスターをバージニア北部リージョン（us-east-1）に作成済み
    - Redshiftアカウントに作成（アカウントID`111111111111`）
    - クラスター名：redshift-cluster-demo
    - DB名：qs_test_db
    - ユーザー名：usera
    - パブリックアクセスを有効にしています
- MWAAのv2.7.2の環境を同じくバージニア北部リージョン（us-east-1）に作成済み
    - MWAAアカウントに作成（アカウントID`999999999999`）

# やってみる

## IAMロール作成 (Redshiftアカウント)
AssumeRole先のIAMロール `Redshift-Query-Role-MWAA` を作成する。
付与するポリシーは以下の通り。

```json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "redshift:GetClusterCredentials",
            "Resource": [
                "arn:aws:redshift:us-east-1:111111111111:dbuser:redshift-cluster-demo/usera",
                "arn:aws:redshift:us-east-1:111111111111:dbname:redshift-cluster-demo/qs_test_db"
            ]
        },
        {
            "Effect": "Allow",
            "Action": "redshift:DescribeClusters",
            "Resource": "*"
        }
    ]
}
```

また、IAMロールの信頼ポリシーは以下のように設定します。
クロスアカウントでのAssumeRoleなので、AssumeRoleしてくるロールに対して明示的に許可を与える必要があります。
`AmazonMWAA-MyAirflowEnvironment-public-us-east-1-272-IPav7H` は、MWAAアカウントに作成されている、MWAAの実行ロールの名前です。

```json
{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "Statement1",
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::999999999999:role/service-role/AmazonMWAA-MyAirflowEnvironment-public-us-east-1-272-IPav7H"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}
```

## Redshiftのセキュリティグループ修正 (Redshiftアカウント)
https://qiita.com/yust0724/items/8841e4284cd5113e30f8#redshift%E3%81%AE%E3%82%BB%E3%82%AD%E3%83%A5%E3%83%AA%E3%83%86%E3%82%A3%E3%82%B0%E3%83%AB%E3%83%BC%E3%83%97%E6%9B%B4%E6%96%B0

上記の記事と同じように、MWAAのVPC内のNAT GatewayのグローバルIPを許可します。
これで、Redshiftアカウント側の作業は完了です。

## MWAA実行ロールの修正 (MWAAアカウント)
AssumeRoleする側のIAMロールに、以下のポリシーを付与します。
こちらは、AssumeRoleすることだけを許可しておけばOKです。

```json
{
    "Version": "2012-10-17",
    "Statement": {
        "Effect": "Allow",
        "Action": "sts:AssumeRole",
        "Resource": "arn:aws:iam::111111111111:role/Redshift-Query-Role-MWAA"
    }
}
```

## AWS Connection作成 (MWAAアカウント)
さて、AssumeRoleするためにはどこかでAssumeRoleする対象のIAMロールを指定しなければなりません。そうしないと、どのIAMロールにAssumeRoleすればよいのか分かりませんよね。
ズバリ、それをここで指定します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2949357/b63dc8f8-75e0-b5fc-e2aa-94145731b28f.png)

Extraの内容で、以下のようにAssumeRoleする際のRoleを指定します。これは別アカウントにあるIAMロールでもOKです。

```
{
  "role_arn": "arn:aws:iam::111111111111:role/Redshift-Query-Role-MWAA",
  "region_name": "us-east-1",
  "assume_role_method": "assume_role"
}
```

AWS Connectionの設定値は以下の公式ドキュメントを参照してください（MWAA v2.7.1に使われている、apache-airflow-providers-amazon v8.7.1のドキュメントになってます）。

https://airflow.apache.org/docs/apache-airflow-providers-amazon/8.7.1/connections/aws.html

## Postgres Connection作成 (MWAAアカウント)
ここでは、Redshiftに接続するのにPostgres Connectionを作成します。
理由は、**Extraの中に "aws_conn_id" を設定するところがあるから**です。これを設定すると、先ほど作成したAWS Connectionを指定することでそこで記載した認証を使ってくれます。つまり、AssumeRoleしてくれます。
Redshift ConnectionにはAWS Connectionを指定する項目がないため、こちらの接続を使っています。もしRedshift Connectionを使ってIAM認証でクロスアカウント接続できる方法をご存じの方はぜひコメントください。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2949357/f8d82bd8-f10f-9a9f-1be2-a2af5c6bd522.png)

Extraの内容で、以下のように先ほど作成したAWS Connectionを指定します。

```
{
  "iam": true,
  "aws_conn_id": "aws_redshift_conn",
  "redshift": true,
  "cluster-identifier": "redshift-cluster-demo"
}
```

Postgres Connectionの設定値は以下の公式ドキュメントを参照してください（MWAA v2.7.1に使われている、apache-airflow-providers-postgres v5.6.1のドキュメントになってます）。

https://airflow.apache.org/docs/apache-airflow-providers-postgres/5.6.1/connections/postgres.html

## DAGコードをS3に格納
以下のコードをS3に格納して、DAGを実行します。
同一アカウントで利用したコードとほぼ同じです。conn_id の指定先のみ異なっています。

```python
from airflow import DAG
from airflow.utils.dates import days_ago

from datetime import datetime

from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator

default_args = {
    'owner': 'airflow',
    'start_date': days_ago(1),
    'depends_on_past': False
}

dag = DAG(
    'Redshift_Query_IAM_postgres_conn',
    description='testDAG',
    schedule_interval=None,
    default_args=default_args,
    start_date=datetime(2023, 10, 11),
    catchup=False
)

execute_query = SQLExecuteQueryOperator(
task_id="execute_query",
sql=f"SELECT 1;",
dag=dag,
conn_id ="postgres_redshift_role",
)

execute_query_2 = SQLExecuteQueryOperator(
task_id="execute_query_2",
sql=f"SELECT * FROM schema_a.table_a",
dag=dag,
conn_id ="postgres_redshift_role",
)

execute_query >> execute_query_2
```

## DAG実行
スクショは省略しますが、DAGは成功しました。

# おわりに
クロスアカウントでRedshiftになかなか接続できずにいましたが、Postgres Connectionを使うことで成功しました。ただ、Redshift Connectionを利用するのが正攻法だと思うので、ぜひこちらのConnectionを使ってクロスアカウント接続をしたいものです。