使用 DAG 在 CLI 中匯入變數

下列範例程式碼使用亞馬遜上的 CLI 匯入變數，適用於 Apache 氣流的受管工作流程。

版本

• 您可以使用此頁面上的程式碼範例阿帕奇氣流 v2 及以上在蟒蛇。

先決條件

• 使用此頁面上的程式碼範例不需要其他權限。

許可

您的AWS帳戶需要存取AmazonMWAAAirflowCliAccess政策。如需進一步了解，請參閱 阿帕奇氣流 CLI 政策:亞馬遜 AirflowCliAccess。

相依性

• 若要將此程式碼範例與 Apache Airflow v2 搭配使用，不需要額外的相依性。該代碼使用阿帕奇氣流 V2 基地安裝在您的環境中。

程式碼範例

下列範例程式碼需要三個輸入：您的 Amazon MWAA 環境名稱 (mwaa_env), 該AWS您環境的區域 (位於aws_region)，以及包含您要匯入之變數的本機檔案 (var_file).

import boto3
import json
import requests 
import base64
import getopt
import sys

argv = sys.argv[1:]
mwaa_env=''
aws_region=''
var_file=''

try:
    opts, args = getopt.getopt(argv, 'e:v:r:', ['environment', 'variable-file','region'])
    #if len(opts) == 0 and len(opts) > 3:
    if len(opts) != 3:
        print ('Usage: -e MWAA environment -v variable file location and filename -r aws region')
    else:
        for opt, arg in opts:
            if opt in ("-e"):
                mwaa_env=arg
            elif opt in ("-r"):
                aws_region=arg
            elif opt in ("-v"):
                var_file=arg

        boto3.setup_default_session(region_name="{}".format(aws_region))
        mwaa_env_name = "{}".format(mwaa_env)

        client = boto3.client('mwaa')
        mwaa_cli_token = client.create_cli_token(
            Name=mwaa_env_name
        )
        
        with open ("{}".format(var_file), "r") as myfile:
            fileconf = myfile.read().replace('\n', '')

        json_dictionary = json.loads(fileconf)
        for key in json_dictionary:
            print(key, " ", json_dictionary[key])
            val = (key + " " + json_dictionary[key])
            mwaa_auth_token = 'Bearer ' + mwaa_cli_token['CliToken']
            mwaa_webserver_hostname = 'https://{0}/aws_mwaa/cli'.format(mwaa_cli_token['WebServerHostname'])
            raw_data = "variables set {0}".format(val)
            mwaa_response = requests.post(
                mwaa_webserver_hostname,
                headers={
                    'Authorization': mwaa_auth_token,
                    'Content-Type': 'text/plain'
                    },
                data=raw_data
                )
            mwaa_std_err_message = base64.b64decode(mwaa_response.json()['stderr']).decode('utf8')
            mwaa_std_out_message = base64.b64decode(mwaa_response.json()['stdout']).decode('utf8')
            print(mwaa_response.status_code)
            print(mwaa_std_err_message)
            print(mwaa_std_out_message)

except:
    print('Use this script with the following options: -e MWAA environment -v variable file location and filename -r aws region')
    print("Unexpected error:", sys.exc_info()[0])
    sys.exit(2)

後續步驟？

• 了解如何將此範例中的 DAG 程式碼上傳至dags您的亞馬遜 S3 存儲桶中的文件夾新增或更新 DAG。