def send_MIME_email(e_from, e_to, mime_msg, dryrun=False):
log = LoggingMixin().logger
SMTP_HOST = configuration.get('smtp', 'SMTP_HOST')
SMTP_PORT = configuration.getint('smtp', 'SMTP_PORT')
SMTP_STARTTLS = configuration.getboolean('smtp', 'SMTP_STARTTLS')
SMTP_SSL = configuration.getboolean('smtp', 'SMTP_SSL')
SMTP_USER = None
SMTP_PASSWORD = None
try:
SMTP_USER = configuration.get('smtp', 'SMTP_USER')
SMTP_PASSWORD = configuration.get('smtp', 'SMTP_PASSWORD')
except AirflowConfigException:
log.debug(
"No user/password found for SMTP, so logging in with no authentication."
)
if not dryrun:
s = smtplib.SMTP_SSL(SMTP_HOST,
SMTP_PORT) if SMTP_SSL else smtplib.SMTP(
SMTP_HOST, SMTP_PORT)
if SMTP_STARTTLS:
s.starttls()
if SMTP_USER and SMTP_PASSWORD:
s.login(SMTP_USER, SMTP_PASSWORD)
log.info("Sent an alert email to %s", e_to)
s.sendmail(e_from, e_to, mime_msg.as_string())
s.quit()
def execute_command(command):
log = LoggingMixin().logger
log.info("Executing command in Celery: %s", command)
try:
subprocess.check_call(command, shell=True)
except subprocess.CalledProcessError as e:
log.error(e)
raise AirflowException('Celery command failed')
def _split_tablename(table_input, default_project_id, var_name=None):
assert default_project_id is not None, "INTERNAL: No default project is specified"
def var_print(var_name):
if var_name is None:
return ""
else:
return "Format exception for {var}: ".format(var=var_name)
if table_input.count('.') + table_input.count(':') > 3:
raise Exception(('{var}Use either : or . to specify project '
'got {input}').format(var=var_print(var_name),
input=table_input))
cmpt = table_input.rsplit(':', 1)
project_id = None
rest = table_input
if len(cmpt) == 1:
project_id = None
rest = cmpt[0]
elif len(cmpt) == 2 and cmpt[0].count(':') <= 1:
if cmpt[-1].count('.') != 2:
project_id = cmpt[0]
rest = cmpt[1]
else:
raise Exception(('{var}Expect format of (<project:)<dataset>.<table>, '
'got {input}').format(var=var_print(var_name),
input=table_input))
cmpt = rest.split('.')
if len(cmpt) == 3:
assert project_id is None, (
"{var}Use either : or . to specify project").format(
var=var_print(var_name))
project_id = cmpt[0]
dataset_id = cmpt[1]
table_id = cmpt[2]
elif len(cmpt) == 2:
dataset_id = cmpt[0]
table_id = cmpt[1]
else:
raise Exception(
('{var}Expect format of (<project.|<project:)<dataset>.<table>, '
'got {input}').format(var=var_print(var_name), input=table_input))
if project_id is None:
if var_name is not None:
log = LoggingMixin().logger
log.info(
'Project not included in {var}: {input}; using project "{project}"'
.format(var=var_name,
input=table_input,
project=default_project_id))
project_id = default_project_id
return project_id, dataset_id, table_id
def GetDefaultExecutor():
"""Creates a new instance of the configured executor if none exists and returns it"""
global DEFAULT_EXECUTOR
if DEFAULT_EXECUTOR is not None:
return DEFAULT_EXECUTOR
executor_name = configuration.get('core', 'EXECUTOR')
DEFAULT_EXECUTOR = _get_executor(executor_name)
log = LoggingMixin().logger
log.info("Using executor %s", executor_name)
return DEFAULT_EXECUTOR
def handle_failure_retry(context):
ti = context['ti']
cmd_id = ti.xcom_pull(key='qbol_cmd_id', task_ids=ti.task_id)
if cmd_id is not None:
cmd = Command.find(cmd_id)
if cmd is not None:
log = LoggingMixin().logger
if cmd.status == 'done':
log.info(
'Command ID: %s has been succeeded, hence marking this '
'TI as Success.', cmd_id)
ti.state = State.SUCCESS
elif cmd.status == 'running':
log.info('Cancelling the Qubole Command Id: %s', cmd_id)
cmd.cancel()
def _poll_with_exponential_delay(request, max_n, is_done_func, is_error_func):
log = LoggingMixin().logger
for i in range(0, max_n):
try:
response = request.execute()
if is_error_func(response):
raise ValueError(
'The response contained an error: {}'.format(response))
elif is_done_func(response):
log.info('Operation is done: %s', response)
return response
else:
time.sleep((2**i) + (random.randint(0, 1000) / 1000))
except errors.HttpError as e:
if e.resp.status != 429:
log.info('Something went wrong. Not retrying: %s', format(e))
raise
else:
time.sleep((2**i) + (random.randint(0, 1000) / 1000))
engine_args = {}
if disable_connection_pool:
engine_args['poolclass'] = NullPool
elif 'sqlite' not in SQL_ALCHEMY_CONN:
# Engine args not supported by sqlite
engine_args['pool_size'] = conf.getint('core', 'SQL_ALCHEMY_POOL_SIZE')
engine_args['pool_recycle'] = conf.getint('core',
'SQL_ALCHEMY_POOL_RECYCLE')
engine = create_engine(SQL_ALCHEMY_CONN, **engine_args)
Session = scoped_session(
sessionmaker(autocommit=False, autoflush=False, bind=engine))
try:
from airflow_local_settings import *
log.info("Loaded airflow_local_settings.")
except:
pass
configure_logging()
configure_vars()
configure_orm()
# TODO: Unify airflow logging setups. Please see AIRFLOW-1457.
logging_config_path = conf.get('core', 'logging_config_path')
try:
from logging_config_path import LOGGING_CONFIG
log.debug("Successfully imported user-defined logging config.")
except Exception as e:
# Import default logging configurations.
log.debug(
def get_connection(cls, conn_id):
conn = random.choice(cls.get_connections(conn_id))
if conn.host:
log = LoggingMixin().logger
log.info("Using connection to: %s", conn.host)
return conn
all_vars = {k: v for d in [globals(), locals()] for k, v in d.items()}
return template.format(**all_vars)
TEST_CONFIG_FILE = AIRFLOW_HOME + '/unittests.cfg'
# only generate a Fernet key if we need to create a new config file
if not os.path.isfile(TEST_CONFIG_FILE) or not os.path.isfile(AIRFLOW_CONFIG):
FERNET_KEY = generate_fernet_key()
else:
FERNET_KEY = ''
TEMPLATE_START = (
'# ----------------------- TEMPLATE BEGINS HERE -----------------------')
if not os.path.isfile(TEST_CONFIG_FILE):
log.info('Creating new Airflow config file for unit tests in: %s',
TEST_CONFIG_FILE)
with open(TEST_CONFIG_FILE, 'w') as f:
cfg = parameterized_config(TEST_CONFIG)
f.write(cfg.split(TEMPLATE_START)[-1].strip())
if not os.path.isfile(AIRFLOW_CONFIG):
log.info('Creating new Airflow config file in: %s', AIRFLOW_CONFIG)
with open(AIRFLOW_CONFIG, 'w') as f:
cfg = parameterized_config(DEFAULT_CONFIG)
f.write(cfg.split(TEMPLATE_START)[-1].strip())
log.info("Reading the config from %s", AIRFLOW_CONFIG)
conf = AirflowConfigParser()
conf.read(AIRFLOW_CONFIG)
