def conf_vars(overrides):
original = {}
original_env_vars = {}
for (section, key), value in overrides.items():
env = conf._env_var_name(section, key)
if env in os.environ:
original_env_vars[env] = os.environ.pop(env)
if conf.has_option(section, key):
original[(section, key)] = conf.get(section, key)
else:
original[(section, key)] = None
if value is not None:
if not conf.has_section(section):
conf.add_section(section)
conf.set(section, key, value)
else:
conf.remove_option(section, key)
settings.configure_vars()
try:
yield
finally:
for (section, key), value in original.items():
if value is not None:
conf.set(section, key, value)
else:
conf.remove_option(section, key)
for env, value in original_env_vars.items():
os.environ[env] = value
settings.configure_vars()
def setUp(self):
if not conf.has_section("kerberos"):
conf.add_section("kerberos")
conf.set("kerberos", "keytab",
os.environ['KRB5_KTNAME'])
keytab_from_cfg = conf.get("kerberos", "keytab")
self.args = Namespace(keytab=keytab_from_cfg, principal=None, pid=None,
daemon=None, stdout=None, stderr=None, log_file=None)
def get_value(args):
"""Get one value from configuration"""
if not conf.has_section(args.section):
print(f'The section [{args.section}] is not found in config.', file=sys.stderr)
sys.exit(1)
if not conf.has_option(args.section, args.option):
print(f'The option [{args.section}/{args.option}] is not found in config.', file=sys.stderr)
sys.exit(1)
value = conf.get(args.section, args.option)
print(value)
def get_value(args):
"""Get one value from configuration"""
if not conf.has_section(args.section):
raise SystemExit(
f'The section [{args.section}] is not found in config.')
if not conf.has_option(args.section, args.option):
raise SystemExit(
f'The option [{args.section}/{args.option}] is not found in config.'
)
value = conf.get(args.section, args.option)
print(value)
def __set_mocked_executor(self):
"""Mock ECS such that there's nothing wrong with anything"""
from airflow.configuration import conf
if not conf.has_section('batch'):
conf.add_section('batch')
conf.set('batch', 'region', 'us-west-1')
conf.set('batch', 'job_name', 'some-job-name')
conf.set('batch', 'job_queue', 'some-job-queue')
conf.set('batch', 'job_definition', 'some-job-def')
executor = AwsBatchExecutor()
executor.start()
# replace boto3 ecs client with mock
batch_mock = mock.Mock(spec=executor.batch)
submit_job_ret_val = {
'jobName': conf.get('batch', 'job_name'),
'jobId': 'ABC'
}
batch_mock.submit_job.return_value = submit_job_ret_val
executor.batch = batch_mock
self.executor = executor
def __set_mocked_executor(self):
"""Mock ECS such that there's nothing wrong with anything"""
from airflow.configuration import conf
if not conf.has_section('ecs_fargate'):
conf.add_section('ecs_fargate')
conf.set('ecs_fargate', 'region', 'us-west-1')
conf.set('ecs_fargate', 'cluster', 'some-ecs-cluster')
conf.set('ecs_fargate', 'task_definition', 'some-ecs-task-definition')
conf.set('ecs_fargate', 'container_name', 'some-ecs-container')
conf.set('ecs_fargate', 'launch_type', 'FARGATE')
executor = AwsEcsFargateExecutor()
executor.start()
# replace boto3 ecs client with mock
ecs_mock = mock.Mock(spec=executor.ecs)
run_task_ret_val = {
'tasks': [{'taskArn': '001'}],
'failures': []
}
ecs_mock.run_task.return_value = run_task_ret_val
executor.ecs = ecs_mock
self.executor = executor
def patch_airflow_config(args):
"""
Updates current Airflow configuration file to include defaults for cwl-airflow.
If something went wrong, restores the original airflow.cfg from the backed up copy.
If update to Airflow 2.0.0 is required, generates new airflow.cfg with some of the
important parameters copied from the old airflow.cfg. Backed up copy is not deleted in
this case.
"""
# TODO: add cwl section with the following parameters:
# - singularity
# - use_container
# CWL-Airflow specific settings
patches = [
["sed", "-i", "-e", "s#^dags_are_paused_at_creation.*#dags_are_paused_at_creation = False#g", args.config],
["sed", "-i", "-e", "s#^load_examples.*#load_examples = False#g", args.config],
["sed", "-i", "-e", "s#^load_default_connections.*#load_default_connections = False#g", args.config],
["sed", "-i", "-e", "s#^logging_config_class.*#logging_config_class = cwl_airflow.config_templates.airflow_local_settings.DEFAULT_LOGGING_CONFIG#g", args.config],
["sed", "-i", "-e", "s#^hide_paused_dags_by_default.*#hide_paused_dags_by_default = True#g", args.config]
]
# Minimum amount of setting that should be enough for starting
# SequentialExecutor, LocalExecutor or CeleryExecutor with
# the same dags and metadata database after updating to Airflow 2.0.0.
# All other user specific settings should be manually updated from the
# backuped airflow.cfg as a lot of them have been refactored.
transferable_settings = [
("core", "dags_folder"),
("core", "default_timezone"),
("core", "executor"),
("core", "sql_alchemy_conn"),
("core", "sql_engine_encoding"), # just in case
("core", "fernet_key"), # to be able to read from the old database
("celery", "broker_url"),
("celery", "result_backend")
]
# create a temporary backup of airflow.cfg to restore from if we failed to apply patches
# this backup will be deleted after all patches applied if it wasn't created right before
# Airflow version update to 2.0.0
airflow_config_backup = args.config + "_backup_" + str(uuid.uuid4())
try:
# reading aiflow.cfg before applying any patches and creating backup
conf.read(args.config)
shutil.copyfile(args.config, airflow_config_backup)
# check if we need to make airflow.cfg correspond to the Airflow 2.0.0
# we search for [logging] section as it's present only Airflow >= 2.0.0
airflow_version_update = not conf.has_section("logging")
if airflow_version_update:
logging.info("Airflow config will be upgraded to correspond to Airflow 2.0.0")
for section, key in transferable_settings:
try:
patches.append(
["sed", "-i", "-e", f"s#^{key}.*#{key} = {conf.get(section, key)}#g", args.config]
)
except AirflowConfigException: # just skip missing in the config section/key
pass
os.remove(args.config) # remove old config
create_airflow_config(args) # create new airflow.cfg with the default values
# Apply all patches
for patch in patches:
logging.debug(f"Applying patch {patch}")
run(
patch,
shell=False, # for proper handling of filenames with spaces
check=True,
stdout=DEVNULL,
stderr=DEVNULL
)
except (CalledProcessError, FileNotFoundError) as err:
logging.error(f"""Failed to patch Airflow configuration file. Restoring from the backup and exiting.\n{err}""")
if os.path.isfile(airflow_config_backup):
shutil.copyfile(airflow_config_backup, args.config)
sys.exit(1)
finally:
if os.path.isfile(airflow_config_backup) and not airflow_version_update:
os.remove(airflow_config_backup)