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]