Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Append region info to S3ToRedshitOperator if present #32328

Merged
merged 6 commits into from
Jul 11, 2023
Merged

Append region info to S3ToRedshitOperator if present #32328

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0907723
Append region info to S3ToRedshitOperator if present
richardbadman Jul 3, 2023
0be52c9
Merge remote-tracking branch 'apache/main' into add-region-to-copy-st…
richardbadman Jul 3, 2023
fa3e9a4
Update tests to make assertion checking valid
richardbadman Jul 3, 2023
958a048
Remove redundant comments
richardbadman Jul 5, 2023
694a3f4
Follow static check formatting
richardbadman Jul 6, 2023
d3cb69b
Remove optional parameter & None reference from method
richardbadman Jul 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions airflow/providers/amazon/aws/transfers/s3_to_redshift.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,13 +119,16 @@ def __init__(
if arg in self.redshift_data_api_kwargs.keys():
raise AirflowException(f"Cannot include param '{arg}' in Redshift Data API kwargs")

def _build_copy_query(self, copy_destination: str, credentials_block: str, copy_options: str) -> str:
def _build_copy_query(
self, copy_destination: str, credentials_block: str, region_info: str, copy_options: str
) -> str:
column_names = "(" + ", ".join(self.column_list) + ")" if self.column_list else ""
return f"""
COPY {copy_destination} {column_names}
FROM 's3://{self.s3_bucket}/{self.s3_key}'
credentials
'{credentials_block}'
{region_info}
Copy link
Member

@potiuk potiuk Jul 11, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will that work when region is None? I srsly doubt

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice catch. It can only be None when user passes null in the extra config. Should be handled explicitly.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good spot, I've updated the parameter to not default to None which prevents the possibility of this happening

{copy_options};
"""

Expand All @@ -139,7 +142,9 @@ def execute(self, context: Context) -> None:
else:
redshift_hook = RedshiftSQLHook(redshift_conn_id=self.redshift_conn_id)
conn = S3Hook.get_connection(conn_id=self.aws_conn_id)

region_info = ""
if conn.extra_dejson.get("region", False):
region_info = f"region '{conn.extra_dejson['region']}'"
if conn.extra_dejson.get("role_arn", False):
credentials_block = f"aws_iam_role={conn.extra_dejson['role_arn']}"
else:
Expand All @@ -151,7 +156,9 @@ def execute(self, context: Context) -> None:
destination = f"{self.schema}.{self.table}"
copy_destination = f"#{self.table}" if self.method == "UPSERT" else destination

copy_statement = self._build_copy_query(copy_destination, credentials_block, copy_options)
copy_statement = self._build_copy_query(
copy_destination, credentials_block, region_info, copy_options
)

sql: str | Iterable[str]

Expand Down
116 changes: 88 additions & 28 deletions tests/providers/amazon/aws/transfers/test_s3_to_redshift.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,17 +63,19 @@ def test_execute(self, mock_run, mock_session, mock_connection, mock_hook):
dag=None,
)
op.execute(None)
copy_query = """
expected_copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_access_key_id=aws_access_key_id;aws_secret_access_key=aws_secret_access_key'
;
"""
actual_copy_query = mock_run.call_args.args[0]

assert mock_run.call_count == 1
assert access_key in copy_query
assert secret_key in copy_query
assert_equal_ignore_multiple_spaces(mock_run.call_args.args[0], copy_query)
assert access_key in actual_copy_query
assert secret_key in actual_copy_query
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)

@mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
@mock.patch("airflow.models.connection.Connection")
Expand Down Expand Up @@ -110,17 +112,19 @@ def test_execute_with_column_list(self, mock_run, mock_session, mock_connection,
dag=None,
)
op.execute(None)
copy_query = """
expected_copy_query = """
COPY schema.table (column_1, column_2)
FROM 's3://bucket/key'
credentials
'aws_access_key_id=aws_access_key_id;aws_secret_access_key=aws_secret_access_key'
;
"""
actual_copy_query = mock_run.call_args.args[0]

assert mock_run.call_count == 1
assert access_key in copy_query
assert secret_key in copy_query
assert_equal_ignore_multiple_spaces(mock_run.call_args.args[0], copy_query)
assert access_key in actual_copy_query
assert secret_key in actual_copy_query
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)

@mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
@mock.patch("airflow.models.connection.Connection")
Expand Down Expand Up @@ -263,18 +267,20 @@ def test_execute_sts_token(self, mock_run, mock_session, mock_connection, mock_h
dag=None,
)
op.execute(None)
copy_statement = """
expected_copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_access_key_id=ASIA_aws_access_key_id;aws_secret_access_key=aws_secret_access_key;token=aws_secret_token'
;
"""
assert access_key in copy_statement
assert secret_key in copy_statement
assert token in copy_statement
actual_copy_query = mock_run.call_args.args[0]

assert access_key in actual_copy_query
assert secret_key in actual_copy_query
assert token in actual_copy_query
assert mock_run.call_count == 1
assert_equal_ignore_multiple_spaces(mock_run.call_args.args[0], copy_statement)
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)

@mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
@mock.patch("airflow.models.connection.Connection")
Expand Down Expand Up @@ -312,17 +318,68 @@ def test_execute_role_arn(self, mock_run, mock_session, mock_connection, mock_ho
dag=None,
)
op.execute(None)
copy_statement = """
expected_copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_iam_role=arn:aws:iam::112233445566:role/myRole'
;
"""
actual_copy_query = mock_run.call_args.args[0]

assert extra["role_arn"] in actual_copy_query
assert mock_run.call_count == 1
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)

@mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
@mock.patch("airflow.models.connection.Connection")
@mock.patch("boto3.session.Session")
@mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook.run")
def test_different_region(self, mock_run, mock_session, mock_connection, mock_hook):
access_key = "aws_access_key_id"
secret_key = "aws_secret_access_key"
extra = {"region": "eu-central-1"}
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

mock_connection.return_value = Connection(extra=extra)
mock_hook.return_value = Connection(extra=extra)

schema = "schema"
table = "table"
s3_bucket = "bucket"
s3_key = "key"
copy_options = ""

op = S3ToRedshiftOperator(
schema=schema,
table=table,
s3_bucket=s3_bucket,
s3_key=s3_key,
copy_options=copy_options,
redshift_conn_id="redshift_conn_id",
aws_conn_id="aws_conn_id",
task_id="task_id",
dag=None,
)
op.execute(None)
expected_copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_access_key_id=aws_access_key_id;aws_secret_access_key=aws_secret_access_key'
region 'eu-central-1'
;
"""
actual_copy_query = mock_run.call_args.args[0]

assert extra["role_arn"] in copy_statement
assert access_key in actual_copy_query
assert secret_key in actual_copy_query
assert extra["region"] in actual_copy_query
assert mock_run.call_count == 1
assert_equal_ignore_multiple_spaces(mock_run.call_args.args[0], copy_statement)
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)

def test_template_fields_overrides(self):
assert S3ToRedshiftOperator.template_fields == (
Expand Down Expand Up @@ -420,19 +477,10 @@ def test_using_redshift_data_api(self, mock_rs, mock_run, mock_session, mock_con
),
)
op.execute(None)
copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_access_key_id=aws_access_key_id;aws_secret_access_key=aws_secret_access_key'
;
"""

mock_run.assert_not_called()
assert access_key in copy_query
assert secret_key in copy_query

mock_rs.execute_statement.assert_called_once()
# test with all args besides sql
_call = deepcopy(mock_rs.execute_statement.call_args.kwargs)
_call.pop("Sql")
assert _call == dict(
Expand All @@ -443,8 +491,20 @@ def test_using_redshift_data_api(self, mock_rs, mock_run, mock_session, mock_con
StatementName=statement_name,
WithEvent=False,
)

expected_copy_query = """
COPY schema.table
FROM 's3://bucket/key'
credentials
'aws_access_key_id=aws_access_key_id;aws_secret_access_key=aws_secret_access_key'
;
"""
actual_copy_query = mock_rs.execute_statement.call_args.kwargs["Sql"]

mock_rs.describe_statement.assert_called_once_with(
Id="STATEMENT_ID",
)
# test sql arg
assert_equal_ignore_multiple_spaces(mock_rs.execute_statement.call_args.kwargs["Sql"], copy_query)

assert access_key in actual_copy_query
assert secret_key in actual_copy_query
assert_equal_ignore_multiple_spaces(actual_copy_query, expected_copy_query)