@@ -101,6 +101,13 @@ def clean_all(self):
101101 rmtree (self .path )
102102 clean_tmp ()
103103
104+ def _get_s3_extra (self , upload = False ):
105+ """Helper to build the extra dict for S3 operations."""
106+ extra = {"ACL" : "bucket-owner-full-control" } if upload else {}
107+ if self .config .get ("requester_pay" , False ):
108+ extra ["RequestPayer" ] = "requester"
109+ return extra
110+
104111 def download_file_from_s3 (self , s3file , working_dir ):
105112 """ Download s3 file to local
106113
@@ -117,7 +124,8 @@ def download_file_from_s3(self, s3file, working_dir):
117124 full path of the downloaded file
118125 """
119126 try :
120- return s3 .download (s3file , working_dir )
127+ extra = self ._get_s3_extra ()
128+ return s3 .download (s3file , working_dir , extra = extra )
121129 except botocore .exceptions .ClientError as ex :
122130 self .logger .error ("Error downloading file %s: %s" % (s3file , working_dir ), exc_info = True )
123131 raise ex
@@ -133,7 +141,8 @@ def upload_file_to_s3(self, filename, uri):
133141 s3 string of file location
134142 """
135143 try :
136- return s3 .upload (filename , uri , extra = {"ACL" : "bucket-owner-full-control" })
144+ extra = self ._get_s3_extra (upload = True )
145+ return s3 .upload (filename , uri , extra = extra )
137146 except botocore .exceptions .ClientError as ex :
138147 self .logger .error ("Error uploading file %s: %s" % (os .path .basename (os .path .basename (filename )), str (ex )), exc_info = True )
139148 raise ex
@@ -370,7 +379,8 @@ def _download_file(self, file_):
370379 """Download the input file from S3."""
371380 input_file = f's3://{ file_ ["bucket" ]} /{ file_ ["key" ]} '
372381 try :
373- return s3 .download (input_file , path = self .path )
382+ extra = self ._get_s3_extra ()
383+ return s3 .download (input_file , path = self .path , extra = extra )
374384 except botocore .exceptions .ClientError as ex :
375385 self .logger .error ("Error downloading file from S3: {}" .format (ex ), exc_info = True )
376386 raise
0 commit comments