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')
