def execute(self, context):
# use the super to list all files in an Google Cloud Storage bucket
files = super().execute(context)
s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id,
verify=self.dest_verify)
if not self.replace:
# if we are not replacing -> list all files in the S3 bucket
# and only keep those files which are present in
# Google Cloud Storage and not in S3
bucket_name, _ = S3Hook.parse_s3_url(self.dest_s3_key)
existing_files = s3_hook.list_keys(bucket_name)
files = list(set(files) - set(existing_files))
if files:
hook = GoogleCloudStorageHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to)
for file in files:
file_bytes = hook.download(self.bucket, file)
dest_key = self.dest_s3_key + file
self.log.info("Saving file to %s", dest_key)
s3_hook.load_bytes(file_bytes,
key=dest_key,
replace=self.replace)
self.log.info("All done, uploaded %d files to S3", len(files))
else:
self.log.info("In sync, no files needed to be uploaded to S3")
return files
def execute(self, context):
# use the super to list all files in an Google Cloud Storage bucket
files = super(GoogleCloudStorageToS3Operator, self).execute(context)
s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id, verify=self.dest_verify)
if not self.replace:
# if we are not replacing -> list all files in the S3 bucket
# and only keep those files which are present in
# Google Cloud Storage and not in S3
bucket_name, _ = S3Hook.parse_s3_url(self.dest_s3_key)
existing_files = s3_hook.list_keys(bucket_name)
files = set(files) - set(existing_files)
if files:
hook = GoogleCloudStorageHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to
)
for file in files:
file_bytes = hook.download(self.bucket, file)
dest_key = self.dest_s3_key + file
self.log.info("Saving file to %s", dest_key)
s3_hook.load_bytes(file_bytes,
key=dest_key,
replace=self.replace)
self.log.info("All done, uploaded %d files to S3", len(files))
else:
self.log.info("In sync, no files needed to be uploaded to S3")
return files
def check_for_url(self, s3url):
"""
check if the s3url exists
:param s3url: S3 url
:type s3url:str
:return: bool
"""
bucket, key = S3Hook.parse_s3_url(s3url)
s3hook = S3Hook(aws_conn_id=self.aws_conn_id)
if not s3hook.check_for_bucket(bucket_name=bucket):
raise AirflowException(
"The input S3 Bucket {} does not exist ".format(bucket))
if not s3hook.check_for_key(key=key, bucket_name=bucket):
raise AirflowException(
"The input S3 Key {} does not exist in the Bucket".format(
s3url, bucket))
return True
def execute(self, context):
# use the super to list all files in an Google Cloud Storage bucket
files = super(GoogleCloudStorageToS3Operator, self).execute(context)
s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id,
verify=self.dest_verify)
if not self.replace:
# if we are not replacing -> list all files in the S3 bucket
# and only keep those files which are present in
# Google Cloud Storage and not in S3
bucket_name, prefix = S3Hook.parse_s3_url(self.dest_s3_key)
# look for the bucket and the prefix to avoid look into
# parent directories/keys
existing_files = s3_hook.list_keys(bucket_name, prefix=prefix)
# in case that no files exists, return an empty array to avoid errors
existing_files = existing_files if existing_files is not None else []
# remove the prefix for the existing files to allow the match
existing_files = [
file.replace(prefix, '', 1) for file in existing_files
]
files = list(set(files) - set(existing_files))
if files:
hook = GoogleCloudStorageHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to)
for file in files:
file_bytes = hook.download(self.bucket, file)
dest_key = self.dest_s3_key + file
self.log.info("Saving file to %s", dest_key)
s3_hook.load_bytes(file_bytes,
key=dest_key,
replace=self.replace)
self.log.info("All done, uploaded %d files to S3", len(files))
else:
self.log.info("In sync, no files needed to be uploaded to S3")
return files
def check_for_url(self, s3url):
"""
check if the s3url exists
:param s3url: S3 url
:type s3url:str
:return: bool
"""
bucket, key = S3Hook.parse_s3_url(s3url)
s3hook = S3Hook(aws_conn_id=self.aws_conn_id)
if not s3hook.check_for_bucket(bucket_name=bucket):
raise AirflowException(
"The input S3 Bucket {} does not exist ".format(bucket))
if key and not s3hook.check_for_key(key=key, bucket_name=bucket)\
and not s3hook.check_for_prefix(
prefix=key, bucket_name=bucket, delimiter='/'):
# check if s3 key exists in the case user provides a single file
# or if s3 prefix exists in the case user provides a prefix for files
raise AirflowException("The input S3 Key "
"or Prefix {} does not exist in the Bucket {}"
.format(s3url, bucket))
return True
def execute(self, context):
# use the super to list all files in an Google Cloud Storage bucket
files = super(MozGoogleCloudStorageToS3Operator, self).execute(context)
s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id,
verify=self.dest_verify)
if not self.replace:
# if we are not replacing -> list all files in the S3 bucket
# and only keep those files which are present in
# Google Cloud Storage and not in S3
bucket_name, _ = S3Hook.parse_s3_url(self.dest_s3_key)
existing_files = s3_hook.list_keys(bucket_name)
files = list(set(files) - set(existing_files))
if files:
gcs_hook = GoogleCloudStorageHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to,
)
def copy_file(source_key):
content = gcs_hook.download(self.bucket, source_key)
dest_key = self.dest_s3_key + source_key
self.log.info("Saving file to %s", dest_key)
s3_hook.load_bytes(content, key=dest_key, replace=self.replace)
if self.num_workers > 0:
pool = ThreadPool(self.num_workers)
try:
pool.map(copy_file, files, chunksize=1)
finally:
pool.close()
pool.join()
else:
for source_key in files:
copy_file(source_key)
self.log.info("All done, uploaded %d files to S3", len(files))
else:
self.log.info("In sync, no files needed to be uploaded to S3")
def execute(self, context):
# use the super to list all files in an Google Cloud Storage bucket
files = super().execute(context)
s3_hook = S3Hook(aws_conn_id=self.dest_aws_conn_id, verify=self.dest_verify)
if not self.replace:
# if we are not replacing -> list all files in the S3 bucket
# and only keep those files which are present in
# Google Cloud Storage and not in S3
bucket_name, prefix = S3Hook.parse_s3_url(self.dest_s3_key)
# look for the bucket and the prefix to avoid look into
# parent directories/keys
existing_files = s3_hook.list_keys(bucket_name, prefix=prefix)
# in case that no files exists, return an empty array to avoid errors
existing_files = existing_files if existing_files is not None else []
# remove the prefix for the existing files to allow the match
existing_files = [file.replace(prefix, '', 1) for file in existing_files]
files = list(set(files) - set(existing_files))
if files:
hook = GoogleCloudStorageHook(
google_cloud_storage_conn_id=self.google_cloud_storage_conn_id,
delegate_to=self.delegate_to
)
for file in files:
file_bytes = hook.download(self.bucket, file)
dest_key = self.dest_s3_key + file
self.log.info("Saving file to %s", dest_key)
s3_hook.load_bytes(file_bytes,
key=dest_key,
replace=self.replace)
self.log.info("All done, uploaded %d files to S3", len(files))
else:
self.log.info("In sync, no files needed to be uploaded to S3")
return files
def check_s3_url(self, s3url):
"""
Check if an S3 URL exists
:param s3url: S3 url
:type s3url: str
:rtype: bool
"""
bucket, key = S3Hook.parse_s3_url(s3url)
if not self.s3_hook.check_for_bucket(bucket_name=bucket):
raise AirflowException(
"The input S3 Bucket {} does not exist ".format(bucket))
if key and not self.s3_hook.check_for_key(key=key, bucket_name=bucket)\
and not self.s3_hook.check_for_prefix(
prefix=key, bucket_name=bucket, delimiter='/'):
# check if s3 key exists in the case user provides a single file
# or if s3 prefix exists in the case user provides multiple files in
# a prefix
raise AirflowException("The input S3 Key "
"or Prefix {} does not exist in the Bucket {}"
.format(s3url, bucket))
return True
def check_s3_url(self, s3url):
"""
Check if an S3 URL exists
:param s3url: S3 url
:type s3url: str
:rtype: bool
"""
bucket, key = S3Hook.parse_s3_url(s3url)
if not self.s3_hook.check_for_bucket(bucket_name=bucket):
raise AirflowException(
"The input S3 Bucket {} does not exist ".format(bucket))
if key and not self.s3_hook.check_for_key(key=key, bucket_name=bucket)\
and not self.s3_hook.check_for_prefix(
prefix=key, bucket_name=bucket, delimiter='/'):
# check if s3 key exists in the case user provides a single file
# or if s3 prefix exists in the case user provides multiple files in
# a prefix
raise AirflowException("The input S3 Key "
"or Prefix {} does not exist in the Bucket {}"
.format(s3url, bucket))
return True
def test_parse_s3_url(self):
parsed = S3Hook.parse_s3_url(self.s3_test_url)
self.assertEqual(parsed, ("test", "this/is/not/a-real-key.txt"),
"Incorrect parsing of the s3 url")
def test_parse_s3_url(self):
parsed = S3Hook.parse_s3_url(self.s3_test_url)
self.assertEqual(parsed,
("test", "this/is/not/a-real-key.txt"),
"Incorrect parsing of the s3 url")
