やりたいこと

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

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

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

前提条件

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

やってみる

IAMロール作成 (Redshiftアカウント)

AssumeRole先のIAMロール Redshift-Query-Role-MWAA を作成する。
付与するポリシーは以下の通り。

{
    "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の実行ロールの名前です。

{
    "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アカウント)

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

MWAA実行ロールの修正 (MWAAアカウント)

AssumeRoleする側のIAMロールに、以下のポリシーを付与します。
こちらは、AssumeRoleすることだけを許可しておけばOKです。

{
    "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すればよいのか分かりませんよね。
ズバリ、それをここで指定します。

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のドキュメントになってます）。

Postgres Connection作成 (MWAAアカウント)

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

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のドキュメントになってます）。

DAGコードをS3に格納

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

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を使ってクロスアカウント接続をしたいものです。

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