From b5ddb19fde71cb8fc047b1d05154b73198b2c5b2 Mon Sep 17 00:00:00 2001 From: Mario Taddeucci Date: Sun, 26 Dec 2021 15:51:23 -0300 Subject: [PATCH 1/2] Enable replace S3 file on MySqlToS3Operator --- airflow/providers/amazon/aws/transfers/mysql_to_s3.py | 6 +++++- tests/providers/amazon/aws/transfers/test_mysql_to_s3.py | 7 +++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py index 992e167834cf7..7bac852818d1b 100644 --- a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py @@ -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 `. :type mysql_conn_id: str :param aws_conn_id: reference to a specific S3 connection @@ -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, @@ -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 @@ -174,7 +178,7 @@ 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) diff --git a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py index 93a301ed113a3..7f5781247d0a3 100644 --- a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py @@ -16,6 +16,7 @@ # specific language governing permissions and limitations # under the License. # +from os import replace import unittest from tempfile import NamedTemporaryFile from unittest import mock @@ -48,6 +49,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, @@ -60,7 +62,7 @@ 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") @@ -85,6 +87,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) @@ -95,7 +98,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): From 0ce5ca4eeb8d397dd05a97a31ae4bd24ad73a448 Mon Sep 17 00:00:00 2001 From: Mario Taddeucci Date: Sun, 26 Dec 2021 16:06:40 -0300 Subject: [PATCH 2/2] pre-commit fix --- airflow/providers/amazon/aws/transfers/mysql_to_s3.py | 4 +++- tests/providers/amazon/aws/transfers/test_mysql_to_s3.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py index 7bac852818d1b..250f8f1f41d08 100644 --- a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py +++ b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py @@ -178,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, replace=self.replace) + 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) diff --git a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py index 7f5781247d0a3..8ef0a43cb5fd4 100644 --- a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py @@ -16,7 +16,6 @@ # specific language governing permissions and limitations # under the License. # -from os import replace import unittest from tempfile import NamedTemporaryFile from unittest import mock @@ -62,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, replace=True, + filename=f.name, + key=s3_key, + bucket_name=s3_bucket, + replace=True, ) @mock.patch("airflow.providers.amazon.aws.transfers.mysql_to_s3.NamedTemporaryFile")