def create_dagrun(self, *args, **kwargs):
run_args = "{}" # TODO extract the run Args from the tasks
marquez_jobrun_id = None
try:
marquez_jobrun_id = self.report_jobrun(run_args,
kwargs['execution_date'])
log.info(f'Successfully recorded job run.',
airflow_dag_id=self.dag_id,
marquez_run_id=marquez_jobrun_id,
marquez_namespace=self.marquez_namespace)
except Exception as e:
log.error(f'Failed to record job run: {e}',
airflow_dag_id=self.dag_id,
marquez_namespace=self.marquez_namespace)
pass
run = super(DAG, self).create_dagrun(*args, **kwargs)
if marquez_jobrun_id:
try:
self._job_id_mapping.set(
JobIdMapping.make_key(run.dag_id, run.run_id),
marquez_jobrun_id)
except Exception as e:
log.error(f'Failed job run lookup: {e}',
airflow_dag_id=self.dag_id,
airflow_run_id=run.run_id,
marquez_run_id=marquez_jobrun_id,
marquez_namespace=self.marquez_namespace)
pass
return run
def report_jobrun_change(self, job_name, run_id, **kwargs):
session = kwargs.get('session')
marquez_job_run_ids = self._job_id_mapping.pop(
JobIdMapping.make_key(job_name, run_id), session)
if marquez_job_run_ids:
log.info('Found job runs.',
airflow_dag_id=self.dag_id,
airflow_job_id=job_name,
airflow_run_id=run_id,
marquez_run_ids=marquez_job_run_ids,
marquez_namespace=self.marquez_namespace)
ids = json.loads(marquez_job_run_ids)
if kwargs.get('success'):
for marquez_job_run_id in ids:
self.get_marquez_client().mark_job_run_as_completed(
marquez_job_run_id)
else:
for marquez_job_run_id in ids:
self.get_marquez_client().mark_job_run_as_failed(
marquez_job_run_id)
state = 'COMPLETED' if kwargs.get('success') else 'FAILED'
log.info(f'Marked job run(s) as {state}.',
airflow_dag_id=self.dag_id,
airflow_job_id=job_name,
airflow_run_id=run_id,
marquez_run_id=marquez_job_run_ids,
marquez_namespace=self.marquez_namespace)
def report_jobrun_change(self, dagrun, **kwargs):
session = kwargs.get('session')
marquez_job_run_id = self._job_id_mapping.pop(
JobIdMapping.make_key(dagrun.dag_id, dagrun.run_id), session)
if marquez_job_run_id:
log.info(f'Found job run.',
airflow_dag_id=dagrun.dag_id,
airflow_run_id=dagrun.run_id,
marquez_run_id=marquez_job_run_id,
marquez_namespace=self.marquez_namespace)
if kwargs.get('success'):
self.get_marquez_client().mark_job_run_as_completed(
marquez_job_run_id)
else:
self.get_marquez_client().mark_job_run_as_failed(
marquez_job_run_id)
state = 'COMPLETED' if kwargs.get('success') else 'FAILED'
log.info(f'Marked job run as {state}.',
airflow_dag_id=dagrun.dag_id,
airflow_run_id=dagrun.run_id,
marquez_run_id=marquez_job_run_id,
marquez_namespace=self.marquez_namespace)
def report_task(self, dag_run_id, execution_date, run_args, task,
extractor):
report_job_start_ms = self._now_ms()
marquez_client = self.get_marquez_client()
if execution_date:
start_time = self._to_iso_8601(execution_date)
end_time = self.compute_endtime(execution_date)
else:
start_time = None
end_time = None
if end_time:
end_time = self._to_iso_8601(end_time)
task_location = None
try:
if hasattr(task, 'file_path') and task.file_path:
task_location = get_location(task.file_path)
else:
task_location = get_location(task.dag.fileloc)
except Exception:
log.warn('Unable to fetch the location')
steps_metadata = []
if extractor:
try:
log.info(f'Using extractor {extractor.__name__}',
task_type=task.__class__.__name__,
airflow_dag_id=self.dag_id,
task_id=task.task_id,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
steps_metadata = extractor(task).extract()
except Exception as e:
log.error(f'Failed to extract metadata {e}',
airflow_dag_id=self.dag_id,
task_id=task.task_id,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
else:
log.warn('Unable to find an extractor.',
task_type=task.__class__.__name__,
airflow_dag_id=self.dag_id,
task_id=task.task_id,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
task_name = f'{self.dag_id}.{task.task_id}'
# If no extractor found or failed to extract metadata,
# report the task metadata
if not steps_metadata:
steps_metadata = [StepMetadata(task_name)]
# store all the JobRuns associated with a task
marquez_jobrun_ids = []
for step in steps_metadata:
input_datasets = []
output_datasets = []
try:
input_datasets = self.register_datasets(step.inputs)
except Exception as e:
log.error(f'Failed to register inputs: {e}',
inputs=str(step.inputs),
airflow_dag_id=self.dag_id,
task_id=task.task_id,
step=step.name,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
try:
output_datasets = self.register_datasets(step.outputs)
except Exception as e:
log.error(f'Failed to register outputs: {e}',
outputs=str(step.outputs),
airflow_dag_id=self.dag_id,
task_id=task.task_id,
step=step.name,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
marquez_client.create_job(
job_name=step.name,
job_type='BATCH', # job type
location=(step.location or task_location),
input_dataset=input_datasets,
output_dataset=output_datasets,
context=step.context,
description=self.description,
namespace_name=self.marquez_namespace)
log.info(f'Successfully recorded job: {step.name}',
airflow_dag_id=self.dag_id,
marquez_namespace=self.marquez_namespace)
marquez_jobrun_id = marquez_client.create_job_run(
step.name,
run_args=run_args,
nominal_start_time=start_time,
nominal_end_time=end_time).get('runId')
if marquez_jobrun_id:
marquez_jobrun_ids.append(marquez_jobrun_id)
marquez_client.mark_job_run_as_started(marquez_jobrun_id)
else:
log.error(f'Failed to get run id: {step.name}',
airflow_dag_id=self.dag_id,
airflow_run_id=dag_run_id,
marquez_namespace=self.marquez_namespace)
log.info(f'Successfully recorded job run: {step.name}',
airflow_dag_id=self.dag_id,
airflow_dag_execution_time=start_time,
marquez_run_id=marquez_jobrun_id,
marquez_namespace=self.marquez_namespace,
duration_ms=(self._now_ms() - report_job_start_ms))
# Store the mapping for all the steps associated with a task
try:
self._job_id_mapping.set(
JobIdMapping.make_key(task_name, dag_run_id),
json.dumps(marquez_jobrun_ids))
except Exception as e:
log.error(f'Failed to set id mapping : {e}',
airflow_dag_id=self.dag_id,
task_id=task.task_id,
airflow_run_id=dag_run_id,
marquez_run_id=marquez_jobrun_ids,
marquez_namespace=self.marquez_namespace)