Skip to content

Commit 3da0dcf

Browse files
bricefotzobricefotzo
committed
chore: make s3_to_sfpt and sftp_to_s3 consistent
1 parent 4335cd3 commit 3da0dcf

6 files changed

Lines changed: 372 additions & 149 deletions

File tree

‎airflow/providers/amazon/aws/transfers/s3_to_sftp.py‎

Lines changed: 37 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,3 @@
1-
#
2-
# Licensed to the Apache Software Foundation (ASF) under one
3-
# or more contributor license agreements. See the NOTICE file
4-
# distributed with this work for additional information
5-
# regarding copyright ownership. The ASF licenses this file
6-
# to you under the Apache License, Version 2.0 (the
7-
# "License"); you may not use this file except in compliance
8-
# with the License. You may obtain a copy of the License at
9-
#
10-
# http://www.apache.org/licenses/LICENSE-2.0
11-
#
12-
# Unless required by applicable law or agreed to in writing,
13-
# software distributed under the License is distributed on an
14-
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15-
# KIND, either express or implied. See the License for the
16-
# specific language governing permissions and limitations
17-
# under the License.
181
from __future__ import annotations
192

203
from tempfile import NamedTemporaryFile
@@ -23,29 +6,25 @@
236

247
from airflow.models import BaseOperator
258
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
26-
from airflow.providers.ssh.hooks.ssh import SSHHook
9+
from airflow.providers.sftp.hooks.sftp import SFTPHook
2710

2811
if TYPE_CHECKING:
2912
from airflow.utils.context import Context
3013

3114

3215
class S3ToSFTPOperator(BaseOperator):
3316
"""
34-
This operator enables the transferring of files from S3 to a SFTP server.
17+
Transfer files from Amazon S3 to an SFTP server using SFTPHook.
3518
36-
.. seealso::
37-
For more information on how to use this operator, take a look at the guide:
38-
:ref:`howto/operator:S3ToSFTPOperator`
39-
40-
:param sftp_conn_id: The sftp connection id. The name or identifier for
41-
establishing a connection to the SFTP server.
42-
:param sftp_path: The sftp remote path. This is the specified file path for
43-
uploading file to the SFTP server.
44-
:param aws_conn_id: aws connection to use
45-
:param s3_bucket: The targeted s3 bucket. This is the S3 bucket from
46-
where the file is downloaded.
47-
:param s3_key: The targeted s3 key. This is the specified file path for
48-
downloading the file from S3.
19+
:param sftp_conn_id: The SFTP connection id.
20+
:param sftp_path: The SFTP remote path for uploading files.
21+
:param aws_conn_id: The AWS connection id for S3.
22+
:param s3_bucket: The source S3 bucket.
23+
:param s3_key: The source S3 key. Can be a single key or a list of keys.
24+
:param sftp_filenames: Optional list of filenames to use when saving files to the SFTP server.
25+
If not provided, the original S3 filenames are used.
26+
:param replace: If True, will replace files on SFTP if they exist. Note: This functionality
27+
depends on the behavior of the SFTPHook and the SFTP server configuration.
4928
"""
5029

5130
template_fields: Sequence[str] = ("s3_key", "sftp_path", "s3_bucket")
@@ -54,33 +33,49 @@ def __init__(
5433
self,
5534
*,
5635
s3_bucket: str,
57-
s3_key: str,
36+
s3_key: str | list[str],
5837
sftp_path: str,
59-
sftp_conn_id: str = "ssh_default",
38+
sftp_filenames: str | list[str] | None = None,
39+
sftp_conn_id: str = "sftp_default",
6040
aws_conn_id: str = "aws_default",
41+
replace: bool = False, # Note: This parameter may not be relevant with SFTPHook's put method.
6142
**kwargs,
6243
) -> None:
6344
super().__init__(**kwargs)
6445
self.sftp_conn_id = sftp_conn_id
6546
self.sftp_path = sftp_path
6647
self.s3_bucket = s3_bucket
6748
self.s3_key = s3_key
49+
self.sftp_filenames = sftp_filenames
6850
self.aws_conn_id = aws_conn_id
51+
self.replace = replace
6952

7053
@staticmethod
7154
def get_s3_key(s3_key: str) -> str:
72-
"""Parse the correct format for S3 keys regardless of how the S3 url is passed."""
55+
"""Parse the correct format for S3 keys regardless of how the S3 URL is passed."""
7356
parsed_s3_key = urlsplit(s3_key)
57+
print(parsed_s3_key)
7458
return parsed_s3_key.path.lstrip("/")
7559

7660
def execute(self, context: Context) -> None:
77-
self.s3_key = self.get_s3_key(self.s3_key)
78-
ssh_hook = SSHHook(ssh_conn_id=self.sftp_conn_id)
7961
s3_hook = S3Hook(self.aws_conn_id)
62+
sftp_hook = SFTPHook(self.sftp_conn_id)
63+
64+
if isinstance(self.s3_key, str):
65+
print("s3_key is a string")
66+
print(self.s3_key)
67+
self.s3_key = [self.s3_key]
68+
self.sftp_filenames = [self.sftp_filenames] if self.sftp_filenames else [s3_key.split('/')[-1] for s3_key in self.s3_key]
8069

81-
s3_client = s3_hook.get_conn()
82-
sftp_client = ssh_hook.get_conn().open_sftp()
70+
for index, s3_key in enumerate(self.s3_key):
71+
print("s3_key is a list")
72+
print(s3_key)
73+
s3_file_key = self.get_s3_key(s3_key)
74+
sftp_filename = self.sftp_filenames[index] if self.sftp_filenames else s3_file_key.split('/')[-1]
75+
sftp_full_path = f"{self.sftp_path.rstrip('/')}/{sftp_filename}"
8376

84-
with NamedTemporaryFile("w") as f:
85-
s3_client.download_file(self.s3_bucket, self.s3_key, f.name)
86-
sftp_client.put(f.name, self.sftp_path)
77+
with NamedTemporaryFile("wb", delete=True) as tmp_file:
78+
print(f"Downloading file {s3_file_key} from S3 to {tmp_file.name}")
79+
s3_hook.download_file(bucket_name=self.s3_bucket, key=s3_file_key, local_path=tmp_file.name)
80+
sftp_hook.put(local_filepath=tmp_file.name, remote_filepath=sftp_full_path)
81+
self.log.info(f"File {s3_file_key} transferred to SFTP as {sftp_full_path}")

‎airflow/providers/amazon/aws/transfers/sftp_to_s3.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ class SFTPToS3Operator(BaseOperator):
4242
establishing a connection to the SFTP server.
4343
:param sftp_path: The sftp remote path. This is the specified file path
4444
for downloading the file from the SFTP server.
45-
:param s3_conn_id: The s3 connection id. The name or identifier for
45+
:param aws_conn_id: The s3 connection id. The name or identifier for
4646
establishing a connection to S3
4747
:param s3_bucket: The targeted s3 bucket. This is the S3 bucket to where
4848
the file is uploaded.
@@ -63,7 +63,7 @@ def __init__(
6363
sftp_filenames: str | list[str] | None = None,
6464
s3_filenames: str | list[str] | None = None,
6565
sftp_conn_id: str = "sftp_default",
66-
s3_conn_id: str = "aws_default",
66+
aws_conn_id: str = "aws_default",
6767
replace: bool = False,
6868
use_temp_file: bool = True,
6969
**kwargs,
@@ -73,7 +73,7 @@ def __init__(
7373
self.sftp_path = sftp_path
7474
self.s3_bucket = s3_bucket
7575
self.s3_key = s3_key
76-
self.s3_conn_id = s3_conn_id
76+
self.aws_conn_id = aws_conn_id
7777
self.use_temp_file = use_temp_file
7878
self.replace = replace
7979
self.sftp_filenames = sftp_filenames
@@ -103,7 +103,7 @@ def __upload_to_s3_from_sftp(self, remote_filename, s3_file_key):
103103

104104
def execute(self, context: Context):
105105
self.sftp_hook = SFTPHook(ssh_conn_id=self.sftp_conn_id)
106-
self.s3_hook = S3Hook(self.s3_conn_id)
106+
self.s3_hook = S3Hook(self.aws_conn_id)
107107

108108
def resolve_filenames(pattern, filenames):
109109
if pattern == "*":
@@ -128,7 +128,7 @@ def generate_s3_key(filename):
128128
s3_file_key = generate_s3_key(file)
129129
self.__upload_to_s3_from_sftp(file, s3_file_key)
130130

131-
else: # Assuming ftp_filenames is a list
131+
else:
132132
for sftp_file, s3_file in zip(self.sftp_filenames, self.s3_filenames if self.s3_filenames else self.sftp_filenames):
133133
sftp_path = self.sftp_path + sftp_file
134134
s3_key = self.s3_key + (s3_file if self.s3_filenames else sftp_file)

‎generated/provider_dependencies.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
"mongo",
6262
"openlineage",
6363
"salesforce",
64+
"sftp",
6465
"ssh"
6566
],
6667
"excluded-python-versions": [],

‎tests/providers/amazon/aws/transfers/test_s3_to_sftp copy.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
from __future__ import annotations
19+
20+
import boto3
21+
import pytest
22+
from moto import mock_aws
23+
24+
from airflow.models import DAG
25+
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
26+
from airflow.providers.amazon.aws.transfers.s3_to_sftp import S3ToSFTPOperator
27+
from airflow.providers.sftp.hooks.sftp import SFTPHook
28+
from airflow.providers.ssh.operators.ssh import SSHOperator
29+
from airflow.utils.timezone import datetime
30+
from tests.test_utils.config import conf_vars
31+
32+
pytestmark = pytest.mark.db_test
33+
34+
35+
TASK_ID = "test_sftp_to_s3"
36+
BUCKET = "test-s3-bucket"
37+
S3_KEY = "test/test_1_file.csv"
38+
SFTP_PATH = "/tmp/remote_path.txt"
39+
AWS_CONN_ID = "aws_default"
40+
SFTP_CONN_ID = "sftp_default"
41+
S3_KEY_MULTIPLE = "test/"
42+
SFTP_PATH_MULTIPLE = "/tmp/"
43+
44+
TEST_DAG_ID = "unit_tests_s3_to_sftp"
45+
DEFAULT_DATE = datetime(2018, 1, 1)
46+
47+
48+
class TestS3ToSFTPOperator:
49+
def setup_method(self):
50+
hook = SSHHook(ssh_conn_id="ssh_default")
51+
hook.no_host_key_check = True
52+
dag = DAG(
53+
f"{TEST_DAG_ID}test_schedule_dag_once",
54+
start_date=DEFAULT_DATE,
55+
schedule="@once",
56+
)
57+
58+
self.hook = hook
59+
60+
self.ssh_client = self.hook.get_conn()
61+
self.sftp_client = self.ssh_client.open_sftp()
62+
63+
self.dag = dag
64+
self.s3_bucket = BUCKET
65+
self.sftp_path = SFTP_PATH
66+
self.s3_key = S3_KEY
67+
68+
@mock_aws
69+
@conf_vars({("core", "enable_xcom_pickling"): "True"})
70+
def test_s3_to_sftp_operation(self):
71+
s3_hook = S3Hook(aws_conn_id=None)
72+
# Setting
73+
test_remote_file_content = (
74+
"This is remote file content \n which is also multiline "
75+
"another line here \n this is last line. EOF"
76+
)
77+
78+
# Test for creation of s3 bucket
79+
conn = boto3.client("s3")
80+
conn.create_bucket(Bucket=self.s3_bucket)
81+
assert s3_hook.check_for_bucket(self.s3_bucket)
82+
83+
with open(LOCAL_FILE_PATH, "w") as file:
84+
file.write(test_remote_file_content)
85+
s3_hook.load_file(LOCAL_FILE_PATH, self.s3_key, bucket_name=BUCKET)
86+
87+
# Check if object was created in s3
88+
objects_in_dest_bucket = conn.list_objects(Bucket=self.s3_bucket, Prefix=self.s3_key)
89+
# there should be object found, and there should only be one object found
90+
assert len(objects_in_dest_bucket["Contents"]) == 1
91+
92+
# the object found should be consistent with dest_key specified earlier
93+
assert objects_in_dest_bucket["Contents"][0]["Key"] == self.s3_key
94+
95+
# get remote file to local
96+
run_task = S3ToSFTPOperator(
97+
s3_bucket=BUCKET,
98+
s3_key=S3_KEY,
99+
sftp_path=SFTP_PATH,
100+
sftp_conn_id=SFTP_CONN_ID,
101+
task_id=TASK_ID,
102+
dag=self.dag,
103+
)
104+
assert run_task is not None
105+
106+
run_task.execute(None)
107+
108+
# Check that the file is created remotely
109+
check_file_task = SSHOperator(
110+
task_id="test_check_file",
111+
ssh_hook=self.hook,
112+
command=f"cat {self.sftp_path}",
113+
do_xcom_push=True,
114+
dag=self.dag,
115+
)
116+
assert check_file_task is not None
117+
result = check_file_task.execute(None)
118+
assert result.strip() == test_remote_file_content.encode("utf-8")
119+
120+
# Clean up after finishing with test
121+
conn.delete_object(Bucket=self.s3_bucket, Key=self.s3_key)
122+
conn.delete_bucket(Bucket=self.s3_bucket)
123+
assert not s3_hook.check_for_bucket(self.s3_bucket)
124+
125+
def delete_remote_resource(self):
126+
# check the remote file content
127+
remove_file_task = SSHOperator(
128+
task_id="test_rm_file",
129+
ssh_hook=self.hook,
130+
command=f"rm {self.sftp_path}",
131+
do_xcom_push=True,
132+
dag=self.dag,
133+
)
134+
assert remove_file_task is not None
135+
remove_file_task.execute(None)
136+
137+
def teardown_method(self):
138+
self.delete_remote_resource()

0 commit comments

Comments
 (0)