やりたいこと

Amazon Managed Workflows for Apache Airflow (MWAA) で同一アカウントにあるRedshiftクラスターに対してクエリを実行します。
また、Redshift内のDBユーザーに対しての認証は、ID/PASS認証ではなくIAM認証を使用します。つまり、MWAAの実行ロールに付与されているアクセス許可に基づいて一時的なDB認証を作成し、接続します。この方式だとDBユーザー名のみでパスワードは不要となります。

前置き

利用するOperator

今回は、Operatorとしてapache-airflow-providers-common-sqlのSQLExecuteQueryOperatorを利用します。

以前まではRedshiftに対してクエリを実行する際はRedshiftSQLOperatorを用いていましたが、apache-airflow-providers-amazonの v8.0.0 でSQLExecuteQueryOperatorを利用するように修正されました。

Removed deprecated RedshiftSQLOperator in favor of the generic SQLExecuteQueryOperator. The parameter that was passed as redshift_conn_id needs to be changed to conn_id, and the behavior should stay the same.

apache-airflow-providers-XXXって何？

Airflowは外部サービスとの連携のためにそれぞれ独自のプロパイダーパッケージを開発して提供しています。その中のAWSとの連携のために開発されているパッケージがapache-airflow-providers-amazonであり、DBサービスに対して共通的にSQLを発行するためのパッケージがapache-airflow-providers-common-sqlです。

このプロパイダーパッケージは独立して開発されているため、それぞれのバージョンを持っています。自身でAirflowをホストした場合はそれぞれのプロパイダーパッケージを個別にバージョンアップできますが、MWAAではバージョンが固定されています。
MWAAの公式ドキュメントから一部抜粋すると、Airflowのバージョンとそれに対応するapache-airflow-providers-amazon、apache-airflow-providers-common-sqlは以下の通りです。

Airflow v2.7.2
- AWS Connection : apache-airflow-providers-amazon[aiobotocore]==8.7.1
- Common SQL : apache-airflow-providers-common-sql==1.7.2
Airflow v2.6.3
- AWS Connection : apache-airflow-providers-amazon[aiobotocore]==8.2.0
- Common SQL : apache-airflow-providers-common-sql==1.5.2
Airflow v2.5.1
- AWS Connection : apache-airflow-providers-amazon[aiobotocore]==7.1.0
- Common SQL : apache-airflow-providers-common-sql==1.3.3

これを見ると、Redshiftに対してクエリを実行するOperatorがSQLExecuteQueryOperatorになったのは、apache-airflow-providers-amazonがv8.0.0の時なので、MWAAのAirflowで言うとv2.6.3からとなりますね。

SQLExecuteQueryOperatorの使い方

上記のドキュメントに記載されたParametersを指定します。
特にRedshiftに接続するためには、conn_idにRedshiftのConnectionを指定する必要があります。その他、sqlに実行するSQLを指定したり、databaseでConnectionに記載しているデータベース名をオーバーライトしたりできます。
今回の記事で、実際にこのOperatorを使ってみます。

前提条件

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

やってみる

実行ロールに権限付与

以下の内容のポリシーを作成して、MWAAの実行ロールに追加します（アカウントIDは999999999999に変更してます）。
今回はGetClusterCredentials-useraという名前のポリシーを作成しました。

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

以下のように、実行ロールに追加します。

Redshift Connection作成

Airflow UI上で、Redshiftの接続に使うConnectionを作成します。
上部のAdmin→Connectionsをクリックすると一覧が表示されますが、今回はその中のredshift_defaultを以下のように修正していきます。

Extraの中の"iam"をtrueに指定することで、ID/PASS認証ではなくIAM認証を利用します。

{
  "iam": true,
  "cluster_identifier": "redshift-cluster-demo",
  "db_user": "usera"
}

DAGコードをS3に格納

以下の.pyファイルをS3のdagsフォルダにアップロードします。
SQLExecuteQueryOperatorを2回使ってDAGを作成しています。2つ目のSQLで先ほどのConnectionで指定したRedshiftのクラスター、DBに同じくConnectionで指定したDBユーザーでIAM認証してSelect文を実行しています。

このままだとSELECT文を実行するだけで終わってしまうので、実際はSELECTしたデータを他のテーブルにINSERTしたり、新たにテーブルをCREATEしたりするケースが多いと思います。

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',
    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;",
split_statements=True,
return_last=False,
dag=dag,
conn_id="redshift_default",
)

execute_query_2 = SQLExecuteQueryOperator(
task_id="execute_query_2",
sql=f"SELECT * FROM schema_a.table_a",
split_statements=True,
return_last=False,
dag=dag,
conn_id ="redshift_default",
)

execute_query >> execute_query_2

上記のDAGは、以下のページを参考にしました。

Redshiftのセキュリティグループ更新

今回は、MWAAのVPC内に作成されている2つのNAT GatewayのグローバルIPを、Redshift側のセキュリティグループで許可します。これでネットワーク的にも接続されますね。

DAG実行

作成されたDAGを実行したところ、以下のように2つのTaskとも成功となりました（色々と実験していたので、何回か失敗の赤い棒グラフが表示されています）。

参考に、2つ目のSELECT文実行したTaskのログは以下のような感じでした。実行したSQL文はログにも出るんですね。

ip-10-192-20-231.ec2.internal
*** Reading remote log from Cloudwatch log_group: airflow-MyAirflowEnvironment-public-272-Task log_stream: dag_id=Redshift_Query_IAM/run_id=manual__2024-02-23T14_27_32.016961+00_00/task_id=execute_query_2/attempt=1.log.
[2024-02-23 14:27:39,775] Dependencies all met for dep_context=non-requeueable deps ti=<TaskInstance: Redshift_Query_IAM.execute_query_2 manual__2024-02-23T14:27:32.016961+00:00 [queued]>
[2024-02-23 14:27:39,791] Dependencies all met for dep_context=requeueable deps ti=<TaskInstance: Redshift_Query_IAM.execute_query_2 manual__2024-02-23T14:27:32.016961+00:00 [queued]>
[2024-02-23 14:27:39,792] Starting attempt 1 of 1
[2024-02-23 14:27:39,820] Executing <Task(SQLExecuteQueryOperator): execute_query_2> on 2024-02-23 14:27:32.016961+00:00
[2024-02-23 14:27:39,823] Started process 15173 to run task
[2024-02-23 14:27:39,825] Running: ['airflow', 'tasks', 'run', 'Redshift_Query_IAM', 'execute_query_2', 'manual__2024-02-23T14:27:32.016961+00:00', '--job-id', '25', '--raw', '--subdir', 'DAGS_FOLDER/redshift_query_dag_IAM.py', '--cfg-path', '/tmp/tmpdn9jl8e6']
[2024-02-23 14:27:39,826] Job 25: Subtask execute_query_2
[2024-02-23 14:27:40,120] Exporting env vars: AIRFLOW_CTX_DAG_OWNER='airflow' AIRFLOW_CTX_DAG_ID='Redshift_Query_IAM' AIRFLOW_CTX_TASK_ID='execute_query_2' AIRFLOW_CTX_EXECUTION_DATE='2024-02-23T14:27:32.016961+00:00' AIRFLOW_CTX_TRY_NUMBER='1' AIRFLOW_CTX_DAG_RUN_ID='manual__2024-02-23T14:27:32.016961+00:00'
[2024-02-23 14:27:40,120] Executing: SELECT * FROM schema_a.table_a
[2024-02-23 14:27:40,645] Running statement: SELECT * FROM schema_a.table_a, parameters: None
[2024-02-23 14:27:40,751] Marking task as SUCCESS. dag_id=Redshift_Query_IAM, task_id=execute_query_2, execution_date=20240223T142732, start_date=20240223T142739, end_date=20240223T142740
[2024-02-23 14:27:40,920] Task exited with return code 0
[2024-02-23 14:27:40,945] 0 downstream tasks scheduled from follow-on schedule check

Redshift側の確認

念のため、Redshiftクラスターの「クエリのモニタリング」タブから、クエリ履歴を見てみます。

意図したユーザーで実行されていますね。

注意

どこかのドキュメントに、RedshiftのConnectionのExtraの欄にprofile:defaultがあったのでそれをそのまま実行していたら、以下のようなエラーが出てしまいました。profileを指定すると、MWAAの実行ロールではなく指定したprofileが利用されるようです（当たり前と言えば当たり前なのですが…）。

redshift_connector.error.InterfaceError: The config profile (default) could not be found

おわりに

今回は、MWAAからRedshiftにIAM認証を使ってクエリを実行してみました。AWS内で完結しているのなら、認証は全てIAMにしたいというユースケースは結構あるのではないでしょうか。
実は別アカウントでも色々試してみて成功したので、次の記事に書ければと思ってます。