def execute(self, context):
xcom_data = xcom_pull(
context, {
'sync_task':
['instance_info', 'zip_blob', 'bin_data_source_blob'],
'setup_task': 'created_instances',
'unzip_task': 'status'
})
zip_blob = xcom_data['sync_task']['zip_blob']
bin_root_blob = xcom_data['sync_task']['bin_data_source_blob']
log.info('xcom data received: {}'.format(xcom_data))
file_paths = download_blob_by_name(
source_blob_name=zip_blob,
bucket_name=BUCKET_NAME,
save_file_root=os.path.expanduser('~/zip_bin_log'))
unzip_roots = []
for path in file_paths:
# unzipping the files
unzip_root = unzip(path)
os.remove(path) # remove the file
walktree_to_upload(
tree_root=unzip_root, root_blob=bin_root_blob, delete=True
) # TODO: Make this upload faster using parallel uploads
unzip_roots.append(unzip_root)
xcom_push(context, {'status': True})
log.info("unzipping complete")
def poke(self, context):
"""
creates an instance and pushes its id into xcom under 'created_instances` and returns False.
check if the unzip instance has returned its status to be true in xcom and then creates the
remaining instances.
"""
xcom_data = xcom_pull(
context, {
'sync_task':
['instance_info', 'bin_data_source_blob', 'zip_blob'],
'setup_task': 'created_instances',
'unzip_task': 'status',
'post_unzip_block_task': "started"
})
instance_info = xcom_data['sync_task']['instance_info']
bin_data_source_blob = xcom_data['sync_task']['bin_data_source_blob']
zip_blob = xcom_data['sync_task']['zip_blob']
created_instances = xcom_data['setup_task']['created_instances'] or []
unzip_status = xcom_data['unzip_task']['status']
post_unzip_block_started = xcom_data['post_unzip_block_task'][
'started']
log.info("xcom data received: {}".format(xcom_data))
EXPORT_ENV_CONFIG_WORKER = [{
"name": "NO_OF_INSTANCES",
"value": len(instance_info['instances'])
}, {
"name": "BIN_DATA_SOURCE_BLOB",
"value": bin_data_source_blob
}, {
"name": "ZIP_BLOB",
"value": zip_blob
}]
if unzip_status == True:
create_instances = instance_info['instances']
for instance in created_instances:
if instance in create_instances:
create_instances.remove(instance)
setup_instances(instances=create_instances,
export_configs=EXPORT_ENV_CONFIG_WORKER)
xcom_push(
context,
{'created_instances': created_instances + create_instances})
if post_unzip_block_started == True:
# return true if unzip is complete and the blocking of worker after unzipping has started
return True
else:
if len(created_instances) >= 1:
# instance is already
return False
create_instance = instance_info['instances'][0]
setup_instances(instances=create_instance,
export_configs=EXPORT_ENV_CONFIG_WORKER)
xcom_push(context, {'created_instances': [create_instance]})
return False
def execute(self, context):
log.info("working")
xcom_data = xcom_pull(context, {'sync_task': ['bin_data_source_blob']})
bin_data_source_blob = xcom_data['sync_task']['bin_data_source_blob']
log.info("xcom_data: {}".format(xcom_data))
worker_task(instance_no=self.operator_param['number'],
total_instances=self.operator_param['total'],
bin_data_source_blob=bin_data_source_blob)
xcom_push(context, {"status": True})
def poke(self, context):
xcom_push(context, {"started": True})
xcom_data = xcom_pull(context, {
'completion_task': ['started'],
})
completion_start = xcom_data['completion_task']['started']
log.info("xcom_data: {}".format(xcom_data))
if completion_start == True:
return True
return False
def execute(self, context):
params = self.operator_param
log.info("Sync in progress...")
# the program must be running in the airflow home thus, airflow_home = current working directory
ignores = ['.git', '/logs/', '.idea']
sync_folders(upload_blob_name=DESTINATION_BLOB_NAME,
folder_root=os.getcwd(),
bucket_name=BUCKET_NAME,
ignores=ignores)
log.info("Sync complete...")
data = {
"instance_info": params['instance_info'],
"bin_data_source_blob": params['bin_data_source_blob'],
"zip_blob": params['zip_blob'],
}
xcom_push(context, data)
log.info("xcom data pushed: {}".format(data))
def poke(self, context):
xcom_push(context, {"started": True})
# result = super(CompletionOperator, self).poke(context)
xcom_data = xcom_pull(
context,
{
'sync_task': ['instance_info'], # 'zip_blob'],
# 'setup_task': 'created_instances',
# 'unzip_task': 'status'
})
instance_info = xcom_data['sync_task']['instance_info']
total_instance = len(instance_info['instances'])
count = 0
for i in range(total_instance):
task_id = 'worker_task' + str(i)
work_status = xcom_pull(context,
{task_id: 'status'})[task_id]['status']
if work_status == True:
count += 1
log.info("count: {}, total_instances: {} ".format(
count, total_instance))
if count == total_instance:
xcom_push(context, {'status': True})
result = True
else:
result = False
if result:
log.info("deleting instances")
xcom_data = xcom_pull(
context, {
'sync_task': 'instance_info',
'setup_task': 'created_instances',
'unzip_task': 'status'
})
instance_info = xcom_data['sync_task']['instance_info']
log.info("Instance info received: " + str(instance_info))
delete_instances(instances=instance_info['instances'])
log.info("Instances deleted")
return result
def poke(self, context):
"""
creates an instance and pushes its id into xcom under 'created_instances` and returns False.
check if the unzip instance has returned its status to be true in xcom and then creates the
remaining instances.
"""
xcom_data = xcom_pull(
context, {
'sync_task': 'instance_info',
'setup_task': 'created_instances',
'unzip_task': 'status',
'post_unzip_block_task': "started"
})
instance_info = xcom_data['sync_task']['instance_info']
created_instances = xcom_data['setup_task']['created_instances'] or []
unzip_status = xcom_data['unzip_task']['status']
post_unzip_block_started = xcom_data['post_unzip_block_task'][
'started']
log.info("xcom data received: {}".format(xcom_data))
if unzip_status == True:
create_instances = instance_info['instances']
for instance in created_instances:
if instance in create_instances:
create_instances.remove(instance)
setup_instances(instances=create_instances)
xcom_push(context,
{'created_instances': instance_info['instances']})
if post_unzip_block_started == True:
# return true if unzip is complete and the blocking of worker after unzipping has started
return True
else:
if len(created_instances) >= 1:
# instance is already
return False
create_instance = instance_info['instances'][0]
setup_instances(instances=create_instance)
xcom_push(context, {'created_instances': [create_instance]})
return False