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

Jump to bottom

Add support to replace S3 file on MySqlToS3Operator #20506

Merged
merged 2 commits into from
Dec 29, 2021
Merged

Add support to replace S3 file on MySqlToS3Operator #20506

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift click to select a range
b5ddb19
Enable replace S3 file on MySqlToS3Operator
mariotaddeucci Dec 26, 2021
0ce5ca4
pre-commit fix
mariotaddeucci Dec 26, 2021
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
8 changes: 7 additions & 1 deletion airflow/providers/amazon/aws/transfers/mysql_to_s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 56,8 @@ class MySQLToS3Operator(BaseOperator):
:type s3_bucket: str
:param s3_key: desired key for the file. It includes the name of the file. (templated)
:type s3_key: str
:param replace: whether or not to replace the file in S3 if it previously existed
:type replace: bool
:param mysql_conn_id: Reference to :ref:`mysql connection id <howto/connection:mysql>`.
:type mysql_conn_id: str
:param aws_conn_id: reference to a specific S3 connection
Expand Down Expand Up @@ -101,6 103,7 @@ def __init__(
query: str,
s3_bucket: str,
s3_key: str,
replace: bool = False,
mysql_conn_id: str = 'mysql_default',
aws_conn_id: str = 'aws_default',
verify: Optional[Union[bool, str]] = None,
Expand All @@ -118,6 121,7 @@ def __init__(
self.mysql_conn_id = mysql_conn_id
self.aws_conn_id = aws_conn_id
self.verify = verify
self.replace = replace

if file_format == "csv":
self.file_format = FILE_FORMAT.CSV
Expand Down Expand Up @@ -174,7 178,9 @@ def execute(self, context) -> None:
data_df.to_csv(tmp_file.name, **self.pd_kwargs)
else:
data_df.to_parquet(tmp_file.name, **self.pd_kwargs)
s3_conn.load_file(filename=tmp_file.name, key=self.s3_key, bucket_name=self.s3_bucket)
s3_conn.load_file(
filename=tmp_file.name, key=self.s3_key, bucket_name=self.s3_bucket, replace=self.replace
)

if s3_conn.check_for_key(self.s3_key, bucket_name=self.s3_bucket):
file_location = os.path.join(self.s3_bucket, self.s3_key)
Expand Down
9 changes: 7 additions & 2 deletions tests/providers/amazon/aws/transfers/test_mysql_to_s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 48,7 @@ def test_execute_csv(self, mock_mysql_hook, mock_s3_hook, temp_mock):
aws_conn_id="aws_conn_id",
task_id="task_id",
index=True,
replace=True,
header=True,
pd_csv_kwargs={'index': False, 'header': False},
dag=None,
Expand All @@ -60,7 61,10 @@ def test_execute_csv(self, mock_mysql_hook, mock_s3_hook, temp_mock):

temp_mock.assert_called_once_with(mode='r ', suffix=".csv")
mock_s3_hook.return_value.load_file.assert_called_once_with(
filename=f.name, key=s3_key, bucket_name=s3_bucket
filename=f.name,
key=s3_key,
bucket_name=s3_bucket,
replace=True,
)

@mock.patch("airflow.providers.amazon.aws.transfers.mysql_to_s3.NamedTemporaryFile")
Expand All @@ -85,6 89,7 @@ def test_execute_parquet(self, mock_mysql_hook, mock_s3_hook, temp_mock):
aws_conn_id="aws_conn_id",
task_id="task_id",
file_format="parquet",
replace=False,
dag=None,
)
op.execute(None)
Expand All @@ -95,7 100,7 @@ def test_execute_parquet(self, mock_mysql_hook, mock_s3_hook, temp_mock):

temp_mock.assert_called_once_with(mode='rb ', suffix=".parquet")
mock_s3_hook.return_value.load_file.assert_called_once_with(
filename=f.name, key=s3_key, bucket_name=s3_bucket
filename=f.name, key=s3_key, bucket_name=s3_bucket, replace=False
)

def test_fix_int_dtypes(self):
Expand Down