Coverage for airflow/operators/dagrun_operator.py: 46%

# Coverage for airflow/operators/dagrun_operator.py : 46%

## 28 statements 13 run 15 missing 0 excluded

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | # -*- coding: utf-8 -*- # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from datetime import datetime import logging from airflow.models import BaseOperator, DagRun from airflow.utils.decorators import apply_defaults from airflow import settings class DagRunOrder(object): def __init__(self, run_id=None, payload=None): self.run_id = run_id self.payload = payload class TriggerDagRunOperator(BaseOperator): """ Triggers a DAG run for a specified ``dag_id`` if a criteria is met :param trigger_dag_id: the dag_id to trigger :type trigger_dag_id: str :param python_callable: a reference to a python function that will be called while passing it the ``context`` object and a placeholder object ``obj`` for your callable to fill and return if you want a DagRun created. This ``obj`` object contains a ``run_id`` and ``payload`` attribute that you can modify in your function. The ``run_id`` should be a unique identifier for that DAG run, and the payload has to be a picklable object that will be made available to your tasks while executing that DAG run. Your function header should look like ``def foo(context, dag_run_obj):`` :type python_callable: python callable """ template_fields = tuple() template_ext = tuple() ui_color = '#ffefeb' @apply_defaults def __init__( self, trigger_dag_id, python_callable, *args, **kwargs): super(TriggerDagRunOperator, self).__init__(*args, **kwargs) self.python_callable = python_callable self.trigger_dag_id = trigger_dag_id def execute(self, context): dro = DagRunOrder(run_id='trig__' + datetime.now().isoformat()) dro = self.python_callable(context, dro) if dro: session = settings.Session() dr = DagRun( dag_id=self.trigger_dag_id, run_id=dro.run_id, conf=dro.payload, external_trigger=True) logging.info("Creating DagRun {}".format(dr)) session.add(dr) session.commit() session.close() else: logging.info("Criteria not met, moving on")