def test_execute( self, table_as_file_name, expected_s3_key, mock_run, mock_session, ): access_key = "aws_access_key_id" secret_key = "aws_secret_access_key" mock_session.return_value = Session(access_key, secret_key) schema = "schema" table = "table" s3_bucket = "bucket" s3_key = "key" unload_options = [ 'HEADER', ] RedshiftToS3Operator( schema=schema, table=table, s3_bucket=s3_bucket, s3_key=s3_key, unload_options=unload_options, include_header=True, redshift_conn_id="redshift_conn_id", aws_conn_id="aws_conn_id", task_id="task_id", table_as_file_name=table_as_file_name, dag=None, ).execute(None) unload_options = '\n\t\t\t'.join(unload_options) select_query = "SELECT * FROM {schema}.{table}".format(schema=schema, table=table) unload_query = """ UNLOAD ('{select_query}') TO 's3://{s3_bucket}/{s3_key}' with credentials 'aws_access_key_id={access_key};aws_secret_access_key={secret_key}' {unload_options}; """.format( select_query=select_query, s3_bucket=s3_bucket, s3_key=expected_s3_key, access_key=access_key, secret_key=secret_key, unload_options=unload_options, ) assert mock_run.call_count == 1 assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
def test_execute_sts_token( self, table_as_file_name, expected_s3_key, mock_run, mock_session, ): access_key = "ASIA_aws_access_key_id" secret_key = "aws_secret_access_key" token = "token" mock_session.return_value = Session(access_key, secret_key, token) mock_session.return_value.access_key = access_key mock_session.return_value.secret_key = secret_key mock_session.return_value.token = token schema = "schema" table = "table" s3_bucket = "bucket" s3_key = "key" unload_options = [ 'HEADER', ] op = RedshiftToS3Operator( schema=schema, table=table, s3_bucket=s3_bucket, s3_key=s3_key, unload_options=unload_options, include_header=True, redshift_conn_id="redshift_conn_id", aws_conn_id="aws_conn_id", task_id="task_id", table_as_file_name=table_as_file_name, dag=None, ) op.execute(None) unload_options = '\n\t\t\t'.join(unload_options) select_query = f"SELECT * FROM {schema}.{table}" credentials_block = build_credentials_block(mock_session.return_value) unload_query = op._build_unload_query( credentials_block, select_query, expected_s3_key, unload_options ) assert mock_run.call_count == 1 assert access_key in unload_query assert secret_key in unload_query assert token in unload_query assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
def test_custom_select_query_unloading( self, table, table_as_file_name, expected_s3_key, mock_run, mock_session, ): access_key = "aws_access_key_id" secret_key = "aws_secret_access_key" mock_session.return_value = Session(access_key, secret_key) mock_session.return_value.access_key = access_key mock_session.return_value.secret_key = secret_key mock_session.return_value.token = None s3_bucket = "bucket" s3_key = "key" unload_options = [ 'HEADER', ] select_query = "select column from table" op = RedshiftToS3Operator( select_query=select_query, table=table, table_as_file_name=table_as_file_name, s3_bucket=s3_bucket, s3_key=s3_key, unload_options=unload_options, include_header=True, redshift_conn_id="redshift_conn_id", aws_conn_id="aws_conn_id", task_id="task_id", dag=None, ) op.execute(None) unload_options = '\n\t\t\t'.join(unload_options) credentials_block = build_credentials_block(mock_session.return_value) unload_query = op._build_unload_query(credentials_block, select_query, expected_s3_key, unload_options) assert mock_run.call_count == 1 assert access_key in unload_query assert secret_key in unload_query assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
# KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from datetime import datetime from os import getenv from airflow import DAG from airflow.providers.amazon.aws.transfers.redshift_to_s3 import RedshiftToS3Operator S3_BUCKET_NAME = getenv("S3_BUCKET_NAME", "s3_bucket_name") S3_KEY = getenv("S3_KEY", "s3_key") REDSHIFT_TABLE = getenv("REDSHIFT_TABLE", "redshift_table") with DAG( dag_id="example_redshift_to_s3", start_date=datetime(2021, 1, 1), schedule_interval=None, catchup=False, tags=['example'], ) as dag: # [START howto_transfer_redshift_to_s3] task_transfer_redshift_to_s3 = RedshiftToS3Operator( task_id='transfer_redshift_to_s3', s3_bucket=S3_BUCKET_NAME, s3_key=S3_KEY, schema='PUBLIC', table=REDSHIFT_TABLE, ) # [END howto_transfer_redshift_to_s3]