Source code for airflow.models

# -*- coding: utf-8 -*-## Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you 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__future__importabsolute_importfrom__future__importdivisionfrom__future__importprint_functionfrom__future__importunicode_literalsfrombuiltinsimportImportErrorasBuiltinImportError,bytes,object,strfromcollectionsimportdefaultdict,namedtuple,OrderedDictimportcopyfromtypingimportIterablefromfuture.standard_libraryimportinstall_aliasesfromairflow.models.baseimportBase,ID_LENtry:# Fix Python > 3.7 deprecationfromcollections.abcimportHashableexceptImportError:# Preserve Python < 3.3 compatibilityfromcollectionsimportHashablefromdatetimeimporttimedeltaimportdillimportfunctoolsimportgetpassimportimpimportimportlibimportzipfileimportjinja2importjsonimportloggingimportosimportpendulumimportpickleimportreimportsignalimportsysimporttimeimporttextwrapimporttracebackimportwarningsimporthashlibfromdatetimeimportdatetimefromurllib.parseimportquotefromsqlalchemyimport(Boolean,Column,DateTime,Float,Index,Integer,PickleType,String,Text,UniqueConstraint,and_,func,or_)fromsqlalchemy.ext.declarativeimportdeclared_attrfromsqlalchemy.ormimportreconstructor,synonymfromcroniterimport(croniter,CroniterBadCronError,CroniterBadDateError,CroniterNotAlphaError)importsixfromairflowimportsettings,utilsfromairflow.executorsimportget_default_executor,LocalExecutorfromairflowimportconfigurationfromairflow.exceptionsimport(AirflowDagCycleException,AirflowException,AirflowSkipException,AirflowTaskTimeout,AirflowRescheduleException)fromairflow.dag.base_dagimportBaseDag,BaseDagBagfromairflow.lineageimportapply_lineage,prepare_lineagefromairflow.models.dagpickleimportDagPicklefromairflow.models.kubernetesimportKubeWorkerIdentifier,KubeResourceVersion# noqa: F401fromairflow.models.logimportLogfromairflow.models.taskfailimportTaskFailfromairflow.models.taskrescheduleimportTaskReschedulefromairflow.models.xcomimportXComfromairflow.ti_deps.deps.not_in_retry_period_depimportNotInRetryPeriodDepfromairflow.ti_deps.deps.prev_dagrun_depimportPrevDagrunDepfromairflow.ti_deps.deps.trigger_rule_depimportTriggerRuleDepfromairflow.ti_deps.dep_contextimportDepContext,QUEUE_DEPS,RUN_DEPSfromairflow.utilsimporttimezonefromairflow.utils.dag_processingimportlist_py_file_pathsfromairflow.utils.datesimportcron_presets,date_rangeasutils_date_rangefromairflow.utils.dbimportprovide_sessionfromairflow.utils.decoratorsimportapply_defaultsfromairflow.utils.emailimportsend_emailfromairflow.utils.helpersimportis_container,validate_key,pprinttablefromairflow.utils.operator_resourcesimportResourcesfromairflow.utils.stateimportStatefromairflow.utils.sqlalchemyimportUtcDateTime,Intervalfromairflow.utils.timeoutimporttimeoutfromairflow.utils.trigger_ruleimportTriggerRulefromairflow.utils.weight_ruleimportWeightRulefromairflow.utils.netimportget_hostnamefromairflow.utils.log.logging_mixinimportLoggingMixininstall_aliases()XCOM_RETURN_KEY='return_value'Stats=settings.Stats

[docs]classInvalidFernetToken(Exception):# If Fernet isn't loaded we need a valid exception class to catch. If it is# loaded this will get reset to the actual class once get_fernet() is calledpass

[docs]classNullFernet(object):""" A "Null" encryptor class that doesn't encrypt or decrypt but that presents a similar interface to Fernet. The purpose of this is to make the rest of the code not have to know the difference, and to only display the message once, not 20 times when `airflow initdb` is ran. """is_encrypted=Falsedefdecrpyt(self,b):returnbdefencrypt(self,b):returnb

_fernet=None

[docs]defget_fernet():""" Deferred load of Fernet key. This function could fail either because Cryptography is not installed or because the Fernet key is invalid. :return: Fernet object :raises: airflow.exceptions.AirflowException if there's a problem trying to load Fernet """global_fernetlog=LoggingMixin().logif_fernet:return_fernettry:fromcryptography.fernetimportFernet,MultiFernet,InvalidTokenglobalInvalidFernetTokenInvalidFernetToken=InvalidTokenexceptBuiltinImportError:log.warning("cryptography not found - values will not be stored encrypted.")_fernet=NullFernet()return_fernettry:fernet_key=configuration.conf.get('core','FERNET_KEY')ifnotfernet_key:log.warning("empty cryptography key - values will not be stored encrypted.")_fernet=NullFernet()else:_fernet=MultiFernet([Fernet(fernet_part.encode('utf-8'))forfernet_partinfernet_key.split(',')])_fernet.is_encrypted=Trueexcept(ValueError,TypeError)asve:raiseAirflowException("Could not create Fernet object: {}".format(ve))return_fernet

# Used by DAG context_managers_CONTEXT_MANAGER_DAG=None

[docs]defclear_task_instances(tis,session,activate_dag_runs=True,dag=None,):""" Clears a set of task instances, but makes sure the running ones get killed. :param tis: a list of task instances :param session: current session :param activate_dag_runs: flag to check for active dag run :param dag: DAG object """job_ids=[]fortiintis:ifti.state==State.RUNNING:ifti.job_id:ti.state=State.SHUTDOWNjob_ids.append(ti.job_id)else:task_id=ti.task_idifdaganddag.has_task(task_id):task=dag.get_task(task_id)task_retries=task.retriesti.max_tries=ti.try_number+task_retries-1else:# Ignore errors when updating max_tries if dag is None or# task not found in dag since database records could be# outdated. We make max_tries the maximum value of its# original max_tries or the current task try number.ti.max_tries=max(ti.max_tries,ti.try_number-1)ti.state=State.NONEsession.merge(ti)ifjob_ids:fromairflow.jobsimportBaseJobasBJforjobinsession.query(BJ).filter(BJ.id.in_(job_ids)).all():job.state=State.SHUTDOWNifactivate_dag_runsandtis:drs=session.query(DagRun).filter(DagRun.dag_id.in_({ti.dag_idfortiintis}),DagRun.execution_date.in_({ti.execution_datefortiintis}),).all()fordrindrs:dr.state=State.RUNNINGdr.start_date=timezone.utcnow()

[docs]defget_last_dagrun(dag_id,session,include_externally_triggered=False):""" Returns the last dag run for a dag, None if there was none. Last dag run can be any type of run eg. scheduled or backfilled. Overridden DagRuns are ignored. """DR=DagRunquery=session.query(DR).filter(DR.dag_id==dag_id)ifnotinclude_externally_triggered:query=query.filter(DR.external_trigger==False)# noqaquery=query.order_by(DR.execution_date.desc())returnquery.first()

[docs]classDagBag(BaseDagBag,LoggingMixin):""" A dagbag is a collection of dags, parsed out of a folder tree and has high level configuration settings, like what database to use as a backend and what executor to use to fire off tasks. This makes it easier to run distinct environments for say production and development, tests, or for different teams or security profiles. What would have been system level settings are now dagbag level so that one system can run multiple, independent settings sets. :param dag_folder: the folder to scan to find DAGs :type dag_folder: unicode :param executor: the executor to use when executing task instances in this DagBag :param include_examples: whether to include the examples that ship with airflow or not :type include_examples: bool :param has_logged: an instance boolean that gets flipped from False to True after a file has been skipped. This is to prevent overloading the user with logging messages about skipped files. Therefore only once per DagBag is a file logged being skipped. """# static class variables to detetct dag cycleCYCLE_NEW=0CYCLE_IN_PROGRESS=1CYCLE_DONE=2def__init__(self,dag_folder=None,executor=None,include_examples=configuration.conf.getboolean('core','LOAD_EXAMPLES'),safe_mode=configuration.conf.getboolean('core','DAG_DISCOVERY_SAFE_MODE')):# do not use default arg in signature, to fix import cycle on plugin loadifexecutorisNone:executor=get_default_executor()dag_folder=dag_folderorsettings.DAGS_FOLDERself.log.info("Filling up the DagBag from %s",dag_folder)self.dag_folder=dag_folderself.dags={}# the file's last modified timestamp when we last read itself.file_last_changed={}self.executor=executorself.import_errors={}self.has_logged=Falseself.collect_dags(dag_folder=dag_folder,include_examples=include_examples,safe_mode=safe_mode)

[docs]defsize(self):""" :return: the amount of dags contained in this dagbag """returnlen(self.dags)

@propertydefdag_ids(self):returnself.dags.keys()

[docs]defget_dag(self,dag_id):""" Gets the DAG out of the dictionary, and refreshes it if expired """# If asking for a known subdag, we want to refresh the parentroot_dag_id=dag_idifdag_idinself.dags:dag=self.dags[dag_id]ifdag.is_subdag:root_dag_id=dag.parent_dag.dag_id# If the dag corresponding to root_dag_id is absent or expiredorm_dag=DagModel.get_current(root_dag_id)iform_dagand(root_dag_idnotinself.dagsor(orm_dag.last_expiredanddag.last_loaded<orm_dag.last_expired)):# Reprocess source filefound_dags=self.process_file(filepath=orm_dag.fileloc,only_if_updated=False)# If the source file no longer exports `dag_id`, delete it from self.dagsiffound_dagsanddag_idin[found_dag.dag_idforfound_daginfound_dags]:returnself.dags[dag_id]elifdag_idinself.dags:delself.dags[dag_id]returnself.dags.get(dag_id)

[docs]defprocess_file(self,filepath,only_if_updated=True,safe_mode=True):""" Given a path to a python module or zip file, this method imports the module and look for dag objects within it. """found_dags=[]# if the source file no longer exists in the DB or in the filesystem,# return an empty list# todo: raise exception?iffilepathisNoneornotos.path.isfile(filepath):returnfound_dagstry:# This failed before in what may have been a git sync# race conditionfile_last_changed_on_disk=datetime.fromtimestamp(os.path.getmtime(filepath))ifonly_if_updated \
andfilepathinself.file_last_changed \
andfile_last_changed_on_disk==self.file_last_changed[filepath]:returnfound_dagsexceptExceptionase:self.log.exception(e)returnfound_dagsmods=[]is_zipfile=zipfile.is_zipfile(filepath)ifnotis_zipfile:ifsafe_modeandos.path.isfile(filepath):withopen(filepath,'rb')asf:content=f.read()ifnotall([sincontentforsin(b'DAG',b'airflow')]):self.file_last_changed[filepath]=file_last_changed_on_disk# Don't want to spam user with skip messagesifnotself.has_logged:self.has_logged=Trueself.log.info("File %s assumed to contain no DAGs. Skipping.",filepath)returnfound_dagsself.log.debug("Importing %s",filepath)org_mod_name,_=os.path.splitext(os.path.split(filepath)[-1])mod_name=('unusual_prefix_'+hashlib.sha1(filepath.encode('utf-8')).hexdigest()+'_'+org_mod_name)ifmod_nameinsys.modules:delsys.modules[mod_name]withtimeout(configuration.conf.getint('core',"DAGBAG_IMPORT_TIMEOUT")):try:m=imp.load_source(mod_name,filepath)mods.append(m)exceptExceptionase:self.log.exception("Failed to import: %s",filepath)self.import_errors[filepath]=str(e)self.file_last_changed[filepath]=file_last_changed_on_diskelse:zip_file=zipfile.ZipFile(filepath)formodinzip_file.infolist():head,_=os.path.split(mod.filename)mod_name,ext=os.path.splitext(mod.filename)ifnotheadand(ext=='.py'orext=='.pyc'):ifmod_name=='__init__':self.log.warning("Found __init__.%s at root of %s",ext,filepath)ifsafe_mode:withzip_file.open(mod.filename)aszf:self.log.debug("Reading %s from %s",mod.filename,filepath)content=zf.read()ifnotall([sincontentforsin(b'DAG',b'airflow')]):self.file_last_changed[filepath]=(file_last_changed_on_disk)# todo: create ignore list# Don't want to spam user with skip messagesifnotself.has_logged:self.has_logged=Trueself.log.info("File %s assumed to contain no DAGs. Skipping.",filepath)ifmod_nameinsys.modules:delsys.modules[mod_name]try:sys.path.insert(0,filepath)m=importlib.import_module(mod_name)mods.append(m)exceptExceptionase:self.log.exception("Failed to import: %s",filepath)self.import_errors[filepath]=str(e)self.file_last_changed[filepath]=file_last_changed_on_diskforminmods:fordaginlist(m.__dict__.values()):ifisinstance(dag,DAG):ifnotdag.full_filepath:dag.full_filepath=filepathifdag.fileloc!=filepathandnotis_zipfile:dag.fileloc=filepathtry:dag.is_subdag=Falseself.bag_dag(dag,parent_dag=dag,root_dag=dag)ifisinstance(dag._schedule_interval,six.string_types):croniter(dag._schedule_interval)found_dags.append(dag)found_dags+=dag.subdagsexcept(CroniterBadCronError,CroniterBadDateError,CroniterNotAlphaError)ascron_e:self.log.exception("Failed to bag_dag: %s",dag.full_filepath)self.import_errors[dag.full_filepath]= \
"Invalid Cron expression: "+str(cron_e)self.file_last_changed[dag.full_filepath]= \
file_last_changed_on_diskexceptAirflowDagCycleExceptionascycle_exception:self.log.exception("Failed to bag_dag: %s",dag.full_filepath)self.import_errors[dag.full_filepath]=str(cycle_exception)self.file_last_changed[dag.full_filepath]= \
file_last_changed_on_diskself.file_last_changed[filepath]=file_last_changed_on_diskreturnfound_dags

[docs]defbag_dag(self,dag,parent_dag,root_dag):""" Adds the DAG into the bag, recurses into sub dags. Throws AirflowDagCycleException if a cycle is detected in this dag or its subdags """dag.test_cycle()# throws if a task cycle is founddag.resolve_template_files()dag.last_loaded=timezone.utcnow()fortaskindag.tasks:settings.policy(task)subdags=dag.subdagstry:forsubdaginsubdags:subdag.full_filepath=dag.full_filepathsubdag.parent_dag=dagsubdag.is_subdag=Trueself.bag_dag(subdag,parent_dag=dag,root_dag=root_dag)self.dags[dag.dag_id]=dagself.log.debug('Loaded DAG %s',dag)exceptAirflowDagCycleExceptionascycle_exception:# There was an error in bagging the dag. Remove it from the list of dagsself.log.exception('Exception bagging dag: {dag.dag_id}'.format(**locals()))# Only necessary at the root level since DAG.subdags automatically# performs DFS to search through all subdagsifdag==root_dag:forsubdaginsubdags:ifsubdag.dag_idinself.dags:delself.dags[subdag.dag_id]raisecycle_exception

[docs]defcollect_dags(self,dag_folder=None,only_if_updated=True,include_examples=configuration.conf.getboolean('core','LOAD_EXAMPLES'),safe_mode=configuration.conf.getboolean('core','DAG_DISCOVERY_SAFE_MODE')):""" Given a file path or a folder, this method looks for python modules, imports them and adds them to the dagbag collection. Note that if a ``.airflowignore`` file is found while processing the directory, it will behave much like a ``.gitignore``, ignoring files that match any of the regex patterns specified in the file. **Note**: The patterns in .airflowignore are treated as un-anchored regexes, not shell-like glob patterns. """start_dttm=timezone.utcnow()dag_folder=dag_folderorself.dag_folder# Used to store stats around DagBag processingstats=[]FileLoadStat=namedtuple('FileLoadStat',"file duration dag_num task_num dags")forfilepathinlist_py_file_paths(dag_folder,safe_mode=safe_mode,include_examples=include_examples):try:ts=timezone.utcnow()found_dags=self.process_file(filepath,only_if_updated=only_if_updated,safe_mode=safe_mode)td=timezone.utcnow()-tstd=td.total_seconds()+(float(td.microseconds)/1000000)stats.append(FileLoadStat(filepath.replace(dag_folder,''),td,len(found_dags),sum([len(dag.tasks)fordaginfound_dags]),str([dag.dag_idfordaginfound_dags]),))exceptExceptionase:self.log.exception(e)Stats.gauge('collect_dags',(timezone.utcnow()-start_dttm).total_seconds(),1)Stats.gauge('dagbag_size',len(self.dags),1)Stats.gauge('dagbag_import_errors',len(self.import_errors),1)self.dagbag_stats=sorted(stats,key=lambdax:x.duration,reverse=True)

[docs]classTaskInstance(Base,LoggingMixin):""" Task instances store the state of a task instance. This table is the authority and single source of truth around what tasks have run and the state they are in. The SqlAlchemy model doesn't have a SqlAlchemy foreign key to the task or dag model deliberately to have more control over transactions. Database transactions on this table should insure double triggers and any confusion around what task instances are or aren't ready to run even while multiple schedulers may be firing task instances. """__tablename__="task_instance"task_id=Column(String(ID_LEN),primary_key=True)dag_id=Column(String(ID_LEN),primary_key=True)execution_date=Column(UtcDateTime,primary_key=True)start_date=Column(UtcDateTime)end_date=Column(UtcDateTime)duration=Column(Float)state=Column(String(20))_try_number=Column('try_number',Integer,default=0)max_tries=Column(Integer)hostname=Column(String(1000))unixname=Column(String(1000))job_id=Column(Integer)pool=Column(String(50))queue=Column(String(50))priority_weight=Column(Integer)operator=Column(String(1000))queued_dttm=Column(UtcDateTime)pid=Column(Integer)executor_config=Column(PickleType(pickler=dill))__table_args__=(Index('ti_dag_state',dag_id,state),Index('ti_dag_date',dag_id,execution_date),Index('ti_state',state),Index('ti_state_lkp',dag_id,task_id,execution_date,state),Index('ti_pool',pool,state,priority_weight),Index('ti_job_id',job_id),)def__init__(self,task,execution_date,state=None):self.dag_id=task.dag_idself.task_id=task.task_idself.task=taskself._log=logging.getLogger("airflow.task")# make sure we have a localized execution_date stored in UTCifexecution_dateandnottimezone.is_localized(execution_date):self.log.warning("execution date %s has no timezone information. Using ""default from dag or system",execution_date)ifself.task.has_dag():execution_date=timezone.make_aware(execution_date,self.task.dag.timezone)else:execution_date=timezone.make_aware(execution_date)execution_date=timezone.convert_to_utc(execution_date)self.execution_date=execution_dateself.queue=task.queueself.pool=task.poolself.priority_weight=task.priority_weight_totalself.try_number=0self.max_tries=self.task.retriesself.unixname=getpass.getuser()self.run_as_user=task.run_as_userifstate:self.state=stateself.hostname=''self.executor_config=task.executor_configself.init_on_load()# Is this TaskInstance being currently running within `airflow run --raw`.# Not persisted to the database so only valid for the current processself.raw=False

[docs]@reconstructordefinit_on_load(self):""" Initialize the attributes that aren't stored in the DB. """self.test_mode=False# can be changed when calling 'run'

@propertydeftry_number(self):""" Return the try number that this task number will be when it is actually run. If the TI is currently running, this will match the column in the databse, in all othercases this will be incremenetd """# This is designed so that task logs end up in the right file.ifself.state==State.RUNNING:returnself._try_numberreturnself._try_number+1@try_number.setterdeftry_number(self,value):self._try_number=value@propertydefnext_try_number(self):returnself._try_number+1

[docs]defcommand(self,mark_success=False,ignore_all_deps=False,ignore_depends_on_past=False,ignore_task_deps=False,ignore_ti_state=False,local=False,pickle_id=None,raw=False,job_id=None,pool=None,cfg_path=None):""" Returns a command that can be executed anywhere where airflow is installed. This command is part of the message sent to executors by the orchestrator. """return" ".join(self.command_as_list(mark_success=mark_success,ignore_all_deps=ignore_all_deps,ignore_depends_on_past=ignore_depends_on_past,ignore_task_deps=ignore_task_deps,ignore_ti_state=ignore_ti_state,local=local,pickle_id=pickle_id,raw=raw,job_id=job_id,pool=pool,cfg_path=cfg_path))

[docs]defcommand_as_list(self,mark_success=False,ignore_all_deps=False,ignore_task_deps=False,ignore_depends_on_past=False,ignore_ti_state=False,local=False,pickle_id=None,raw=False,job_id=None,pool=None,cfg_path=None):""" Returns a command that can be executed anywhere where airflow is installed. This command is part of the message sent to executors by the orchestrator. """dag=self.task.dagshould_pass_filepath=notpickle_idanddagifshould_pass_filepathanddag.full_filepath!=dag.filepath:path="DAGS_FOLDER/{}".format(dag.filepath)elifshould_pass_filepathanddag.full_filepath:path=dag.full_filepathelse:path=NonereturnTaskInstance.generate_command(self.dag_id,self.task_id,self.execution_date,mark_success=mark_success,ignore_all_deps=ignore_all_deps,ignore_task_deps=ignore_task_deps,ignore_depends_on_past=ignore_depends_on_past,ignore_ti_state=ignore_ti_state,local=local,pickle_id=pickle_id,file_path=path,raw=raw,job_id=job_id,pool=pool,cfg_path=cfg_path)

[docs]@provide_sessiondefcurrent_state(self,session=None):""" Get the very latest state from the database, if a session is passed, we use and looking up the state becomes part of the session, otherwise a new session is used. """TI=TaskInstanceti=session.query(TI).filter(TI.dag_id==self.dag_id,TI.task_id==self.task_id,TI.execution_date==self.execution_date,).all()ifti:state=ti[0].stateelse:state=Nonereturnstate

[docs]@provide_sessiondeferror(self,session=None):""" Forces the task instance's state to FAILED in the database. """self.log.error("Recording the task instance as FAILED")self.state=State.FAILEDsession.merge(self)session.commit()

[docs]@provide_sessiondefrefresh_from_db(self,session=None,lock_for_update=False):""" Refreshes the task instance from the database based on the primary key :param lock_for_update: if True, indicates that the database should lock the TaskInstance (issuing a FOR UPDATE clause) until the session is committed. """TI=TaskInstanceqry=session.query(TI).filter(TI.dag_id==self.dag_id,TI.task_id==self.task_id,TI.execution_date==self.execution_date)iflock_for_update:ti=qry.with_for_update().first()else:ti=qry.first()ifti:self.state=ti.stateself.start_date=ti.start_dateself.end_date=ti.end_date# Get the raw value of try_number column, don't read through the# accessor here otherwise it will be incremeneted by one already.self.try_number=ti._try_numberself.max_tries=ti.max_triesself.hostname=ti.hostnameself.pid=ti.pidself.executor_config=ti.executor_configelse:self.state=None

[docs]@provide_sessiondefclear_xcom_data(self,session=None):""" Clears all XCom data from the database for the task instance """session.query(XCom).filter(XCom.dag_id==self.dag_id,XCom.task_id==self.task_id,XCom.execution_date==self.execution_date).delete()session.commit()

@propertydefkey(self):""" Returns a tuple that identifies the task instance uniquely """returnself.dag_id,self.task_id,self.execution_date,self.try_number@provide_sessiondefset_state(self,state,session=None):self.state=stateself.start_date=timezone.utcnow()self.end_date=timezone.utcnow()session.merge(self)session.commit()@propertydefis_premature(self):""" Returns whether a task is in UP_FOR_RETRY state and its retry interval has elapsed. """# is the task still in the retry waiting period?returnself.state==State.UP_FOR_RETRYandnotself.ready_for_retry()

[docs]@provide_sessiondefare_dependents_done(self,session=None):""" Checks whether the dependents of this task instance have all succeeded. This is meant to be used by wait_for_downstream. This is useful when you do not want to start processing the next schedule of a task until the dependents are done. For instance, if the task DROPs and recreates a table. """task=self.taskifnottask.downstream_task_ids:returnTrueti=session.query(func.count(TaskInstance.task_id)).filter(TaskInstance.dag_id==self.dag_id,TaskInstance.task_id.in_(task.downstream_task_ids),TaskInstance.execution_date==self.execution_date,TaskInstance.state==State.SUCCESS,)count=ti[0][0]returncount==len(task.downstream_task_ids)

@provide_sessiondef_get_previous_ti(self,session=None):dag=self.task.dagifdag:dr=self.get_dagrun(session=session)# LEGACY: most likely running from unit testsifnotdr:# Means that this TI is NOT being run from a DR, but from a catchupprevious_scheduled_date=dag.previous_schedule(self.execution_date)ifnotprevious_scheduled_date:returnNonereturnTaskInstance(task=self.task,execution_date=previous_scheduled_date)dr.dag=dagifdag.catchup:last_dagrun=dr.get_previous_scheduled_dagrun(session=session)else:last_dagrun=dr.get_previous_dagrun(session=session)iflast_dagrun:returnlast_dagrun.get_task_instance(self.task_id,session=session)returnNone@propertydefprevious_ti(self):"""The task instance for the task that ran before this task instance."""returnself._get_previous_ti()

[docs]@provide_sessiondefare_dependencies_met(self,dep_context=None,session=None,verbose=False):""" Returns whether or not all the conditions are met for this task instance to be run given the context for the dependencies (e.g. a task instance being force run from the UI will ignore some dependencies). :param dep_context: The execution context that determines the dependencies that should be evaluated. :type dep_context: DepContext :param session: database session :type session: sqlalchemy.orm.session.Session :param verbose: whether log details on failed dependencies on info or debug log level :type verbose: bool """dep_context=dep_contextorDepContext()failed=Falseverbose_aware_logger=self.log.infoifverboseelseself.log.debugfordep_statusinself.get_failed_dep_statuses(dep_context=dep_context,session=session):failed=Trueverbose_aware_logger("Dependencies not met for %s, dependency '%s' FAILED: %s",self,dep_status.dep_name,dep_status.reason)iffailed:returnFalseverbose_aware_logger("Dependencies all met for %s",self)returnTrue

[docs]defnext_retry_datetime(self):""" Get datetime of the next retry if the task instance fails. For exponential backoff, retry_delay is used as base and will be converted to seconds. """delay=self.task.retry_delayifself.task.retry_exponential_backoff:min_backoff=int(delay.total_seconds()*(2**(self.try_number-2)))# deterministic per task instancehash=int(hashlib.sha1("{}#{}#{}#{}".format(self.dag_id,self.task_id,self.execution_date,self.try_number).encode('utf-8')).hexdigest(),16)# between 0.5 * delay * (2^retry_number) and 1.0 * delay * (2^retry_number)modded_hash=min_backoff+hash%min_backoff# timedelta has a maximum representable value. The exponentiation# here means this value can be exceeded after a certain number# of tries (around 50 if the initial delay is 1s, even fewer if# the delay is larger). Cap the value here before creating a# timedelta object so the operation doesn't fail.delay_backoff_in_seconds=min(modded_hash,timedelta.max.total_seconds()-1)delay=timedelta(seconds=delay_backoff_in_seconds)ifself.task.max_retry_delay:delay=min(self.task.max_retry_delay,delay)returnself.end_date+delay

[docs]defready_for_retry(self):""" Checks on whether the task instance is in the right state and timeframe to be retried. """return(self.state==State.UP_FOR_RETRYandself.next_retry_datetime()<timezone.utcnow())

[docs]@provide_sessiondefpool_full(self,session):""" Returns a boolean as to whether the slot pool has room for this task to run """ifnotself.task.pool:returnFalsepool=(session.query(Pool).filter(Pool.pool==self.task.pool).first())ifnotpool:returnFalseopen_slots=pool.open_slots(session=session)returnopen_slots<=0

@provide_sessiondef_check_and_change_state_before_execution(self,verbose=True,ignore_all_deps=False,ignore_depends_on_past=False,ignore_task_deps=False,ignore_ti_state=False,mark_success=False,test_mode=False,job_id=None,pool=None,session=None):""" Checks dependencies and then sets state to RUNNING if they are met. Returns True if and only if state is set to RUNNING, which implies that task should be executed, in preparation for _run_raw_task :param verbose: whether to turn on more verbose logging :type verbose: bool :param ignore_all_deps: Ignore all of the non-critical dependencies, just runs :type ignore_all_deps: bool :param ignore_depends_on_past: Ignore depends_on_past DAG attribute :type ignore_depends_on_past: bool :param ignore_task_deps: Don't check the dependencies of this TI's task :type ignore_task_deps: bool :param ignore_ti_state: Disregards previous task instance state :type ignore_ti_state: bool :param mark_success: Don't run the task, mark its state as success :type mark_success: bool :param test_mode: Doesn't record success or failure in the DB :type test_mode: bool :param pool: specifies the pool to use to run the task instance :type pool: str :return: whether the state was changed to running or not :rtype: bool """task=self.taskself.pool=poolortask.poolself.test_mode=test_modeself.refresh_from_db(session=session,lock_for_update=True)self.job_id=job_idself.hostname=get_hostname()self.operator=task.__class__.__name__ifnotignore_all_depsandnotignore_ti_stateandself.state==State.SUCCESS:Stats.incr('previously_succeeded',1,1)queue_dep_context=DepContext(deps=QUEUE_DEPS,ignore_all_deps=ignore_all_deps,ignore_ti_state=ignore_ti_state,ignore_depends_on_past=ignore_depends_on_past,ignore_task_deps=ignore_task_deps)ifnotself.are_dependencies_met(dep_context=queue_dep_context,session=session,verbose=True):session.commit()returnFalse# TODO: Logging needs cleanup, not clear what is being printedhr="\n"+("-"*80)# Line break# For reporting purposes, we report based on 1-indexed,# not 0-indexed lists (i.e. Attempt 1 instead of# Attempt 0 for the first attempt).# Set the task start date. In case it was re-scheduled use the initial# start date that is recorded in task_reschedule tableself.start_date=timezone.utcnow()task_reschedules=TaskReschedule.find_for_task_instance(self,session)iftask_reschedules:self.start_date=task_reschedules[0].start_datedep_context=DepContext(deps=RUN_DEPS-QUEUE_DEPS,ignore_all_deps=ignore_all_deps,ignore_depends_on_past=ignore_depends_on_past,ignore_task_deps=ignore_task_deps,ignore_ti_state=ignore_ti_state)runnable=self.are_dependencies_met(dep_context=dep_context,session=session,verbose=True)ifnotrunnableandnotmark_success:# FIXME: we might have hit concurrency limits, which means we probably# have been running prematurely. This should be handled in the# scheduling mechanism.self.state=State.NONEself.log.warning(hr)self.log.warning("FIXME: Rescheduling due to concurrency limits reached at task runtime. Attempt %s of ""%s. State set to NONE.",self.try_number,self.max_tries+1)self.log.warning(hr)self.queued_dttm=timezone.utcnow()self.log.info("Queuing into pool %s",self.pool)session.merge(self)session.commit()returnFalse# Another worker might have started running this task instance while# the current worker process was blocked on refresh_from_dbifself.state==State.RUNNING:self.log.warning("Task Instance already running %s",self)session.commit()returnFalse# print status messageself.log.info(hr)self.log.info("Starting attempt %s of %s",self.try_number,self.max_tries+1)self.log.info(hr)self._try_number+=1ifnottest_mode:session.add(Log(State.RUNNING,self))self.state=State.RUNNINGself.pid=os.getpid()self.end_date=Noneifnottest_mode:session.merge(self)session.commit()# Closing all pooled connections to prevent# "max number of connections reached"settings.engine.dispose()ifverbose:ifmark_success:self.log.info("Marking success for %s on %s",self.task,self.execution_date)else:self.log.info("Executing %s on %s",self.task,self.execution_date)returnTrue@provide_sessiondef_run_raw_task(self,mark_success=False,test_mode=False,job_id=None,pool=None,session=None):""" Immediately runs the task (without checking or changing db state before execution) and then sets the appropriate final state after completion and runs any post-execute callbacks. Meant to be called only after another function changes the state to running. :param mark_success: Don't run the task, mark its state as success :type mark_success: bool :param test_mode: Doesn't record success or failure in the DB :type test_mode: bool :param pool: specifies the pool to use to run the task instance :type pool: str """task=self.taskself.pool=poolortask.poolself.test_mode=test_modeself.refresh_from_db(session=session)self.job_id=job_idself.hostname=get_hostname()self.operator=task.__class__.__name__context={}actual_start_date=timezone.utcnow()try:ifnotmark_success:context=self.get_template_context()task_copy=copy.copy(task)self.task=task_copydefsignal_handler(signum,frame):self.log.error("Received SIGTERM. Terminating subprocesses.")task_copy.on_kill()raiseAirflowException("Task received SIGTERM signal")signal.signal(signal.SIGTERM,signal_handler)# Don't clear Xcom until the task is certain to executeself.clear_xcom_data()start_time=time.time()self.render_templates()task_copy.pre_execute(context=context)# If a timeout is specified for the task, make it fail# if it goes beyondresult=Noneiftask_copy.execution_timeout:try:withtimeout(int(task_copy.execution_timeout.total_seconds())):result=task_copy.execute(context=context)exceptAirflowTaskTimeout:task_copy.on_kill()raiseelse:result=task_copy.execute(context=context)# If the task returns a result, push an XCom containing itiftask_copy.do_xcom_pushandresultisnotNone:self.xcom_push(key=XCOM_RETURN_KEY,value=result)task_copy.post_execute(context=context,result=result)end_time=time.time()duration=end_time-start_timeStats.timing('dag.{dag_id}.{task_id}.duration'.format(dag_id=task_copy.dag_id,task_id=task_copy.task_id),duration)Stats.incr('operator_successes_{}'.format(self.task.__class__.__name__),1,1)Stats.incr('ti_successes')self.refresh_from_db(lock_for_update=True)self.state=State.SUCCESSexceptAirflowSkipException:self.refresh_from_db(lock_for_update=True)self.state=State.SKIPPEDexceptAirflowRescheduleExceptionasreschedule_exception:self.refresh_from_db()self._handle_reschedule(actual_start_date,reschedule_exception,test_mode,context)returnexceptAirflowExceptionase:self.refresh_from_db()# for case when task is marked as success/failed externally# current behavior doesn't hit the success callbackifself.statein{State.SUCCESS,State.FAILED}:returnelse:self.handle_failure(e,test_mode,context)raiseexcept(Exception,KeyboardInterrupt)ase:self.handle_failure(e,test_mode,context)raise# Success callbacktry:iftask.on_success_callback:task.on_success_callback(context)exceptExceptionase3:self.log.error("Failed when executing success callback")self.log.exception(e3)# Recording SUCCESSself.end_date=timezone.utcnow()self.set_duration()ifnottest_mode:session.add(Log(self.state,self))session.merge(self)session.commit()@provide_sessiondefrun(self,verbose=True,ignore_all_deps=False,ignore_depends_on_past=False,ignore_task_deps=False,ignore_ti_state=False,mark_success=False,test_mode=False,job_id=None,pool=None,session=None):res=self._check_and_change_state_before_execution(verbose=verbose,ignore_all_deps=ignore_all_deps,ignore_depends_on_past=ignore_depends_on_past,ignore_task_deps=ignore_task_deps,ignore_ti_state=ignore_ti_state,mark_success=mark_success,test_mode=test_mode,job_id=job_id,pool=pool,session=session)ifres:self._run_raw_task(mark_success=mark_success,test_mode=test_mode,job_id=job_id,pool=pool,session=session)defdry_run(self):task=self.tasktask_copy=copy.copy(task)self.task=task_copyself.render_templates()task_copy.dry_run()@provide_sessiondef_handle_reschedule(self,actual_start_date,reschedule_exception,test_mode=False,context=None,session=None):# Don't record reschedule request in test modeiftest_mode:returnself.end_date=timezone.utcnow()self.set_duration()# Log reschedule requestsession.add(TaskReschedule(self.task,self.execution_date,self._try_number,actual_start_date,self.end_date,reschedule_exception.reschedule_date))# set stateself.state=State.UP_FOR_RESCHEDULE# Decrement try_number so subsequent runs will use the same try number and write# to same log file.self._try_number-=1session.merge(self)session.commit()self.log.info('Rescheduling task, marking task as UP_FOR_RESCHEDULE')@provide_sessiondefhandle_failure(self,error,test_mode=False,context=None,session=None):self.log.exception(error)task=self.taskself.end_date=timezone.utcnow()self.set_duration()Stats.incr('operator_failures_{}'.format(task.__class__.__name__),1,1)Stats.incr('ti_failures')ifnottest_mode:session.add(Log(State.FAILED,self))# Log failure durationsession.add(TaskFail(task,self.execution_date,self.start_date,self.end_date))ifcontextisnotNone:context['exception']=error# Let's go deepertry:# Since this function is called only when the TI state is running,# try_number contains the current try_number (not the next). We# only mark task instance as FAILED if the next task instance# try_number exceeds the max_tries.ifself.is_eligible_to_retry():self.state=State.UP_FOR_RETRYself.log.info('Marking task as UP_FOR_RETRY')iftask.email_on_retryandtask.email:self.email_alert(error)else:self.state=State.FAILEDiftask.retries:self.log.info('All retries failed; marking task as FAILED')else:self.log.info('Marking task as FAILED.')iftask.email_on_failureandtask.email:self.email_alert(error)exceptExceptionase2:self.log.error('Failed to send email to: %s',task.email)self.log.exception(e2)# Handling callbacks pessimisticallytry:ifself.state==State.UP_FOR_RETRYandtask.on_retry_callback:task.on_retry_callback(context)ifself.state==State.FAILEDandtask.on_failure_callback:task.on_failure_callback(context)exceptExceptionase3:self.log.error("Failed at executing callback")self.log.exception(e3)ifnottest_mode:session.merge(self)session.commit()

@provide_sessiondefget_template_context(self,session=None):task=self.taskfromairflowimportmacrostables=Noneif'tables'intask.params:tables=task.params['tables']params={}run_id=''dag_run=Noneifhasattr(task,'dag'):iftask.dag.params:params.update(task.dag.params)dag_run=(session.query(DagRun).filter_by(dag_id=task.dag.dag_id,execution_date=self.execution_date).first())run_id=dag_run.run_idifdag_runelseNonesession.expunge_all()session.commit()ds=self.execution_date.strftime('%Y-%m-%d')ts=self.execution_date.isoformat()yesterday_ds=(self.execution_date-timedelta(1)).strftime('%Y-%m-%d')tomorrow_ds=(self.execution_date+timedelta(1)).strftime('%Y-%m-%d')# For manually triggered dagruns that aren't run on a schedule, next/previous# schedule dates don't make sense, and should be set to execution date for# consistency with how execution_date is set for manually triggered tasks, i.e.# triggered_date == execution_date.ifdag_runanddag_run.external_trigger:prev_execution_date=self.execution_datenext_execution_date=self.execution_dateelse:prev_execution_date=task.dag.previous_schedule(self.execution_date)next_execution_date=task.dag.following_schedule(self.execution_date)next_ds=Nonenext_ds_nodash=Noneifnext_execution_date:next_ds=next_execution_date.strftime('%Y-%m-%d')next_ds_nodash=next_ds.replace('-','')prev_ds=Noneprev_ds_nodash=Noneifprev_execution_date:prev_ds=prev_execution_date.strftime('%Y-%m-%d')prev_ds_nodash=prev_ds.replace('-','')ds_nodash=ds.replace('-','')ts_nodash=self.execution_date.strftime('%Y%m%dT%H%M%S')ts_nodash_with_tz=ts.replace('-','').replace(':','')yesterday_ds_nodash=yesterday_ds.replace('-','')tomorrow_ds_nodash=tomorrow_ds.replace('-','')ti_key_str="{task.dag_id}__{task.task_id}__{ds_nodash}"ti_key_str=ti_key_str.format(**locals())iftask.params:params.update(task.params)ifconfiguration.getboolean('core','dag_run_conf_overrides_params'):self.overwrite_params_with_dag_run_conf(params=params,dag_run=dag_run)classVariableAccessor:""" Wrapper around Variable. This way you can get variables in templates by using {var.value.your_variable_name}. """def__init__(self):self.var=Nonedef__getattr__(self,item):self.var=Variable.get(item)returnself.vardef__repr__(self):returnstr(self.var)classVariableJsonAccessor:""" Wrapper around deserialized Variables. This way you can get variables in templates by using {var.json.your_variable_name}. """def__init__(self):self.var=Nonedef__getattr__(self,item):self.var=Variable.get(item,deserialize_json=True)returnself.vardef__repr__(self):returnstr(self.var)return{'dag':task.dag,'ds':ds,'next_ds':next_ds,'next_ds_nodash':next_ds_nodash,'prev_ds':prev_ds,'prev_ds_nodash':prev_ds_nodash,'ds_nodash':ds_nodash,'ts':ts,'ts_nodash':ts_nodash,'ts_nodash_with_tz':ts_nodash_with_tz,'yesterday_ds':yesterday_ds,'yesterday_ds_nodash':yesterday_ds_nodash,'tomorrow_ds':tomorrow_ds,'tomorrow_ds_nodash':tomorrow_ds_nodash,'END_DATE':ds,'end_date':ds,'dag_run':dag_run,'run_id':run_id,'execution_date':self.execution_date,'prev_execution_date':prev_execution_date,'next_execution_date':next_execution_date,'latest_date':ds,'macros':macros,'params':params,'tables':tables,'task':task,'task_instance':self,'ti':self,'task_instance_key_str':ti_key_str,'conf':configuration,'test_mode':self.test_mode,'var':{'value':VariableAccessor(),'json':VariableJsonAccessor()},'inlets':task.inlets,'outlets':task.outlets,}defoverwrite_params_with_dag_run_conf(self,params,dag_run):ifdag_runanddag_run.conf:params.update(dag_run.conf)defrender_templates(self):task=self.taskjinja_context=self.get_template_context()ifhasattr(self,'task')andhasattr(self.task,'dag'):ifself.task.dag.user_defined_macros:jinja_context.update(self.task.dag.user_defined_macros)rt=self.task.render_template# shortcut to methodforattrintask.__class__.template_fields:content=getattr(task,attr)ifcontent:rendered_content=rt(attr,content,jinja_context)setattr(task,attr,rendered_content)defemail_alert(self,exception):exception_html=str(exception).replace('\n','<br>')jinja_context=self.get_template_context()# This function is called after changing the state# from State.RUNNING so need to subtract 1 from self.try_number.jinja_context.update(dict(exception=exception,exception_html=exception_html,try_number=self.try_number-1,max_tries=self.max_tries))jinja_env=self.task.get_template_env()default_subject='Airflow alert: {{ti}}'# For reporting purposes, we report based on 1-indexed,# not 0-indexed lists (i.e. Try 1 instead of# Try 0 for the first attempt).default_html_content=('Try {{try_number}} out of {{max_tries + 1}}<br>''Exception:<br>{{exception_html}}<br>''Log: <a href="{{ti.log_url}}">Link</a><br>''Host: {{ti.hostname}}<br>''Log file: {{ti.log_filepath}}<br>''Mark success: <a href="{{ti.mark_success_url}}">Link</a><br>')defrender(key,content):ifconfiguration.has_option('email',key):path=configuration.get('email',key)withopen(path)asf:content=f.read()returnjinja_env.from_string(content).render(**jinja_context)subject=render('subject_template',default_subject)html_content=render('html_content_template',default_html_content)send_email(self.task.email,subject,html_content)defset_duration(self):ifself.end_dateandself.start_date:self.duration=(self.end_date-self.start_date).total_seconds()else:self.duration=None

[docs]defxcom_push(self,key,value,execution_date=None):""" Make an XCom available for tasks to pull. :param key: A key for the XCom :type key: str :param value: A value for the XCom. The value is pickled and stored in the database. :type value: any pickleable object :param execution_date: if provided, the XCom will not be visible until this date. This can be used, for example, to send a message to a task on a future date without it being immediately visible. :type execution_date: datetime """ifexecution_dateandexecution_date<self.execution_date:raiseValueError('execution_date can not be in the past (current ''execution_date is {}; received {})'.format(self.execution_date,execution_date))XCom.set(key=key,value=value,task_id=self.task_id,dag_id=self.dag_id,execution_date=execution_dateorself.execution_date)

[docs]defxcom_pull(self,task_ids=None,dag_id=None,key=XCOM_RETURN_KEY,include_prior_dates=False):""" Pull XComs that optionally meet certain criteria. The default value for `key` limits the search to XComs that were returned by other tasks (as opposed to those that were pushed manually). To remove this filter, pass key=None (or any desired value). If a single task_id string is provided, the result is the value of the most recent matching XCom from that task_id. If multiple task_ids are provided, a tuple of matching values is returned. None is returned whenever no matches are found. :param key: A key for the XCom. If provided, only XComs with matching keys will be returned. The default key is 'return_value', also available as a constant XCOM_RETURN_KEY. This key is automatically given to XComs returned by tasks (as opposed to being pushed manually). To remove the filter, pass key=None. :type key: str :param task_ids: Only XComs from tasks with matching ids will be pulled. Can pass None to remove the filter. :type task_ids: str or iterable of strings (representing task_ids) :param dag_id: If provided, only pulls XComs from this DAG. If None (default), the DAG of the calling task is used. :type dag_id: str :param include_prior_dates: If False, only XComs from the current execution_date are returned. If True, XComs from previous dates are returned as well. :type include_prior_dates: bool """ifdag_idisNone:dag_id=self.dag_idpull_fn=functools.partial(XCom.get_one,execution_date=self.execution_date,key=key,dag_id=dag_id,include_prior_dates=include_prior_dates)ifis_container(task_ids):returntuple(pull_fn(task_id=t)fortintask_ids)else:returnpull_fn(task_id=task_ids)

[docs]@functools.total_orderingclassBaseOperator(LoggingMixin):""" Abstract base class for all operators. Since operators create objects that become nodes in the dag, BaseOperator contains many recursive methods for dag crawling behavior. To derive this class, you are expected to override the constructor as well as the 'execute' method. Operators derived from this class should perform or trigger certain tasks synchronously (wait for completion). Example of operators could be an operator that runs a Pig job (PigOperator), a sensor operator that waits for a partition to land in Hive (HiveSensorOperator), or one that moves data from Hive to MySQL (Hive2MySqlOperator). Instances of these operators (tasks) target specific operations, running specific scripts, functions or data transfers. This class is abstract and shouldn't be instantiated. Instantiating a class derived from this one results in the creation of a task object, which ultimately becomes a node in DAG objects. Task dependencies should be set by using the set_upstream and/or set_downstream methods. :param task_id: a unique, meaningful id for the task :type task_id: str :param owner: the owner of the task, using the unix username is recommended :type owner: str :param retries: the number of retries that should be performed before failing the task :type retries: int :param retry_delay: delay between retries :type retry_delay: datetime.timedelta :param retry_exponential_backoff: allow progressive longer waits between retries by using exponential backoff algorithm on retry delay (delay will be converted into seconds) :type retry_exponential_backoff: bool :param max_retry_delay: maximum delay interval between retries :type max_retry_delay: datetime.timedelta :param start_date: The ``start_date`` for the task, determines the ``execution_date`` for the first task instance. The best practice is to have the start_date rounded to your DAG's ``schedule_interval``. Daily jobs have their start_date some day at 00:00:00, hourly jobs have their start_date at 00:00 of a specific hour. Note that Airflow simply looks at the latest ``execution_date`` and adds the ``schedule_interval`` to determine the next ``execution_date``. It is also very important to note that different tasks' dependencies need to line up in time. If task A depends on task B and their start_date are offset in a way that their execution_date don't line up, A's dependencies will never be met. If you are looking to delay a task, for example running a daily task at 2AM, look into the ``TimeSensor`` and ``TimeDeltaSensor``. We advise against using dynamic ``start_date`` and recommend using fixed ones. Read the FAQ entry about start_date for more information. :type start_date: datetime.datetime :param end_date: if specified, the scheduler won't go beyond this date :type end_date: datetime.datetime :param depends_on_past: when set to true, task instances will run sequentially while relying on the previous task's schedule to succeed. The task instance for the start_date is allowed to run. :type depends_on_past: bool :param wait_for_downstream: when set to true, an instance of task X will wait for tasks immediately downstream of the previous instance of task X to finish successfully before it runs. This is useful if the different instances of a task X alter the same asset, and this asset is used by tasks downstream of task X. Note that depends_on_past is forced to True wherever wait_for_downstream is used. :type wait_for_downstream: bool :param queue: which queue to target when running this job. Not all executors implement queue management, the CeleryExecutor does support targeting specific queues. :type queue: str :param dag: a reference to the dag the task is attached to (if any) :type dag: airflow.models.DAG :param priority_weight: priority weight of this task against other task. This allows the executor to trigger higher priority tasks before others when things get backed up. Set priority_weight as a higher number for more important tasks. :type priority_weight: int :param weight_rule: weighting method used for the effective total priority weight of the task. Options are: ``{ downstream | upstream | absolute }`` default is ``downstream`` When set to ``downstream`` the effective weight of the task is the aggregate sum of all downstream descendants. As a result, upstream tasks will have higher weight and will be scheduled more aggressively when using positive weight values. This is useful when you have multiple dag run instances and desire to have all upstream tasks to complete for all runs before each dag can continue processing downstream tasks. When set to ``upstream`` the effective weight is the aggregate sum of all upstream ancestors. This is the opposite where downtream tasks have higher weight and will be scheduled more aggressively when using positive weight values. This is useful when you have multiple dag run instances and prefer to have each dag complete before starting upstream tasks of other dags. When set to ``absolute``, the effective weight is the exact ``priority_weight`` specified without additional weighting. You may want to do this when you know exactly what priority weight each task should have. Additionally, when set to ``absolute``, there is bonus effect of significantly speeding up the task creation process as for very large DAGS. Options can be set as string or using the constants defined in the static class ``airflow.utils.WeightRule`` :type weight_rule: str :param pool: the slot pool this task should run in, slot pools are a way to limit concurrency for certain tasks :type pool: str :param sla: time by which the job is expected to succeed. Note that this represents the ``timedelta`` after the period is closed. For example if you set an SLA of 1 hour, the scheduler would send an email soon after 1:00AM on the ``2016-01-02`` if the ``2016-01-01`` instance has not succeeded yet. The scheduler pays special attention for jobs with an SLA and sends alert emails for sla misses. SLA misses are also recorded in the database for future reference. All tasks that share the same SLA time get bundled in a single email, sent soon after that time. SLA notification are sent once and only once for each task instance. :type sla: datetime.timedelta :param execution_timeout: max time allowed for the execution of this task instance, if it goes beyond it will raise and fail. :type execution_timeout: datetime.timedelta :param on_failure_callback: a function to be called when a task instance of this task fails. a context dictionary is passed as a single parameter to this function. Context contains references to related objects to the task instance and is documented under the macros section of the API. :type on_failure_callback: callable :param on_retry_callback: much like the ``on_failure_callback`` except that it is executed when retries occur. :type on_retry_callback: callable :param on_success_callback: much like the ``on_failure_callback`` except that it is executed when the task succeeds. :type on_success_callback: callable :param trigger_rule: defines the rule by which dependencies are applied for the task to get triggered. Options are: ``{ all_success | all_failed | all_done | one_success | one_failed | none_failed | dummy}`` default is ``all_success``. Options can be set as string or using the constants defined in the static class ``airflow.utils.TriggerRule`` :type trigger_rule: str :param resources: A map of resource parameter names (the argument names of the Resources constructor) to their values. :type resources: dict :param run_as_user: unix username to impersonate while running the task :type run_as_user: str :param task_concurrency: When set, a task will be able to limit the concurrent runs across execution_dates :type task_concurrency: int :param executor_config: Additional task-level configuration parameters that are interpreted by a specific executor. Parameters are namespaced by the name of executor. **Example**: to run this task in a specific docker container through the KubernetesExecutor :: MyOperator(..., executor_config={ "KubernetesExecutor": {"image": "myCustomDockerImage"} } ) :type executor_config: dict :param do_xcom_push: if True, an XCom is pushed containing the Operator's result :type do_xcom_push: bool """# For derived classes to define which fields will get jinjaifiedtemplate_fields=[]# type: Iterable[str]# Defines which files extensions to look for in the templated fieldstemplate_ext=[]# type: Iterable[str]# Defines the color in the UIui_color='#fff'ui_fgcolor='#000'# base list which includes all the attrs that don't need deep copy._base_operator_shallow_copy_attrs=('user_defined_macros','user_defined_filters','params','_log',)# each operator should override this class attr for shallow copy attrs.shallow_copy_attrs=()# type: Iterable[str]@apply_defaultsdef__init__(self,task_id,owner=configuration.conf.get('operators','DEFAULT_OWNER'),email=None,email_on_retry=True,email_on_failure=True,retries=0,retry_delay=timedelta(seconds=300),retry_exponential_backoff=False,max_retry_delay=None,start_date=None,end_date=None,schedule_interval=None,# not hooked as of nowdepends_on_past=False,wait_for_downstream=False,dag=None,params=None,default_args=None,adhoc=False,priority_weight=1,weight_rule=WeightRule.DOWNSTREAM,queue=configuration.conf.get('celery','default_queue'),pool=None,sla=None,execution_timeout=None,on_failure_callback=None,on_success_callback=None,on_retry_callback=None,trigger_rule=TriggerRule.ALL_SUCCESS,resources=None,run_as_user=None,task_concurrency=None,executor_config=None,do_xcom_push=True,inlets=None,outlets=None,*args,**kwargs):ifargsorkwargs:# TODO remove *args and **kwargs in Airflow 2.0warnings.warn('Invalid arguments were passed to {c} (task_id: {t}). ''Support for passing such arguments will be dropped in ''Airflow 2.0. Invalid arguments were:''\n*args: {a}\n**kwargs: {k}'.format(c=self.__class__.__name__,a=args,k=kwargs,t=task_id),category=PendingDeprecationWarning,stacklevel=3)validate_key(task_id)self.task_id=task_idself.owner=ownerself.email=emailself.email_on_retry=email_on_retryself.email_on_failure=email_on_failureself.start_date=start_dateifstart_dateandnotisinstance(start_date,datetime):self.log.warning("start_date for %s isn't datetime.datetime",self)elifstart_date:self.start_date=timezone.convert_to_utc(start_date)self.end_date=end_dateifend_date:self.end_date=timezone.convert_to_utc(end_date)ifnotTriggerRule.is_valid(trigger_rule):raiseAirflowException("The trigger_rule must be one of {all_triggers},""'{d}.{t}'; received '{tr}'.".format(all_triggers=TriggerRule.all_triggers(),d=dag.dag_idifdagelse"",t=task_id,tr=trigger_rule))self.trigger_rule=trigger_ruleself.depends_on_past=depends_on_pastself.wait_for_downstream=wait_for_downstreamifwait_for_downstream:self.depends_on_past=Trueifschedule_interval:self.log.warning("schedule_interval is used for %s, though it has ""been deprecated as a task parameter, you need to ""specify it as a DAG parameter instead",self)self._schedule_interval=schedule_intervalself.retries=retriesself.queue=queueself.pool=poolself.sla=slaself.execution_timeout=execution_timeoutself.on_failure_callback=on_failure_callbackself.on_success_callback=on_success_callbackself.on_retry_callback=on_retry_callbackifisinstance(retry_delay,timedelta):self.retry_delay=retry_delayelse:self.log.debug("Retry_delay isn't timedelta object, assuming secs")self.retry_delay=timedelta(seconds=retry_delay)self.retry_exponential_backoff=retry_exponential_backoffself.max_retry_delay=max_retry_delayself.params=paramsor{}# Available in templates!self.adhoc=adhocself.priority_weight=priority_weightifnotWeightRule.is_valid(weight_rule):raiseAirflowException("The weight_rule must be one of {all_weight_rules},""'{d}.{t}'; received '{tr}'.".format(all_weight_rules=WeightRule.all_weight_rules,d=dag.dag_idifdagelse"",t=task_id,tr=weight_rule))self.weight_rule=weight_ruleself.resources=Resources(**(resourcesor{}))self.run_as_user=run_as_userself.task_concurrency=task_concurrencyself.executor_config=executor_configor{}self.do_xcom_push=do_xcom_push# Private attributesself._upstream_task_ids=set()self._downstream_task_ids=set()ifnotdagand_CONTEXT_MANAGER_DAG:dag=_CONTEXT_MANAGER_DAGifdag:self.dag=dagself._log=logging.getLogger("airflow.task.operators")# lineageself.inlets=[]self.outlets=[]self.lineage_data=Noneself._inlets={"auto":False,"task_ids":[],"datasets":[],}self._outlets={"datasets":[],}ifinlets:self._inlets.update(inlets)ifoutlets:self._outlets.update(outlets)self._comps={'task_id','dag_id','owner','email','email_on_retry','retry_delay','retry_exponential_backoff','max_retry_delay','start_date','schedule_interval','depends_on_past','wait_for_downstream','adhoc','priority_weight','sla','execution_timeout','on_failure_callback','on_success_callback','on_retry_callback','do_xcom_push',}def__eq__(self,other):if(type(self)==type(other)andself.task_id==other.task_id):returnall(self.__dict__.get(c,None)==other.__dict__.get(c,None)forcinself._comps)returnFalsedef__ne__(self,other):returnnotself==otherdef__lt__(self,other):returnself.task_id<other.task_iddef__hash__(self):hash_components=[type(self)]forcinself._comps:val=getattr(self,c,None)try:hash(val)hash_components.append(val)exceptTypeError:hash_components.append(repr(val))returnhash(tuple(hash_components))# Composing Operators -----------------------------------------------def__rshift__(self,other):""" Implements Self >> Other == self.set_downstream(other) If "Other" is a DAG, the DAG is assigned to the Operator. """ifisinstance(other,DAG):# if this dag is already assigned, do nothing# otherwise, do normal dag assignmentifnot(self.has_dag()andself.dagisother):self.dag=otherelse:self.set_downstream(other)returnotherdef__lshift__(self,other):""" Implements Self << Other == self.set_upstream(other) If "Other" is a DAG, the DAG is assigned to the Operator. """ifisinstance(other,DAG):# if this dag is already assigned, do nothing# otherwise, do normal dag assignmentifnot(self.has_dag()andself.dagisother):self.dag=otherelse:self.set_upstream(other)returnotherdef__rrshift__(self,other):""" Called for [DAG] >> [Operator] because DAGs don't have __rshift__ operators. """self.__lshift__(other)returnselfdef__rlshift__(self,other):""" Called for [DAG] << [Operator] because DAGs don't have __lshift__ operators. """self.__rshift__(other)returnself# /Composing Operators ---------------------------------------------@propertydefdag(self):""" Returns the Operator's DAG if set, otherwise raises an error """ifself.has_dag():returnself._dagelse:raiseAirflowException('Operator {} has not been assigned to a DAG yet'.format(self))@dag.setterdefdag(self,dag):""" Operators can be assigned to one DAG, one time. Repeat assignments to that same DAG are ok. """ifnotisinstance(dag,DAG):raiseTypeError('Expected DAG; received {}'.format(dag.__class__.__name__))elifself.has_dag()andself.dagisnotdag:raiseAirflowException("The DAG assigned to {} can not be changed.".format(self))elifself.task_idnotindag.task_dict:dag.add_task(self)self._dag=dag

[docs]defhas_dag(self):""" Returns True if the Operator has been assigned to a DAG. """returngetattr(self,'_dag',None)isnotNone

@propertydefdag_id(self):ifself.has_dag():returnself.dag.dag_idelse:return'adhoc_'+self.owner@propertydefdeps(self):""" Returns the list of dependencies for the operator. These differ from execution context dependencies in that they are specific to tasks and can be extended/overridden by subclasses. """return{NotInRetryPeriodDep(),PrevDagrunDep(),TriggerRuleDep(),}@propertydefschedule_interval(self):""" The schedule interval of the DAG always wins over individual tasks so that tasks within a DAG always line up. The task still needs a schedule_interval as it may not be attached to a DAG. """ifself.has_dag():returnself.dag._schedule_intervalelse:returnself._schedule_interval@propertydefpriority_weight_total(self):ifself.weight_rule==WeightRule.ABSOLUTE:returnself.priority_weightelifself.weight_rule==WeightRule.DOWNSTREAM:upstream=Falseelifself.weight_rule==WeightRule.UPSTREAM:upstream=Trueelse:upstream=Falsereturnself.priority_weight+sum(map(lambdatask_id:self._dag.task_dict[task_id].priority_weight,self.get_flat_relative_ids(upstream=upstream)))

[docs]@prepare_lineagedefpre_execute(self,context):""" This hook is triggered right before self.execute() is called. """pass

[docs]defexecute(self,context):""" This is the main method to derive when creating an operator. Context is the same dictionary used as when rendering jinja templates. Refer to get_template_context for more context. """raiseNotImplementedError()

[docs]@apply_lineagedefpost_execute(self,context,result=None):""" This hook is triggered right after self.execute() is called. It is passed the execution context and any results returned by the operator. """pass

[docs]defon_kill(self):""" Override this method to cleanup subprocesses when a task instance gets killed. Any use of the threading, subprocess or multiprocessing module within an operator needs to be cleaned up or it will leave ghost processes behind. """pass

[docs]defrender_template_from_field(self,attr,content,context,jinja_env):""" Renders a template from a field. If the field is a string, it will simply render the string and return the result. If it is a collection or nested set of collections, it will traverse the structure and render all elements in it. If the field has another type, it will return it as it is. """rt=self.render_templateifisinstance(content,six.string_types):result=jinja_env.from_string(content).render(**context)elifisinstance(content,(list,tuple)):result=[rt(attr,e,context)foreincontent]elifisinstance(content,dict):result={k:rt("{}[{}]".format(attr,k),v,context)fork,vinlist(content.items())}else:result=contentreturnresult

[docs]defrender_template(self,attr,content,context):""" Renders a template either from a file or directly in a field, and returns the rendered result. """jinja_env=self.get_template_env()exts=self.__class__.template_extif(isinstance(content,six.string_types)andany([content.endswith(ext)forextinexts])):returnjinja_env.get_template(content).render(**context)else:returnself.render_template_from_field(attr,content,context,jinja_env)

[docs]defprepare_template(self):""" Hook that is triggered after the templated fields get replaced by their content. If you need your operator to alter the content of the file before the template is rendered, it should override this method to do so. """pass

[docs]@provide_sessiondefclear(self,start_date=None,end_date=None,upstream=False,downstream=False,session=None):""" Clears the state of task instances associated with the task, following the parameters specified. """TI=TaskInstanceqry=session.query(TI).filter(TI.dag_id==self.dag_id)ifstart_date:qry=qry.filter(TI.execution_date>=start_date)ifend_date:qry=qry.filter(TI.execution_date<=end_date)tasks=[self.task_id]ifupstream:tasks+=[t.task_idfortinself.get_flat_relatives(upstream=True)]ifdownstream:tasks+=[t.task_idfortinself.get_flat_relatives(upstream=False)]qry=qry.filter(TI.task_id.in_(tasks))count=qry.count()clear_task_instances(qry.all(),session,dag=self.dag)session.commit()returncount

[docs]@provide_sessiondefget_task_instances(self,start_date=None,end_date=None,session=None):""" Get a set of task instance related to this task for a specific date range. """end_date=end_dateortimezone.utcnow()returnsession.query(TaskInstance)\
.filter(TaskInstance.dag_id==self.dag_id)\
.filter(TaskInstance.task_id==self.task_id)\
.filter(TaskInstance.execution_date>=start_date)\
.filter(TaskInstance.execution_date<=end_date)\
.order_by(TaskInstance.execution_date)\
.all()

[docs]defget_flat_relative_ids(self,upstream=False,found_descendants=None):""" Get a flat list of relatives' ids, either upstream or downstream. """ifnotfound_descendants:found_descendants=set()relative_ids=self.get_direct_relative_ids(upstream)forrelative_idinrelative_ids:ifrelative_idnotinfound_descendants:found_descendants.add(relative_id)relative_task=self._dag.task_dict[relative_id]relative_task.get_flat_relative_ids(upstream,found_descendants)returnfound_descendants

[docs]defget_flat_relatives(self,upstream=False):""" Get a flat list of relatives, either upstream or downstream. """returnlist(map(lambdatask_id:self._dag.task_dict[task_id],self.get_flat_relative_ids(upstream)))

[docs]defrun(self,start_date=None,end_date=None,ignore_first_depends_on_past=False,ignore_ti_state=False,mark_success=False):""" Run a set of task instances for a date range. """start_date=start_dateorself.start_dateend_date=end_dateorself.end_dateortimezone.utcnow()fordtinself.dag.date_range(start_date,end_date=end_date):TaskInstance(self,dt).run(mark_success=mark_success,ignore_depends_on_past=(dt==start_dateandignore_first_depends_on_past),ignore_ti_state=ignore_ti_state)

[docs]defget_direct_relative_ids(self,upstream=False):""" Get the direct relative ids to the current task, upstream or downstream. """ifupstream:returnself._upstream_task_idselse:returnself._downstream_task_ids

[docs]defget_direct_relatives(self,upstream=False):""" Get the direct relatives to the current task, upstream or downstream. """ifupstream:returnself.upstream_listelse:returnself.downstream_list

def__repr__(self):return"<Task({self.__class__.__name__}): {self.task_id}>".format(self=self)@propertydeftask_type(self):returnself.__class__.__name__defadd_only_new(self,item_set,item):ifiteminitem_set:self.log.warning('Dependency {self}, {item} already registered'''.format(**locals()))else:item_set.add(item)def_set_relatives(self,task_or_task_list,upstream=False):try:task_list=list(task_or_task_list)exceptTypeError:task_list=[task_or_task_list]fortintask_list:ifnotisinstance(t,BaseOperator):raiseAirflowException("Relationships can only be set between ""Operators; received {}".format(t.__class__.__name__))# relationships can only be set if the tasks share a single DAG. Tasks# without a DAG are assigned to that DAG.dags={t._dag.dag_id:t._dagfortin[self]+task_listift.has_dag()}iflen(dags)>1:raiseAirflowException('Tried to set relationships between tasks in ''more than one DAG: {}'.format(dags.values()))eliflen(dags)==1:dag=dags.popitem()[1]else:raiseAirflowException("Tried to create relationships between tasks that don't have ""DAGs yet. Set the DAG for at least one ""task and try again: {}".format([self]+task_list))ifdagandnotself.has_dag():self.dag=dagfortaskintask_list:ifdagandnottask.has_dag():task.dag=dagifupstream:task.add_only_new(task._downstream_task_ids,self.task_id)self.add_only_new(self._upstream_task_ids,task.task_id)else:self.add_only_new(self._downstream_task_ids,task.task_id)task.add_only_new(task._upstream_task_ids,self.task_id)

[docs]defset_downstream(self,task_or_task_list):""" Set a task or a task list to be directly downstream from the current task. """self._set_relatives(task_or_task_list,upstream=False)

[docs]defset_upstream(self,task_or_task_list):""" Set a task or a task list to be directly upstream from the current task. """self._set_relatives(task_or_task_list,upstream=True)

[docs]classDagModel(Base):__tablename__="dag"""" These items are stored in the database for state related information """dag_id=Column(String(ID_LEN),primary_key=True)# A DAG can be paused from the UI / DB# Set this default value of is_paused based on a configuration value!is_paused_at_creation=configuration.conf\
.getboolean('core','dags_are_paused_at_creation')is_paused=Column(Boolean,default=is_paused_at_creation)# Whether the DAG is a subdagis_subdag=Column(Boolean,default=False)# Whether that DAG was seen on the last DagBag loadis_active=Column(Boolean,default=False)# Last time the scheduler startedlast_scheduler_run=Column(UtcDateTime)# Last time this DAG was pickledlast_pickled=Column(UtcDateTime)# Time when the DAG last received a refresh signal# (e.g. the DAG's "refresh" button was clicked in the web UI)last_expired=Column(UtcDateTime)# Whether (one of) the scheduler is scheduling this DAG at the momentscheduler_lock=Column(Boolean)# Foreign key to the latest pickle_idpickle_id=Column(Integer)# The location of the file containing the DAG objectfileloc=Column(String(2000))# String representing the ownersowners=Column(String(2000))# Description of the dagdescription=Column(Text)# Default view of the inside the webserverdefault_view=Column(String(25))# Schedule intervalschedule_interval=Column(Interval)def__repr__(self):return"<DAG: {self.dag_id}>".format(self=self)@propertydeftimezone(self):returnsettings.TIMEZONE@staticmethod@provide_sessiondefget_dagmodel(dag_id,session=None):returnsession.query(DagModel).filter(DagModel.dag_id==dag_id).first()@classmethod@provide_sessiondefget_current(cls,dag_id,session=None):returnsession.query(cls).filter(cls.dag_id==dag_id).first()defget_default_view(self):ifself.default_viewisNone:returnconfiguration.conf.get('webserver','dag_default_view').lower()else:returnself.default_view@provide_sessiondefget_last_dagrun(self,session=None,include_externally_triggered=False):returnget_last_dagrun(self.dag_id,session=session,include_externally_triggered=include_externally_triggered)@propertydefsafe_dag_id(self):returnself.dag_id.replace('.','__dot__')defget_dag(self):returnDagBag(dag_folder=self.fileloc).get_dag(self.dag_id)

[docs]@provide_sessiondefcreate_dagrun(self,run_id,state,execution_date,start_date=None,external_trigger=False,conf=None,session=None):""" Creates a dag run from this dag including the tasks associated with this dag. Returns the dag run. :param run_id: defines the the run id for this dag run :type run_id: str :param execution_date: the execution date of this dag run :type execution_date: datetime.datetime :param state: the state of the dag run :type state: airflow.utils.state.State :param start_date: the date this dag run should be evaluated :type start_date: datetime.datetime :param external_trigger: whether this dag run is externally triggered :type external_trigger: bool :param session: database session :type session: sqlalchemy.orm.session.Session """returnself.get_dag().create_dagrun(run_id=run_id,state=state,execution_date=execution_date,start_date=start_date,external_trigger=external_trigger,conf=conf,session=session)

[docs]@functools.total_orderingclassDAG(BaseDag,LoggingMixin):""" A dag (directed acyclic graph) is a collection of tasks with directional dependencies. A dag also has a schedule, a start date and an end date (optional). For each schedule, (say daily or hourly), the DAG needs to run each individual tasks as their dependencies are met. Certain tasks have the property of depending on their own past, meaning that they can't run until their previous schedule (and upstream tasks) are completed. DAGs essentially act as namespaces for tasks. A task_id can only be added once to a DAG. :param dag_id: The id of the DAG :type dag_id: str :param description: The description for the DAG to e.g. be shown on the webserver :type description: str :param schedule_interval: Defines how often that DAG runs, this timedelta object gets added to your latest task instance's execution_date to figure out the next schedule :type schedule_interval: datetime.timedelta or dateutil.relativedelta.relativedelta or str that acts as a cron expression :param start_date: The timestamp from which the scheduler will attempt to backfill :type start_date: datetime.datetime :param end_date: A date beyond which your DAG won't run, leave to None for open ended scheduling :type end_date: datetime.datetime :param template_searchpath: This list of folders (non relative) defines where jinja will look for your templates. Order matters. Note that jinja/airflow includes the path of your DAG file by default :type template_searchpath: str or list[str] :param user_defined_macros: a dictionary of macros that will be exposed in your jinja templates. For example, passing ``dict(foo='bar')`` to this argument allows you to ``{{ foo }}`` in all jinja templates related to this DAG. Note that you can pass any type of object here. :type user_defined_macros: dict :param user_defined_filters: a dictionary of filters that will be exposed in your jinja templates. For example, passing ``dict(hello=lambda name: 'Hello %s' % name)`` to this argument allows you to ``{{ 'world' | hello }}`` in all jinja templates related to this DAG. :type user_defined_filters: dict :param default_args: A dictionary of default parameters to be used as constructor keyword parameters when initialising operators. Note that operators have the same hook, and precede those defined here, meaning that if your dict contains `'depends_on_past': True` here and `'depends_on_past': False` in the operator's call `default_args`, the actual value will be `False`. :type default_args: dict :param params: a dictionary of DAG level parameters that are made accessible in templates, namespaced under `params`. These params can be overridden at the task level. :type params: dict :param concurrency: the number of task instances allowed to run concurrently :type concurrency: int :param max_active_runs: maximum number of active DAG runs, beyond this number of DAG runs in a running state, the scheduler won't create new active DAG runs :type max_active_runs: int :param dagrun_timeout: specify how long a DagRun should be up before timing out / failing, so that new DagRuns can be created. The timeout is only enforced for scheduled DagRuns, and only once the # of active DagRuns == max_active_runs. :type dagrun_timeout: datetime.timedelta :param sla_miss_callback: specify a function to call when reporting SLA timeouts. :type sla_miss_callback: types.FunctionType :param default_view: Specify DAG default view (tree, graph, duration, gantt, landing_times) :type default_view: str :param orientation: Specify DAG orientation in graph view (LR, TB, RL, BT) :type orientation: str :param catchup: Perform scheduler catchup (or only run latest)? Defaults to True :type catchup: bool :param on_failure_callback: A function to be called when a DagRun of this dag fails. A context dictionary is passed as a single parameter to this function. :type on_failure_callback: callable :param on_success_callback: Much like the ``on_failure_callback`` except that it is executed when the dag succeeds. :type on_success_callback: callable :param access_control: Specify optional DAG-level permissions, e.g., "{'role1': {'can_dag_read'}, 'role2': {'can_dag_read', 'can_dag_edit'}}" :type access_control: dict """def__init__(self,dag_id,description='',schedule_interval=timedelta(days=1),start_date=None,end_date=None,full_filepath=None,template_searchpath=None,user_defined_macros=None,user_defined_filters=None,default_args=None,concurrency=configuration.conf.getint('core','dag_concurrency'),max_active_runs=configuration.conf.getint('core','max_active_runs_per_dag'),dagrun_timeout=None,sla_miss_callback=None,default_view=None,orientation=configuration.conf.get('webserver','dag_orientation'),catchup=configuration.conf.getboolean('scheduler','catchup_by_default'),on_success_callback=None,on_failure_callback=None,doc_md=None,params=None,access_control=None):self.user_defined_macros=user_defined_macrosself.user_defined_filters=user_defined_filtersself.default_args=default_argsor{}self.params=paramsor{}# merging potentially conflicting default_args['params'] into paramsif'params'inself.default_args:self.params.update(self.default_args['params'])delself.default_args['params']validate_key(dag_id)# Properties from BaseDagself._dag_id=dag_idself._full_filepath=full_filepathiffull_filepathelse''self._concurrency=concurrencyself._pickle_id=Noneself._description=description# set file location to caller source pathself.fileloc=sys._getframe().f_back.f_code.co_filenameself.task_dict=dict()# set timezoneifstart_dateandstart_date.tzinfo:self.timezone=start_date.tzinfoelif'start_date'inself.default_argsandself.default_args['start_date']:ifisinstance(self.default_args['start_date'],six.string_types):self.default_args['start_date']=(timezone.parse(self.default_args['start_date']))self.timezone=self.default_args['start_date'].tzinfoifnothasattr(self,'timezone')ornotself.timezone:self.timezone=settings.TIMEZONEself.start_date=timezone.convert_to_utc(start_date)self.end_date=timezone.convert_to_utc(end_date)# also convert tasksif'start_date'inself.default_args:self.default_args['start_date']=(timezone.convert_to_utc(self.default_args['start_date']))if'end_date'inself.default_args:self.default_args['end_date']=(timezone.convert_to_utc(self.default_args['end_date']))self.schedule_interval=schedule_intervalifisinstance(schedule_interval,Hashable)andschedule_intervalincron_presets:self._schedule_interval=cron_presets.get(schedule_interval)elifschedule_interval=='@once':self._schedule_interval=Noneelse:self._schedule_interval=schedule_intervalifisinstance(template_searchpath,six.string_types):template_searchpath=[template_searchpath]self.template_searchpath=template_searchpathself.parent_dag=None# Gets set when DAGs are loadedself.last_loaded=timezone.utcnow()self.safe_dag_id=dag_id.replace('.','__dot__')self.max_active_runs=max_active_runsself.dagrun_timeout=dagrun_timeoutself.sla_miss_callback=sla_miss_callbackself._default_view=default_viewself.orientation=orientationself.catchup=catchupself.is_subdag=False# DagBag.bag_dag() will set this to True if appropriateself.partial=Falseself.on_success_callback=on_success_callbackself.on_failure_callback=on_failure_callbackself.doc_md=doc_mdself._old_context_manager_dags=[]self._access_control=access_controlself._comps={'dag_id','task_ids','parent_dag','start_date','schedule_interval','full_filepath','template_searchpath','last_loaded',}def__repr__(self):return"<DAG: {self.dag_id}>".format(self=self)def__eq__(self,other):if(type(self)==type(other)andself.dag_id==other.dag_id):# Use getattr() instead of __dict__ as __dict__ doesn't return# correct values for properties.returnall(getattr(self,c,None)==getattr(other,c,None)forcinself._comps)returnFalsedef__ne__(self,other):returnnotself==otherdef__lt__(self,other):returnself.dag_id<other.dag_iddef__hash__(self):hash_components=[type(self)]forcinself._comps:# task_ids returns a list and lists can't be hashedifc=='task_ids':val=tuple(self.task_dict.keys())else:val=getattr(self,c,None)try:hash(val)hash_components.append(val)exceptTypeError:hash_components.append(repr(val))returnhash(tuple(hash_components))# Context Manager -----------------------------------------------def__enter__(self):global_CONTEXT_MANAGER_DAGself._old_context_manager_dags.append(_CONTEXT_MANAGER_DAG)_CONTEXT_MANAGER_DAG=selfreturnselfdef__exit__(self,_type,_value,_tb):global_CONTEXT_MANAGER_DAG_CONTEXT_MANAGER_DAG=self._old_context_manager_dags.pop()# /Context Manager ----------------------------------------------

[docs]defget_default_view(self):"""This is only there for backward compatible jinja2 templates"""ifself._default_viewisNone:returnconfiguration.conf.get('webserver','dag_default_view').lower()else:returnself._default_view

[docs]defis_fixed_time_schedule(self):""" Figures out if the DAG schedule has a fixed time (e.g. 3 AM). :return: True if the schedule has a fixed time, False if not. """now=datetime.now()cron=croniter(self._schedule_interval,now)start=cron.get_next(datetime)cron_next=cron.get_next(datetime)ifcron_next.minute==start.minuteandcron_next.hour==start.hour:returnTruereturnFalse

[docs]deffollowing_schedule(self,dttm):""" Calculates the following schedule for this dag in UTC. :param dttm: utc datetime :return: utc datetime """ifisinstance(self._schedule_interval,six.string_types):# we don't want to rely on the transitions created by# croniter as they are not always correctdttm=pendulum.instance(dttm)naive=timezone.make_naive(dttm,self.timezone)cron=croniter(self._schedule_interval,naive)# We assume that DST transitions happen on the minute/hourifnotself.is_fixed_time_schedule():# relative offset (eg. every 5 minutes)delta=cron.get_next(datetime)-naivefollowing=dttm.in_timezone(self.timezone).add_timedelta(delta)else:# absolute (e.g. 3 AM)naive=cron.get_next(datetime)tz=pendulum.timezone(self.timezone.name)following=timezone.make_aware(naive,tz)returntimezone.convert_to_utc(following)elifself._schedule_intervalisnotNone:returndttm+self._schedule_interval

[docs]defprevious_schedule(self,dttm):""" Calculates the previous schedule for this dag in UTC :param dttm: utc datetime :return: utc datetime """ifisinstance(self._schedule_interval,six.string_types):# we don't want to rely on the transitions created by# croniter as they are not always correctdttm=pendulum.instance(dttm)naive=timezone.make_naive(dttm,self.timezone)cron=croniter(self._schedule_interval,naive)# We assume that DST transitions happen on the minute/hourifnotself.is_fixed_time_schedule():# relative offset (eg. every 5 minutes)delta=naive-cron.get_prev(datetime)previous=dttm.in_timezone(self.timezone).subtract_timedelta(delta)else:# absolute (e.g. 3 AM)naive=cron.get_prev(datetime)tz=pendulum.timezone(self.timezone.name)previous=timezone.make_aware(naive,tz)returntimezone.convert_to_utc(previous)elifself._schedule_intervalisnotNone:returndttm-self._schedule_interval

[docs]defget_run_dates(self,start_date,end_date=None):""" Returns a list of dates between the interval received as parameter using this dag's schedule interval. Returned dates can be used for execution dates. :param start_date: the start date of the interval :type start_date: datetime :param end_date: the end date of the interval, defaults to timezone.utcnow() :type end_date: datetime :return: a list of dates within the interval following the dag's schedule :rtype: list """run_dates=[]using_start_date=start_dateusing_end_date=end_date# dates for dag runsusing_start_date=using_start_dateormin([t.start_datefortinself.tasks])using_end_date=using_end_dateortimezone.utcnow()# next run date for a subdag isn't relevant (schedule_interval for subdags# is ignored) so we use the dag run's start date in the case of a subdagnext_run_date=(self.normalize_schedule(using_start_date)ifnotself.is_subdagelseusing_start_date)whilenext_run_dateandnext_run_date<=using_end_date:run_dates.append(next_run_date)next_run_date=self.following_schedule(next_run_date)returnrun_dates

[docs]defnormalize_schedule(self,dttm):""" Returns dttm + interval unless dttm is first interval then it returns dttm """following=self.following_schedule(dttm)# in case of @onceifnotfollowing:returndttmifself.previous_schedule(following)!=dttm:returnfollowingreturndttm

@provide_sessiondefget_last_dagrun(self,session=None,include_externally_triggered=False):returnget_last_dagrun(self.dag_id,session=session,include_externally_triggered=include_externally_triggered)@propertydefdag_id(self):returnself._dag_id@dag_id.setterdefdag_id(self,value):self._dag_id=value@propertydeffull_filepath(self):returnself._full_filepath@full_filepath.setterdeffull_filepath(self,value):self._full_filepath=value@propertydefconcurrency(self):returnself._concurrency@concurrency.setterdefconcurrency(self,value):self._concurrency=value@propertydefaccess_control(self):returnself._access_control@access_control.setterdefaccess_control(self,value):self._access_control=value@propertydefdescription(self):returnself._description@propertydefpickle_id(self):returnself._pickle_id@pickle_id.setterdefpickle_id(self,value):self._pickle_id=value@propertydeftasks(self):returnlist(self.task_dict.values())@tasks.setterdeftasks(self,val):raiseAttributeError('DAG.tasks can not be modified. Use dag.add_task() instead.')@propertydeftask_ids(self):returnlist(self.task_dict.keys())@propertydefactive_task_ids(self):returnlist(kfork,vinself.task_dict.items()ifnotv.adhoc)@propertydefactive_tasks(self):return[tfortinself.tasksifnott.adhoc]@propertydeffilepath(self):""" File location of where the dag object is instantiated """fn=self.full_filepath.replace(settings.DAGS_FOLDER+'/','')fn=fn.replace(os.path.dirname(__file__)+'/','')returnfn@propertydeffolder(self):""" Folder location of where the dag object is instantiated """returnos.path.dirname(self.full_filepath)@propertydefowner(self):return", ".join(list(set([t.ownerfortinself.tasks])))@provide_sessiondef_get_concurrency_reached(self,session=None):TI=TaskInstanceqry=session.query(func.count(TI.task_id)).filter(TI.dag_id==self.dag_id,TI.state==State.RUNNING,)returnqry.scalar()>=self.concurrency@propertydefconcurrency_reached(self):""" Returns a boolean indicating whether the concurrency limit for this DAG has been reached """returnself._get_concurrency_reached()@provide_sessiondef_get_is_paused(self,session=None):qry=session.query(DagModel).filter(DagModel.dag_id==self.dag_id)returnqry.value('is_paused')@propertydefis_paused(self):""" Returns a boolean indicating whether this DAG is paused """returnself._get_is_paused()

[docs]@provide_sessiondefhandle_callback(self,dagrun,success=True,reason=None,session=None):""" Triggers the appropriate callback depending on the value of success, namely the on_failure_callback or on_success_callback. This method gets the context of a single TaskInstance part of this DagRun and passes that to the callable along with a 'reason', primarily to differentiate DagRun failures. .. note: The logs end up in ``$AIRFLOW_HOME/logs/scheduler/latest/PROJECT/DAG_FILE.py.log`` :param dagrun: DagRun object :param success: Flag to specify if failure or success callback should be called :param reason: Completion reason :param session: Database session """callback=self.on_success_callbackifsuccesselseself.on_failure_callbackifcallback:self.log.info('Executing dag callback function: {}'.format(callback))tis=dagrun.get_task_instances()ti=tis[-1]# get first TaskInstance of DagRunti.task=self.get_task(ti.task_id)context=ti.get_template_context(session=session)context.update({'reason':reason})callback(context)

[docs]@provide_sessiondefget_dagrun(self,execution_date,session=None):""" Returns the dag run for a given execution date if it exists, otherwise none. :param execution_date: The execution date of the DagRun to find. :param session: :return: The DagRun if found, otherwise None. """dagrun=(session.query(DagRun).filter(DagRun.dag_id==self.dag_id,DagRun.execution_date==execution_date).first())returndagrun

@provide_sessiondef_get_latest_execution_date(self,session=None):returnsession.query(func.max(DagRun.execution_date)).filter(DagRun.dag_id==self.dag_id).scalar()@propertydeflatest_execution_date(self):""" Returns the latest date for which at least one dag run exists """returnself._get_latest_execution_date()@propertydefsubdags(self):""" Returns a list of the subdag objects associated to this DAG """# Check SubDag for class but don't check class directlyfromairflow.operators.subdag_operatorimportSubDagOperatorsubdag_lst=[]fortaskinself.tasks:if(isinstance(task,SubDagOperator)or# TODO remove in Airflow 2.0type(task).__name__=='SubDagOperator'):subdag_lst.append(task.subdag)subdag_lst+=task.subdag.subdagsreturnsubdag_lstdefresolve_template_files(self):fortinself.tasks:t.resolve_template_files()

[docs]defget_template_env(self):""" Returns a jinja2 Environment while taking into account the DAGs template_searchpath, user_defined_macros and user_defined_filters """searchpath=[self.folder]ifself.template_searchpath:searchpath+=self.template_searchpathenv=jinja2.Environment(loader=jinja2.FileSystemLoader(searchpath),extensions=["jinja2.ext.do"],cache_size=0)ifself.user_defined_macros:env.globals.update(self.user_defined_macros)ifself.user_defined_filters:env.filters.update(self.user_defined_filters)returnenv

[docs]defset_dependency(self,upstream_task_id,downstream_task_id):""" Simple utility method to set dependency between two tasks that already have been added to the DAG using add_task() """self.get_task(upstream_task_id).set_downstream(self.get_task(downstream_task_id))

[docs]deftopological_sort(self):""" Sorts tasks in topographical order, such that a task comes after any of its upstream dependencies. Heavily inspired by: http://blog.jupo.org/2012/04/06/topological-sorting-acyclic-directed-graphs/ :return: list of tasks in topological order """# convert into an OrderedDict to speedup lookup while keeping order the samegraph_unsorted=OrderedDict((task.task_id,task)fortaskinself.tasks)graph_sorted=[]# special caseiflen(self.tasks)==0:returntuple(graph_sorted)# Run until the unsorted graph is empty.whilegraph_unsorted:# Go through each of the node/edges pairs in the unsorted# graph. If a set of edges doesn't contain any nodes that# haven't been resolved, that is, that are still in the# unsorted graph, remove the pair from the unsorted graph,# and append it to the sorted graph. Note here that by using# using the items() method for iterating, a copy of the# unsorted graph is used, allowing us to modify the unsorted# graph as we move through it. We also keep a flag for# checking that that graph is acyclic, which is true if any# nodes are resolved during each pass through the graph. If# not, we need to bail out as the graph therefore can't be# sorted.acyclic=Falsefornodeinlist(graph_unsorted.values()):foredgeinnode.upstream_list:ifedge.task_idingraph_unsorted:break# no edges in upstream taskselse:acyclic=Truedelgraph_unsorted[node.task_id]graph_sorted.append(node)ifnotacyclic:raiseAirflowException("A cyclic dependency occurred in dag: {}".format(self.dag_id))returntuple(graph_sorted)

[docs]@provide_sessiondefclear(self,start_date=None,end_date=None,only_failed=False,only_running=False,confirm_prompt=False,include_subdags=True,include_parentdag=True,reset_dag_runs=True,dry_run=False,session=None,get_tis=False,):""" Clears a set of task instances associated with the current dag for a specified date range. """TI=TaskInstancetis=session.query(TI)ifinclude_subdags:# Crafting the right filter for dag_id and task_ids comboconditions=[]fordaginself.subdags+[self]:conditions.append(TI.dag_id.like(dag.dag_id)&TI.task_id.in_(dag.task_ids))tis=tis.filter(or_(*conditions))else:tis=session.query(TI).filter(TI.dag_id==self.dag_id)tis=tis.filter(TI.task_id.in_(self.task_ids))ifinclude_parentdagandself.is_subdag:p_dag=self.parent_dag.sub_dag(task_regex=self.dag_id.split('.')[1],include_upstream=False,include_downstream=True)tis=tis.union(p_dag.clear(start_date=start_date,end_date=end_date,only_failed=only_failed,only_running=only_running,confirm_prompt=confirm_prompt,include_subdags=include_subdags,include_parentdag=False,reset_dag_runs=reset_dag_runs,get_tis=True,session=session,))ifstart_date:tis=tis.filter(TI.execution_date>=start_date)ifend_date:tis=tis.filter(TI.execution_date<=end_date)ifonly_failed:tis=tis.filter(or_(TI.state==State.FAILED,TI.state==State.UPSTREAM_FAILED))ifonly_running:tis=tis.filter(TI.state==State.RUNNING)ifget_tis:returntisifdry_run:tis=tis.all()session.expunge_all()returntiscount=tis.count()do_it=Trueifcount==0:return0ifconfirm_prompt:ti_list="\n".join([str(t)fortintis])question=("You are about to delete these {count} tasks:\n""{ti_list}\n\n""Are you sure? (yes/no): ").format(**locals())do_it=utils.helpers.ask_yesno(question)ifdo_it:clear_task_instances(tis.all(),session,dag=self,)ifreset_dag_runs:self.set_dag_runs_state(session=session,start_date=start_date,end_date=end_date,)else:count=0print("Bail. Nothing was cleared.")session.commit()returncount

@classmethoddefclear_dags(cls,dags,start_date=None,end_date=None,only_failed=False,only_running=False,confirm_prompt=False,include_subdags=True,include_parentdag=False,reset_dag_runs=True,dry_run=False,):all_tis=[]fordagindags:tis=dag.clear(start_date=start_date,end_date=end_date,only_failed=only_failed,only_running=only_running,confirm_prompt=False,include_subdags=include_subdags,include_parentdag=include_parentdag,reset_dag_runs=reset_dag_runs,dry_run=True)all_tis.extend(tis)ifdry_run:returnall_tiscount=len(all_tis)do_it=Trueifcount==0:print("Nothing to clear.")return0ifconfirm_prompt:ti_list="\n".join([str(t)fortinall_tis])question=("You are about to delete these {} tasks:\n""{}\n\n""Are you sure? (yes/no): ").format(count,ti_list)do_it=utils.helpers.ask_yesno(question)ifdo_it:fordagindags:dag.clear(start_date=start_date,end_date=end_date,only_failed=only_failed,only_running=only_running,confirm_prompt=False,include_subdags=include_subdags,reset_dag_runs=reset_dag_runs,dry_run=False,)else:count=0print("Bail. Nothing was cleared.")returncountdef__deepcopy__(self,memo):# Swiwtcharoo to go around deepcopying objects coming through the# backdoorcls=self.__class__result=cls.__new__(cls)memo[id(self)]=resultfork,vinlist(self.__dict__.items()):ifknotin('user_defined_macros','user_defined_filters','params'):setattr(result,k,copy.deepcopy(v,memo))result.user_defined_macros=self.user_defined_macrosresult.user_defined_filters=self.user_defined_filtersresult.params=self.paramsreturnresult

[docs]defsub_dag(self,task_regex,include_downstream=False,include_upstream=True):""" Returns a subset of the current dag as a deep copy of the current dag based on a regex that should match one or many tasks, and includes upstream and downstream neighbours based on the flag passed. """# deep-copying self.task_dict takes a long time, and we don't want all# the tasks anyway, so we copy the tasks manually latertask_dict=self.task_dictself.task_dict={}dag=copy.deepcopy(self)self.task_dict=task_dictregex_match=[tfortinself.tasksifre.findall(task_regex,t.task_id)]also_include=[]fortinregex_match:ifinclude_downstream:also_include+=t.get_flat_relatives(upstream=False)ifinclude_upstream:also_include+=t.get_flat_relatives(upstream=True)# Compiling the unique list of tasks that made the cut# Make sure to not recursively deepcopy the dag while copying the taskdag.task_dict={t.task_id:copy.deepcopy(t,{id(t.dag):t.dag})fortinregex_match+also_include}fortindag.tasks:# Removing upstream/downstream references to tasks that did not# made the cutt._upstream_task_ids=t._upstream_task_ids.intersection(dag.task_dict.keys())t._downstream_task_ids=t._downstream_task_ids.intersection(dag.task_dict.keys())iflen(dag.tasks)<len(self.tasks):dag.partial=Truereturndag

[docs]defadd_task(self,task):""" Add a task to the DAG :param task: the task you want to add :type task: task """ifnotself.start_dateandnottask.start_date:raiseAirflowException("Task is missing the start_date parameter")# if the task has no start date, assign it the same as the DAGelifnottask.start_date:task.start_date=self.start_date# otherwise, the task will start on the later of its own start date and# the DAG's start dateelifself.start_date:task.start_date=max(task.start_date,self.start_date)# if the task has no end date, assign it the same as the dagifnottask.end_date:task.end_date=self.end_date# otherwise, the task will end on the earlier of its own end date and# the DAG's end dateeliftask.end_dateandself.end_date:task.end_date=min(task.end_date,self.end_date)iftask.task_idinself.task_dict:# TODO: raise an error in Airflow 2.0warnings.warn('The requested task could not be added to the DAG because a ''task with task_id {} is already in the DAG. Starting in ''Airflow 2.0, trying to overwrite a task will raise an ''exception.'.format(task.task_id),category=PendingDeprecationWarning)else:self.task_dict[task.task_id]=tasktask.dag=selfself.task_count=len(self.task_dict)

[docs]defadd_tasks(self,tasks):""" Add a list of tasks to the DAG :param tasks: a lit of tasks you want to add :type tasks: list of tasks """fortaskintasks:self.add_task(task)

[docs]defrun(self,start_date=None,end_date=None,mark_success=False,local=False,executor=None,donot_pickle=configuration.conf.getboolean('core','donot_pickle'),ignore_task_deps=False,ignore_first_depends_on_past=False,pool=None,delay_on_limit_secs=1.0,verbose=False,conf=None,rerun_failed_tasks=False,run_backwards=False,):""" Runs the DAG. :param start_date: the start date of the range to run :type start_date: datetime.datetime :param end_date: the end date of the range to run :type end_date: datetime.datetime :param mark_success: True to mark jobs as succeeded without running them :type mark_success: bool :param local: True to run the tasks using the LocalExecutor :type local: bool :param executor: The executor instance to run the tasks :type executor: airflow.executor.BaseExecutor :param donot_pickle: True to avoid pickling DAG object and send to workers :type donot_pickle: bool :param ignore_task_deps: True to skip upstream tasks :type ignore_task_deps: bool :param ignore_first_depends_on_past: True to ignore depends_on_past dependencies for the first set of tasks only :type ignore_first_depends_on_past: bool :param pool: Resource pool to use :type pool: str :param delay_on_limit_secs: Time in seconds to wait before next attempt to run dag run when max_active_runs limit has been reached :type delay_on_limit_secs: float :param verbose: Make logging output more verbose :type verbose: bool :param conf: user defined dictionary passed from CLI :type conf: dict :param rerun_failed_tasks: :type: bool :param run_backwards: :type: bool """fromairflow.jobsimportBackfillJobifnotexecutorandlocal:executor=LocalExecutor()elifnotexecutor:executor=get_default_executor()job=BackfillJob(self,start_date=start_date,end_date=end_date,mark_success=mark_success,executor=executor,donot_pickle=donot_pickle,ignore_task_deps=ignore_task_deps,ignore_first_depends_on_past=ignore_first_depends_on_past,pool=pool,delay_on_limit_secs=delay_on_limit_secs,verbose=verbose,conf=conf,rerun_failed_tasks=rerun_failed_tasks,run_backwards=run_backwards,)job.run()

[docs]defcli(self):""" Exposes a CLI specific to this DAG """fromairflow.binimportcliparser=cli.CLIFactory.get_parser(dag_parser=True)args=parser.parse_args()args.func(args,self)

[docs]@provide_sessiondefcreate_dagrun(self,run_id,state,execution_date=None,start_date=None,external_trigger=False,conf=None,session=None):""" Creates a dag run from this dag including the tasks associated with this dag. Returns the dag run. :param run_id: defines the the run id for this dag run :type run_id: str :param execution_date: the execution date of this dag run :type execution_date: datetime.datetime :param state: the state of the dag run :type state: airflow.utils.state.State :param start_date: the date this dag run should be evaluated :type start_date: datetime :param external_trigger: whether this dag run is externally triggered :type external_trigger: bool :param session: database session :type session: sqlalchemy.orm.session.Session """run=DagRun(dag_id=self.dag_id,run_id=run_id,execution_date=execution_date,start_date=start_date,external_trigger=external_trigger,conf=conf,state=state)session.add(run)session.commit()run.dag=self# create the associated task instances# state is None at the moment of creationrun.verify_integrity(session=session)run.refresh_from_db()returnrun

[docs]@provide_sessiondefsync_to_db(self,owner=None,sync_time=None,session=None):""" Save attributes about this DAG to the DB. Note that this method can be called for both DAGs and SubDAGs. A SubDag is actually a SubDagOperator. :param dag: the DAG object to save to the DB :type dag: airflow.models.DAG :param sync_time: The time that the DAG should be marked as sync'ed :type sync_time: datetime :return: None """ifownerisNone:owner=self.ownerifsync_timeisNone:sync_time=timezone.utcnow()orm_dag=session.query(DagModel).filter(DagModel.dag_id==self.dag_id).first()ifnotorm_dag:orm_dag=DagModel(dag_id=self.dag_id)self.log.info("Creating ORM DAG for %s",self.dag_id)orm_dag.fileloc=self.filelocorm_dag.is_subdag=self.is_subdagorm_dag.owners=ownerorm_dag.is_active=Trueorm_dag.last_scheduler_run=sync_timeorm_dag.default_view=self._default_vieworm_dag.description=self.descriptionorm_dag.schedule_interval=self.schedule_intervalsession.merge(orm_dag)session.commit()forsubdaginself.subdags:subdag.sync_to_db(owner=owner,sync_time=sync_time,session=session)

[docs]@staticmethod@provide_sessiondefdeactivate_unknown_dags(active_dag_ids,session=None):""" Given a list of known DAGs, deactivate any other DAGs that are marked as active in the ORM :param active_dag_ids: list of DAG IDs that are active :type active_dag_ids: list[unicode] :return: None """iflen(active_dag_ids)==0:returnfordaginsession.query(DagModel).filter(~DagModel.dag_id.in_(active_dag_ids)).all():dag.is_active=Falsesession.merge(dag)session.commit()

[docs]@staticmethod@provide_sessiondefdeactivate_stale_dags(expiration_date,session=None):""" Deactivate any DAGs that were last touched by the scheduler before the expiration date. These DAGs were likely deleted. :param expiration_date: set inactive DAGs that were touched before this time :type expiration_date: datetime :return: None """log=LoggingMixin().logfordaginsession.query(DagModel).filter(DagModel.last_scheduler_run<expiration_date,DagModel.is_active).all():log.info("Deactivating DAG ID %s since it was last touched by the scheduler at %s",dag.dag_id,dag.last_scheduler_run.isoformat())dag.is_active=Falsesession.merge(dag)session.commit()

[docs]@staticmethod@provide_sessiondefget_num_task_instances(dag_id,task_ids,states=None,session=None):""" Returns the number of task instances in the given DAG. :param session: ORM session :param dag_id: ID of the DAG to get the task concurrency of :type dag_id: unicode :param task_ids: A list of valid task IDs for the given DAG :type task_ids: list[unicode] :param states: A list of states to filter by if supplied :type states: list[state] :return: The number of running tasks :rtype: int """qry=session.query(func.count(TaskInstance.task_id)).filter(TaskInstance.dag_id==dag_id,TaskInstance.task_id.in_(task_ids))ifstatesisnotNone:ifNoneinstates:qry=qry.filter(or_(TaskInstance.state.in_(states),TaskInstance.state.is_(None)))else:qry=qry.filter(TaskInstance.state.in_(states))returnqry.scalar()

[docs]deftest_cycle(self):""" Check to see if there are any cycles in the DAG. Returns False if no cycle found, otherwise raises exception. """# default of int is 0 which corresponds to CYCLE_NEWvisit_map=defaultdict(int)fortask_idinself.task_dict.keys():# print('starting %s' % task_id)ifvisit_map[task_id]==DagBag.CYCLE_NEW:self._test_cycle_helper(visit_map,task_id)returnFalse

[docs]@classmethoddefsetdefault(cls,key,default,deserialize_json=False):""" Like a Python builtin dict object, setdefault returns the current value for a key, and if it isn't there, stores the default value and returns it. :param key: Dict key for this Variable :type key: str :param default: Default value to set and return if the variable isn't already in the DB :type default: Mixed :param deserialize_json: Store this as a JSON encoded value in the DB and un-encode it when retrieving a value :return: Mixed """obj=Variable.get(key,default_var=None,deserialize_json=deserialize_json)ifobjisNone:ifdefaultisnotNone:Variable.set(key,default,serialize_json=deserialize_json)returndefaultelse:raiseValueError('Default Value must be set')else:returnobj

@classmethod@provide_sessiondefget(cls,key,default_var=__NO_DEFAULT_SENTINEL,deserialize_json=False,session=None):obj=session.query(cls).filter(cls.key==key).first()ifobjisNone:ifdefault_varisnotcls.__NO_DEFAULT_SENTINEL:returndefault_varelse:raiseKeyError('Variable {} does not exist'.format(key))else:ifdeserialize_json:returnjson.loads(obj.val)else:returnobj.val@classmethod@provide_sessiondefset(cls,key,value,serialize_json=False,session=None):ifserialize_json:stored_value=json.dumps(value)else:stored_value=str(value)session.query(cls).filter(cls.key==key).delete()session.add(Variable(key=key,val=stored_value))session.flush()defrotate_fernet_key(self):fernet=get_fernet()ifself._valandself.is_encrypted:self._val=fernet.rotate(self._val.encode('utf-8')).decode()

[docs]classDagRun(Base,LoggingMixin):""" DagRun describes an instance of a Dag. It can be created by the scheduler (for regular runs) or by an external trigger """__tablename__="dag_run"ID_PREFIX='scheduled__'ID_FORMAT_PREFIX=ID_PREFIX+'{0}'id=Column(Integer,primary_key=True)dag_id=Column(String(ID_LEN))execution_date=Column(UtcDateTime,default=timezone.utcnow)start_date=Column(UtcDateTime,default=timezone.utcnow)end_date=Column(UtcDateTime)_state=Column('state',String(50),default=State.RUNNING)run_id=Column(String(ID_LEN))external_trigger=Column(Boolean,default=True)conf=Column(PickleType)dag=None__table_args__=(Index('dag_id_state',dag_id,_state),UniqueConstraint('dag_id','execution_date'),UniqueConstraint('dag_id','run_id'),)def__repr__(self):return('<DagRun {dag_id} @ {execution_date}: {run_id}, ''externally triggered: {external_trigger}>').format(dag_id=self.dag_id,execution_date=self.execution_date,run_id=self.run_id,external_trigger=self.external_trigger)defget_state(self):returnself._statedefset_state(self,state):ifself._state!=state:self._state=stateself.end_date=timezone.utcnow()ifself._stateinState.finished()elseNone@declared_attrdefstate(self):returnsynonym('_state',descriptor=property(self.get_state,self.set_state))@classmethoddefid_for_date(cls,date,prefix=ID_FORMAT_PREFIX):returnprefix.format(date.isoformat()[:19])

[docs]@provide_sessiondefrefresh_from_db(self,session=None):""" Reloads the current dagrun from the database :param session: database session """DR=DagRunexec_date=func.cast(self.execution_date,DateTime)dr=session.query(DR).filter(DR.dag_id==self.dag_id,func.cast(DR.execution_date,DateTime)==exec_date,DR.run_id==self.run_id).one()self.id=dr.idself.state=dr.state

[docs]@staticmethod@provide_sessiondeffind(dag_id=None,run_id=None,execution_date=None,state=None,external_trigger=None,no_backfills=False,session=None):""" Returns a set of dag runs for the given search criteria. :param dag_id: the dag_id to find dag runs for :type dag_id: int, list :param run_id: defines the the run id for this dag run :type run_id: str :param execution_date: the execution date :type execution_date: datetime.datetime :param state: the state of the dag run :type state: airflow.utils.state.State :param external_trigger: whether this dag run is externally triggered :type external_trigger: bool :param no_backfills: return no backfills (True), return all (False). Defaults to False :type no_backfills: bool :param session: database session :type session: sqlalchemy.orm.session.Session """DR=DagRunqry=session.query(DR)ifdag_id:qry=qry.filter(DR.dag_id==dag_id)ifrun_id:qry=qry.filter(DR.run_id==run_id)ifexecution_date:ifisinstance(execution_date,list):qry=qry.filter(DR.execution_date.in_(execution_date))else:qry=qry.filter(DR.execution_date==execution_date)ifstate:qry=qry.filter(DR.state==state)ifexternal_triggerisnotNone:qry=qry.filter(DR.external_trigger==external_trigger)ifno_backfills:# in order to prevent a circular dependencyfromairflow.jobsimportBackfillJobqry=qry.filter(DR.run_id.notlike(BackfillJob.ID_PREFIX+'%'))dr=qry.order_by(DR.execution_date).all()returndr

[docs]@provide_sessiondefget_task_instances(self,state=None,session=None):""" Returns the task instances for this dag run """tis=session.query(TaskInstance).filter(TaskInstance.dag_id==self.dag_id,TaskInstance.execution_date==self.execution_date,)ifstate:ifisinstance(state,six.string_types):tis=tis.filter(TaskInstance.state==state)else:# this is required to deal with NULL valuesifNoneinstate:tis=tis.filter(or_(TaskInstance.state.in_(state),TaskInstance.state.is_(None)))else:tis=tis.filter(TaskInstance.state.in_(state))ifself.dagandself.dag.partial:tis=tis.filter(TaskInstance.task_id.in_(self.dag.task_ids))returntis.all()

[docs]@provide_sessiondefget_task_instance(self,task_id,session=None):""" Returns the task instance specified by task_id for this dag run :param task_id: the task id """TI=TaskInstanceti=session.query(TI).filter(TI.dag_id==self.dag_id,TI.execution_date==self.execution_date,TI.task_id==task_id).first()returnti

[docs]defget_dag(self):""" Returns the Dag associated with this DagRun. :return: DAG """ifnotself.dag:raiseAirflowException("The DAG (.dag) for {} needs to be set".format(self))returnself.dag

[docs]@provide_sessiondefget_previous_dagrun(self,session=None):"""The previous DagRun, if there is one"""returnsession.query(DagRun).filter(DagRun.dag_id==self.dag_id,DagRun.execution_date<self.execution_date).order_by(DagRun.execution_date.desc()).first()

[docs]@provide_sessiondefget_previous_scheduled_dagrun(self,session=None):"""The previous, SCHEDULED DagRun, if there is one"""dag=self.get_dag()returnsession.query(DagRun).filter(DagRun.dag_id==self.dag_id,DagRun.execution_date==dag.previous_schedule(self.execution_date)).first()

[docs]@provide_sessiondefupdate_state(self,session=None):""" Determines the overall state of the DagRun based on the state of its TaskInstances. :return: State """dag=self.get_dag()tis=self.get_task_instances(session=session)self.log.debug("Updating state for %s considering %s task(s)",self,len(tis))fortiinlist(tis):# skip in db?ifti.state==State.REMOVED:tis.remove(ti)else:ti.task=dag.get_task(ti.task_id)# pre-calculate# db is fasterstart_dttm=timezone.utcnow()unfinished_tasks=self.get_task_instances(state=State.unfinished(),session=session)none_depends_on_past=all(nott.task.depends_on_pastfortinunfinished_tasks)none_task_concurrency=all(t.task.task_concurrencyisNonefortinunfinished_tasks)# small speed upifunfinished_tasksandnone_depends_on_pastandnone_task_concurrency:# todo: this can actually get pretty slow: one task costs between 0.01-015sno_dependencies_met=Trueforutinunfinished_tasks:# We need to flag upstream and check for changes because upstream# failures/re-schedules can result in deadlock false positivesold_state=ut.statedeps_met=ut.are_dependencies_met(dep_context=DepContext(flag_upstream_failed=True,ignore_in_retry_period=True,ignore_in_reschedule_period=True),session=session)ifdeps_metorold_state!=ut.current_state(session=session):no_dependencies_met=Falsebreakduration=(timezone.utcnow()-start_dttm).total_seconds()*1000Stats.timing("dagrun.dependency-check.{}".format(self.dag_id),duration)# future: remove the check on adhoc tasks (=active_tasks)iflen(tis)==len(dag.active_tasks):root_ids=[t.task_idfortindag.roots]roots=[tfortintisift.task_idinroot_ids]# if all roots finished and at least one failed, the run failedif(notunfinished_tasksandany(r.statein(State.FAILED,State.UPSTREAM_FAILED)forrinroots)):self.log.info('Marking run %s failed',self)self.set_state(State.FAILED)dag.handle_callback(self,success=False,reason='task_failure',session=session)# if all roots succeeded and no unfinished tasks, the run succeededelifnotunfinished_tasksandall(r.statein(State.SUCCESS,State.SKIPPED)forrinroots):self.log.info('Marking run %s successful',self)self.set_state(State.SUCCESS)dag.handle_callback(self,success=True,reason='success',session=session)# if *all tasks* are deadlocked, the run failedelif(unfinished_tasksandnone_depends_on_pastandnone_task_concurrencyandno_dependencies_met):self.log.info('Deadlock; marking run %s failed',self)self.set_state(State.FAILED)dag.handle_callback(self,success=False,reason='all_tasks_deadlocked',session=session)# finally, if the roots aren't done, the dag is still runningelse:self.set_state(State.RUNNING)self._emit_duration_stats_for_finished_state()# todo: determine we want to use with_for_update to make sure to lock the runsession.merge(self)session.commit()returnself.state

[docs]@provide_sessiondefverify_integrity(self,session=None):""" Verifies the DagRun by checking for removed tasks or tasks that are not in the database yet. It will set state to removed or add the task if required. """dag=self.get_dag()tis=self.get_task_instances(session=session)# check for removed or restored taskstask_ids=[]fortiintis:task_ids.append(ti.task_id)task=Nonetry:task=dag.get_task(ti.task_id)exceptAirflowException:ifti.state==State.REMOVED:pass# ti has already been removed, just ignore itelifself.stateisnotState.RUNNINGandnotdag.partial:self.log.warning("Failed to get task '{}' for dag '{}'. ""Marking it as removed.".format(ti,dag))Stats.incr("task_removed_from_dag.{}".format(dag.dag_id),1,1)ti.state=State.REMOVEDis_task_in_dag=taskisnotNoneshould_restore_task=is_task_in_dagandti.state==State.REMOVEDifshould_restore_task:self.log.info("Restoring task '{}' which was previously ""removed from DAG '{}'".format(ti,dag))Stats.incr("task_restored_to_dag.{}".format(dag.dag_id),1,1)ti.state=State.NONE# check for missing tasksfortaskinsix.itervalues(dag.task_dict):iftask.adhoc:continueiftask.start_date>self.execution_dateandnotself.is_backfill:continueiftask.task_idnotintask_ids:Stats.incr("task_instance_created-{}".format(task.__class__.__name__),1,1)ti=TaskInstance(task,self.execution_date)session.add(ti)session.commit()

[docs]@classmethod@provide_sessiondefget_latest_runs(cls,session):"""Returns the latest DagRun for each DAG. """subquery=(session.query(cls.dag_id,func.max(cls.execution_date).label('execution_date')).group_by(cls.dag_id).subquery())dagruns=(session.query(cls).join(subquery,and_(cls.dag_id==subquery.c.dag_id,cls.execution_date==subquery.c.execution_date)).all())returndagruns

[docs]@provide_sessiondefused_slots(self,session):""" Returns the number of slots used at the moment """running=(session.query(TaskInstance).filter(TaskInstance.pool==self.pool).filter(TaskInstance.state==State.RUNNING).count())returnrunning

[docs]@provide_sessiondefqueued_slots(self,session):""" Returns the number of slots used at the moment """return(session.query(TaskInstance).filter(TaskInstance.pool==self.pool).filter(TaskInstance.state==State.QUEUED).count())

[docs]@provide_sessiondefopen_slots(self,session):""" Returns the number of slots open at the moment """used_slots=self.used_slots(session=session)queued_slots=self.queued_slots(session=session)returnself.slots-used_slots-queued_slots