diff --git a/airflow/providers/amazon/aws/transfers/mysql_to_s3.py b/airflow/providers/amazon/aws/transfers/mysql_to_s3.py index 992e167834cf7..250f8f1f41d08 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,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) 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..8ef0a43cb5fd4 100644 --- a/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py +++ b/tests/providers/amazon/aws/transfers/test_mysql_to_s3.py @@ -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, @@ -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") @@ -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) @@ -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):