def execute(self, context: 'Context') -> str:
glacier_hook = GlacierHook(aws_conn_id=self.aws_conn_id)
gcs_hook = GCSHook(
gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
impersonation_chain=self.impersonation_chain,
)
job_id = glacier_hook.retrieve_inventory(vault_name=self.vault_name)
with tempfile.NamedTemporaryFile() as temp_file:
glacier_data = glacier_hook.retrieve_inventory_results(
vault_name=self.vault_name, job_id=job_id["jobId"])
# Read the file content in chunks using StreamingBody
# https://botocore.amazonaws.com/v1/documentation/api/latest/reference/response.html
stream = glacier_data["body"]
for chunk in stream.iter_chunk(chunk_size=self.chunk_size):
temp_file.write(chunk)
temp_file.flush()
gcs_hook.upload(
bucket_name=self.bucket_name,
object_name=self.object_name,
filename=temp_file.name,
gzip=self.gzip,
)
return f"gs://{self.bucket_name}/{self.object_name}"
def poke(self, context: 'Context') -> bool:
hook = GlacierHook(aws_conn_id=self.aws_conn_id)
response = hook.describe_job(vault_name=self.vault_name, job_id=self.job_id)
if response["StatusCode"] == JobStatus.SUCCEEDED.value:
self.log.info("Job status: %s, code status: %s", response["Action"], response["StatusCode"])
self.log.info("Job finished successfully")
return True
elif response["StatusCode"] == JobStatus.IN_PROGRESS.value:
self.log.info("Processing...")
self.log.warning("Code status: %s", response["StatusCode"])
return False
else:
raise AirflowException(
f'Sensor failed. Job status: {response["Action"]}, code status: {response["StatusCode"]}'
)
def setUp(self):
with mock.patch(
"airflow.providers.amazon.aws.hooks.glacier.GlacierHook.__init__",
return_value=None):
self.hook = GlacierHook(aws_conn_id="aws_default")
def execute(self, context: 'Context'):
hook = GlacierHook(aws_conn_id=self.aws_conn_id)
response = hook.retrieve_inventory(vault_name=self.vault_name)
return response