def main(): try: from airflow import conf except ImportError: from airflow.configuration import conf from airflow.jobs.scheduler_job import SchedulerJob from airflow.models import DagBag from dbnd import dbnd_bootstrap from dbnd._core.log.logging_utils import create_file_handler from dbnd_airflow.executors.simple_executor import InProcessExecutor from test_dbnd_airflow.scenarios.scheduler_perf_experiment import ( dag_folder, dag_id, log_scheduler, ) dbnd_bootstrap() conf.set("core", "unit_test_mode", "True") logging.root.addHandler(create_file_handler(log_file=log_scheduler)) dag_bag = DagBag(dag_folder=dag_folder) scheduler_job = SchedulerJob( dag_ids=[dag_id], subdir=dag_folder, do_pickle=False, num_runs=3, executor=InProcessExecutor(dag_bag=dag_bag), ) scheduler_job.run()
def export_data_directly( self, since, include_logs, include_task_args, include_xcom, dag_ids, quantity, incomplete_offset, dags_only, ): from airflow import models, settings, conf from airflow.settings import STORE_SERIALIZED_DAGS from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from dbnd_airflow_export.dbnd_airflow_export_plugin import get_airflow_data conf.set("core", "sql_alchemy_conn", value=self.sql_conn_string) dagbag = models.DagBag( self.dag_folder if self.dag_folder else settings.DAGS_FOLDER, include_examples=True, store_serialized_dags=STORE_SERIALIZED_DAGS, ) engine = create_engine(self.sql_conn_string) session = sessionmaker(bind=engine) result = get_airflow_data( dagbag=dagbag, since=since, include_logs=include_logs, include_task_args=include_task_args, include_xcom=include_xcom, dag_ids=dag_ids, quantity=quantity, incomplete_offset=incomplete_offset, dags_only=dags_only, session=session(), ) return result
def _get_session(self): import airflow if hasattr(airflow, "conf"): from airflow import conf else: from airflow.configuration import conf from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker if not self._engine: conf.set("core", "sql_alchemy_conn", value=self.sql_conn_string) self._engine = create_engine(self.sql_conn_string) self._session = sessionmaker(bind=self._engine) session = self._session() try: yield session session.commit() except Exception: session.rollback() raise
def update_config(args): logging.info("Update Airflow configuration") with open(conf.AIRFLOW_CONFIG, 'w') as output_stream: try: conf.conf.add_section('cwl') except configparser.DuplicateSectionError: pass conf.set('core', 'dags_are_paused_at_creation', 'False') conf.set('core', 'load_examples', 'False') conf.set('cwl', 'jobs', str(args.jobs)) conf.set('cwl', 'limit', str(args.limit)) conf.set('cwl', 'logging_level', conf_get_default('cwl', 'logging_level', 'ERROR') ) # To supress all useless output from cwltool's functions conf.set('core', 'dagbag_import_timeout', str(args.dag_timeout)) conf.set('scheduler', 'max_threads', str(args.threads)) conf.set('webserver', 'worker_refresh_interval', str(args.web_interval)) conf.set('webserver', 'worker_refresh_batch_size', str(args.web_workers)) conf.set('webserver', 'hide_paused_dags_by_default', 'True') conf.conf.write(output_stream)
def setUp(self): conf.set('smtp', 'SMTP_SSL', 'False')