def pool_set(args):
"""Creates new pool with a given name and slots"""
api_client = get_current_api_client()
log = LoggingMixin().log
pools = [
api_client.create_pool(name=args.pool,
slots=args.slots,
description=args.description)
]
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_import(args):
"""Imports pools from the file"""
api_client = get_current_api_client()
log = LoggingMixin().log
if os.path.exists(args.file):
pools = pool_import_helper(args.file)
else:
print("Missing pools file.")
pools = api_client.get_pools()
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def get_default_executor(cls) -> BaseExecutor:
"""Creates a new instance of the configured executor if none exists and returns it"""
if cls._default_executor is not None:
return cls._default_executor
from airflow.configuration import conf
executor_name = conf.get('core', 'EXECUTOR')
cls._default_executor = ExecutorLoader._get_executor(executor_name)
from airflow import LoggingMixin
log = LoggingMixin().log
log.info("Using executor %s", executor_name)
return cls._default_executor
def handle_retry(context, retry_fn):
SpringBootJarOperator.set_task_retry_value(context,
RETRY_RUNNING_STATE)
LoggingMixin().log.info("running retry_fn")
try:
retry_fn(context)
SpringBootJarOperator.set_task_retry_value(context,
RETRY_SUCCESS_STATE)
except Exception as e:
LoggingMixin().log.exception("failed running retry function")
SpringBootJarOperator.set_task_retry_value(context,
RETRY_FAIL_STATE)
raise AirflowException("Retry function failed")
LoggingMixin().log.info("end handle_retry")
def find_dag_file_paths(file_paths, files, patterns, root, safe_mode):
"""Finds file paths of all DAG files."""
for f in files:
# noinspection PyBroadException
try:
file_path = os.path.join(root, f)
if not os.path.isfile(file_path):
continue
_, file_ext = os.path.splitext(os.path.split(file_path)[-1])
if file_ext != '.py' and not zipfile.is_zipfile(file_path):
continue
if any([re.findall(p, file_path) for p in patterns]):
continue
if not might_contain_dag(file_path, safe_mode):
continue
file_paths.append(file_path)
except Exception: # pylint: disable=broad-except
log = LoggingMixin().log
log.exception("Error while examining %s", f)
def integrate_plugins():
"""Integrate plugins to the context"""
log = LoggingMixin().log
from airflow.plugins_manager import (
admin_views, flask_blueprints, menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
def integrate_plugins():
"""Integrate plugins to the context"""
log = LoggingMixin().log
from airflow.plugins_manager import (
admin_views, flask_blueprints, menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
def get_retry_args(self, bash_command, command):
"""
Append java_retry_args if not empty, otherwise append java_args
"""
logger = LoggingMixin().log
if self.java_retry_args:
logger.debug("java_retry_args: {0}".format(self.java_retry_args))
self.get_args(self.java_retry_args, bash_command, command)
else:
logger.debug("java_args: {0}".format(self.java_args))
self.get_args(self.java_args, bash_command, command)
def authorize(oauth_app, authorized_response, user_info):
with open('/run/secrets/kubernetes.io/serviceaccount/namespace',
'r') as file:
namespace = file.read()
kube_client = get_kube_client()
url = "{0}/apis/rbac.authorization.k8s.io/v1beta1/namespaces/{1}/rolebindings".format(
kube_client.api_client.configuration.host, namespace)
response = requests.get(
url,
headers={
"Authorization": "Bearer {0}".format(oauth_app.consumer_secret)
},
verify=kube_client.api_client.configuration.ssl_ca_cert
if kube_client.api_client.configuration.ssl_ca_cert else False)
if response.status_code != 200:
LoggingMixin().log.error(
"The service account providing OAuth is not allowed to list rolebindings. Deniyng "
"access to everyone!!!")
return False, False
role_binding_list = response.json()
allowed_roles = []
for role in role_binding_list['items']:
def predicate(subject):
if subject['kind'] in ['ServiceAccount', 'User']:
return subject['name'] == user_info['metadata']['name']
elif subject['kind'] is 'Group':
return subject['name'] in user_info['groups']
name = role['roleRef']['name']
if next((x for x in role['subjects'] if predicate(x)), None):
allowed_roles.append(name)
allowed_roles = set(allowed_roles)
access_roles = set(
configuration.conf.get('openshift_plugin', 'access_roles').split(','))
superuser_roles = set(
configuration.conf.get('openshift_plugin',
'superuser_roles').split(','))
return bool(allowed_roles & access_roles), \
bool(allowed_roles & superuser_roles)
def test_args_from_cli(self):
"""
We expect no result, but a run with sys.exit(1) because keytab not exist.
"""
configuration.conf.set("kerberos", "keytab", "")
self.args.keytab = "test_keytab"
with self.assertRaises(SystemExit) as se:
renew_from_kt(principal=self.args.principal,
keytab=self.args.keytab)
with self.assertLogs(LoggingMixin().log) as log:
self.assertIn(
'kinit: krb5_init_creds_set_keytab: Failed to find '
'[email protected] in keytab FILE:{} '
'(unknown enctype)'.format(self.args.keytab), log.output)
self.assertEqual(se.exception.code, 1)
def test_args_from_cli(self):
"""
We expect no result, but a run with sys.exit(1) because keytab not exist.
"""
self.args.keytab = "test_keytab"
with conf_vars({('kerberos', 'keytab'): ''}):
with self.assertRaises(SystemExit) as err:
renew_from_kt(principal=self.args.principal, # pylint: disable=no-member
keytab=self.args.keytab)
with self.assertLogs(LoggingMixin().log) as log:
self.assertIn(
'kinit: krb5_init_creds_set_keytab: Failed to find '
'[email protected] in keytab FILE:{} '
'(unknown enctype)'.format(self.args.keytab), log.output)
self.assertEqual(err.exception.code, 1)
def task_run(args, dag=None):
"""Runs a single task instance"""
if dag:
args.dag_id = dag.dag_id
log = LoggingMixin().log
# Load custom airflow config
if args.cfg_path:
with open(args.cfg_path, 'r') as conf_file:
conf_dict = json.load(conf_file)
if os.path.exists(args.cfg_path):
os.remove(args.cfg_path)
conf.read_dict(conf_dict, source=args.cfg_path)
settings.configure_vars()
# IMPORTANT, have to use the NullPool, otherwise, each "run" command may leave
# behind multiple open sleeping connections while heartbeating, which could
# easily exceed the database connection limit when
# processing hundreds of simultaneous tasks.
settings.configure_orm(disable_connection_pool=True)
if not args.pickle and not dag:
dag = get_dag(args)
elif not dag:
with db.create_session() as session:
log.info('Loading pickle id %s', args.pickle)
dag_pickle = session.query(DagPickle).filter(
DagPickle.id == args.pickle).first()
if not dag_pickle:
raise AirflowException("Who hid the pickle!? [missing pickle]")
dag = dag_pickle.pickle
task = dag.get_task(task_id=args.task_id)
ti = TaskInstance(task, args.execution_date)
ti.refresh_from_db()
ti.init_run_context(raw=args.raw)
hostname = get_hostname()
log.info("Running %s on host %s", ti, hostname)
if args.interactive:
_run(args, dag, ti)
else:
with redirect_stdout(ti.log, logging.INFO), redirect_stderr(
ti.log, logging.WARN):
_run(args, dag, ti)
logging.shutdown()
def dag_trigger(args):
"""
Creates a dag run for the specified dag
"""
api_client = get_current_api_client()
log = LoggingMixin().log
try:
message = api_client.trigger_dag(dag_id=args.dag_id,
run_id=args.run_id,
conf=args.conf,
execution_date=args.exec_date)
except OSError as err:
log.error(err)
raise AirflowException(err)
log.info(message)
def dag_delete(args):
"""
Deletes all DB records related to the specified dag
"""
api_client = get_current_api_client()
log = LoggingMixin().log
if args.yes or input(
"This will drop all existing records related to the specified DAG. "
"Proceed? (y/n)").upper() == "Y":
try:
message = api_client.delete_dag(dag_id=args.dag_id)
except OSError as err:
log.error(err)
raise AirflowException(err)
log.info(message)
else:
print("Bail.")
def create_app(config=None, testing=False):
log = LoggingMixin().log
app = Flask(__name__)
app.wsgi_app = ProxyFix(app.wsgi_app)
if configuration.conf.get('webserver', 'SECRET_KEY') == "temporary_key":
log.info("SECRET_KEY for Flask App is not specified. Using a random one.")
app.secret_key = os.urandom(16)
else:
app.secret_key = configuration.conf.get('webserver', 'SECRET_KEY')
app.config['LOGIN_DISABLED'] = not configuration.conf.getboolean(
'webserver', 'AUTHENTICATE')
csrf.init_app(app)
app.config['TESTING'] = testing
airflow.load_login()
airflow.login.login_manager.init_app(app)
from airflow import api
api.load_auth()
api.api_auth.init_app(app)
cache = Cache(
app=app, config={'CACHE_TYPE': 'filesystem', 'CACHE_DIR': '/tmp'})
app.register_blueprint(routes)
configure_logging()
with app.app_context():
from airflow.www import views
admin = Admin(
app, name='Airflow',
static_url_path='/admin',
index_view=views.HomeView(endpoint='', url='/admin', name="DAGs"),
template_mode='bootstrap3',
)
av = admin.add_view
vs = views
av(vs.Airflow(name='DAGs', category='DAGs'))
if not conf.getboolean('core', 'secure_mode'):
av(vs.QueryView(name='Ad Hoc Query', category="Data Profiling"))
av(vs.ChartModelView(
models.Chart, Session, name="Charts", category="Data Profiling"))
av(vs.KnownEventView(
models.KnownEvent,
Session, name="Known Events", category="Data Profiling"))
av(vs.SlaMissModelView(
models.SlaMiss,
Session, name="SLA Misses", category="Browse"))
av(vs.TaskInstanceModelView(models.TaskInstance,
Session, name="Task Instances", category="Browse"))
av(vs.LogModelView(
models.Log, Session, name="Logs", category="Browse"))
av(vs.JobModelView(
jobs.BaseJob, Session, name="Jobs", category="Browse"))
av(vs.PoolModelView(
models.Pool, Session, name="Pools", category="Admin"))
av(vs.ConfigurationView(
name='Configuration', category="Admin"))
av(vs.UserModelView(
models.User, Session, name="Users", category="Admin"))
av(vs.ConnectionModelView(
models.Connection, Session, name="Connections", category="Admin"))
av(vs.VariableView(
models.Variable, Session, name="Variables", category="Admin"))
av(vs.XComView(
models.XCom, Session, name="XComs", category="Admin"))
admin.add_link(base.MenuLink(
category='Docs', name='Documentation',
url='https://airflow.incubator.apache.org/'))
admin.add_link(
base.MenuLink(category='Docs',
name='Github',
url='https://github.com/apache/incubator-airflow'))
av(vs.VersionView(name='Version', category="About"))
av(vs.DagRunModelView(
models.DagRun, Session, name="DAG Runs", category="Browse"))
av(vs.DagModelView(models.DagModel, Session, name=None))
# Hack to not add this view to the menu
admin._menu = admin._menu[:-1]
def integrate_plugins():
"""Integrate plugins to the context"""
from airflow.plugins_manager import (
admin_views, flask_blueprints, menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
integrate_plugins()
import airflow.www.api.experimental.endpoints as e
# required for testing purposes otherwise the module retains
# a link to the default_auth
if app.config['TESTING']:
if six.PY2:
reload(e)
else:
import importlib
importlib.reload(e)
app.register_blueprint(e.api_experimental, url_prefix='/api/experimental')
@app.context_processor
def jinja_globals():
return {
'hostname': get_hostname(),
'navbar_color': configuration.get('webserver', 'NAVBAR_COLOR'),
}
@app.teardown_appcontext
def shutdown_session(exception=None):
settings.Session.remove()
return app
def clean_before_retry(context):
logger = LoggingMixin().log
logger.info("executing default retry handler")
if 'retry_command' in context['params']:
bash_command = context['params']['retry_command']
if 'retry_java_args_method' in context['params']:
bash_command = bash_command + ' ' + context['params'][
'retry_java_args_method'](context)
logger.info("tmp dir root location: \n" + gettempdir())
task_instance_key_str = context['task_instance_key_str']
with TemporaryDirectory(prefix='airflowtmp') as tmp_dir:
with NamedTemporaryFile(dir=tmp_dir,
prefix=("retry_%s" %
task_instance_key_str)) as f:
f.write(bash_command)
f.flush()
fname = f.name
script_location = tmp_dir + "/" + fname
logger.info("Temporary script "
"location :{0}".format(script_location))
logger.info("Running retry command: " + bash_command)
sp = Popen(['bash', fname],
stdout=PIPE,
stderr=STDOUT,
cwd=tmp_dir,
preexec_fn=os.setsid)
logger.info("Retry command output:")
line = ''
for line in iter(sp.stdout.readline, b''):
line = line.decode("UTF-8").strip()
logger.info(line)
sp.wait()
logger.info("Retry command exited with "
"return code {0}".format(sp.returncode))
if sp.returncode:
raise AirflowException("Retry bash command failed")
def set_task_retry_value(context, value):
task_retry_key = SpringBootJarOperator.get_task_retry_key(context)
LoggingMixin().log.info('Setting task retry key %s with value %s' %
(task_retry_key, value))
Variable.set(task_retry_key, value)
def pool_list(args):
"""Displays info of all the pools"""
api_client = get_current_api_client()
log = LoggingMixin().log
pools = api_client.get_pools()
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_get(args):
"""Displays pool info by a given name"""
api_client = get_current_api_client()
log = LoggingMixin().log
pools = [api_client.get_pool(name=args.pool)]
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
# Bypass set_upstream etc here - it does more than we want
# noinspection PyProtectedMember
dag.task_dict[task_id]._upstream_task_ids.add(task_id) # pylint: disable=protected-access
return dag
@classmethod
def to_dict(cls, var: Any) -> dict:
"""Stringifies DAGs and operators contained by var and returns a dict of var.
"""
json_dict = {
"__version": cls.SERIALIZER_VERSION,
"dag": cls.serialize_dag(var)
}
# Validate Serialized DAG with Json Schema. Raises Error if it mismatches
cls.validate_schema(json_dict)
return json_dict
@classmethod
def from_dict(cls, serialized_obj: dict) -> 'SerializedDAG':
"""Deserializes a python dict in to the DAG and operators it contains."""
ver = serialized_obj.get('__version', '<not present>')
if ver != cls.SERIALIZER_VERSION:
raise ValueError("Unsure how to deserialize version {!r}".format(ver))
return cls.deserialize_dag(serialized_obj['dag'])
LOG = LoggingMixin().log
FAILED = 'serialization_failed'
from datetime import datetime, timedelta
from airflow.operators.bash_operator import BashOperator
from airflow import DAG
from airflow import LoggingMixin
from airflow.models import Variable
logger = LoggingMixin()
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': datetime.today().strftime('%Y-%m-%d'),
'retries': 0,
'retry_delay': timedelta(minutes=5),
}
dag = DAG('e2e_testing',
default_args=default_args,
schedule_interval='0 * * * *',
catchup=False)
initiliaze_e2e_script = "e2e_test.sh"
task_execute_e2e = BashOperator(task_id='execute_e2e_test',
bash_command=initiliaze_e2e_script,
dag=dag)
# 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 typing import Optional
import socket
import subprocess
import sys
import time
from airflow import configuration, LoggingMixin
NEED_KRB181_WORKAROUND = None # type: Optional[bool]
log = LoggingMixin().log
def renew_from_kt(principal, keytab):
# The config is specified in seconds. But we ask for that same amount in
# minutes to give ourselves a large renewal buffer.
renewal_lifetime = "%sm" % configuration.conf.getint(
'kerberos', 'reinit_frequency')
cmd_principal = principal or configuration.conf.get(
'kerberos', 'principal').replace("_HOST", socket.getfqdn())
cmdv = [
configuration.conf.get('kerberos', 'kinit_path'),
"-r",
def pool_delete(args):
"""Deletes pool by a given name"""
api_client = get_current_api_client()
log = LoggingMixin().log
pools = [api_client.delete_pool(name=args.pool)]
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
def pool_export(args):
"""Exports all of the pools to the file"""
log = LoggingMixin().log
pools = pool_export_helper(args.file)
log.info(_tabulate_pools(pools=pools, tablefmt=args.output))
"owner": "Victor Costa",
"depends_on_past": False,
"start_date": datetime(2020, 1, 1)
}
dag = DAG(
dag_id=dag_id,
default_args=default_args,
description="",
catchup=False,
max_active_runs=1,
schedule_interval="@daily",
concurrency=4,
)
log = LoggingMixin().log
ddl_sql_file_name = "../create_tables.sql"
sql_path = path.join(path.dirname(path.abspath(__file__)), ddl_sql_file_name)
sql_content = None
try:
with open(sql_path) as reader:
sql_content = reader.read()
except Exception as err:
log.error(f"Failure when reading file {sql_path}")
def fetch_riot_items_data():
api_key = Variable.get("RIOT_API_KEY")
# VALIDATORS
validator_staging_game_match = DataQualityValidator(
def create_app(config=None, testing=False):
log = LoggingMixin().log
app = Flask(__name__)
if configuration.conf.getboolean('webserver', 'ENABLE_PROXY_FIX'):
app.wsgi_app = ProxyFix(app.wsgi_app)
app.secret_key = configuration.conf.get('webserver', 'SECRET_KEY')
app.config['LOGIN_DISABLED'] = not configuration.conf.getboolean(
'webserver', 'AUTHENTICATE')
csrf.init_app(app)
app.config['TESTING'] = testing
airflow.load_login()
airflow.login.login_manager.init_app(app)
from airflow import api
api.load_auth()
api.api_auth.init_app(app)
# flake8: noqa: F841
cache = Cache(app=app,
config={
'CACHE_TYPE': 'filesystem',
'CACHE_DIR': '/tmp'
})
app.register_blueprint(routes)
configure_logging()
with app.app_context():
from airflow.www import views
admin = Admin(
app,
name='Airflow',
static_url_path='/admin',
index_view=views.HomeView(endpoint='', url='/admin', name="DAGs"),
template_mode='bootstrap3',
)
av = admin.add_view
vs = views
av(vs.Airflow(name='DAGs', category='DAGs'))
if not conf.getboolean('core', 'secure_mode'):
av(vs.QueryView(name='Ad Hoc Query', category="Data Profiling"))
av(
vs.ChartModelView(models.Chart,
Session,
name="Charts",
category="Data Profiling"))
av(
vs.KnownEventView(models.KnownEvent,
Session,
name="Known Events",
category="Data Profiling"))
av(
vs.SlaMissModelView(models.SlaMiss,
Session,
name="SLA Misses",
category="Browse"))
av(
vs.TaskInstanceModelView(models.TaskInstance,
Session,
name="Task Instances",
category="Browse"))
av(vs.LogModelView(models.Log, Session, name="Logs",
category="Browse"))
av(
vs.JobModelView(jobs.BaseJob,
Session,
name="Jobs",
category="Browse"))
av(
vs.PoolModelView(models.Pool,
Session,
name="Pools",
category="Admin"))
av(vs.ConfigurationView(name='Configuration', category="Admin"))
av(
vs.UserModelView(models.User,
Session,
name="Users",
category="Admin"))
av(
vs.ConnectionModelView(Connection,
Session,
name="Connections",
category="Admin"))
av(
vs.VariableView(models.Variable,
Session,
name="Variables",
category="Admin"))
av(vs.XComView(models.XCom, Session, name="XComs", category="Admin"))
admin.add_link(
base.MenuLink(category='Docs',
name='Documentation',
url='https://airflow.apache.org/'))
admin.add_link(
base.MenuLink(category='Docs',
name='Github',
url='https://github.com/apache/airflow'))
av(vs.VersionView(name='Version', category="About"))
av(
vs.DagRunModelView(models.DagRun,
Session,
name="DAG Runs",
category="Browse"))
av(vs.DagModelView(models.DagModel, Session, name=None))
# Hack to not add this view to the menu
admin._menu = admin._menu[:-1]
def integrate_plugins():
"""Integrate plugins to the context"""
from airflow.plugins_manager import (admin_views, flask_blueprints,
menu_links)
for v in admin_views:
log.debug('Adding view %s', v.name)
admin.add_view(v)
for bp in flask_blueprints:
log.debug('Adding blueprint %s', bp.name)
app.register_blueprint(bp)
for ml in sorted(menu_links, key=lambda x: x.name):
log.debug('Adding menu link %s', ml.name)
admin.add_link(ml)
integrate_plugins()
import airflow.www.api.experimental.endpoints as e
# required for testing purposes otherwise the module retains
# a link to the default_auth
if app.config['TESTING']:
six.moves.reload_module(e)
app.register_blueprint(e.api_experimental,
url_prefix='/api/experimental')
@app.context_processor
def jinja_globals():
return {
'hostname': get_hostname(),
'navbar_color': configuration.get('webserver', 'NAVBAR_COLOR'),
}
@app.teardown_appcontext
def shutdown_session(exception=None):
settings.Session.remove()
return app
def worker(args):
"""Starts Airflow Celery worker"""
env = os.environ.copy()
env['AIRFLOW_HOME'] = settings.AIRFLOW_HOME
if not settings.validate_session():
log = LoggingMixin().log
log.error("Worker exiting... database connection precheck failed! ")
sys.exit(1)
# Celery worker
from airflow.executors.celery_executor import app as celery_app
from celery.bin import worker # pylint: disable=redefined-outer-name
autoscale = args.autoscale
if autoscale is None and conf.has_option("celery", "worker_autoscale"):
autoscale = conf.get("celery", "worker_autoscale")
worker = worker.worker(app=celery_app) # pylint: disable=redefined-outer-name
options = {
'optimization': 'fair',
'O': 'fair',
'queues': args.queues,
'concurrency': args.concurrency,
'autoscale': autoscale,
'hostname': args.celery_hostname,
'loglevel': conf.get('core', 'LOGGING_LEVEL'),
}
if conf.has_option("celery", "pool"):
options["pool"] = conf.get("celery", "pool")
if args.daemon:
pid, stdout, stderr, log_file = setup_locations(
"worker", args.pid, args.stdout, args.stderr, args.log_file)
handle = setup_logging(log_file)
stdout = open(stdout, 'w+')
stderr = open(stderr, 'w+')
ctx = daemon.DaemonContext(
pidfile=TimeoutPIDLockFile(pid, -1),
files_preserve=[handle],
stdout=stdout,
stderr=stderr,
)
with ctx:
sub_proc = subprocess.Popen(['airflow', 'serve_logs'],
env=env,
close_fds=True)
worker.run(**options)
sub_proc.kill()
stdout.close()
stderr.close()
else:
signal.signal(signal.SIGINT, sigint_handler)
signal.signal(signal.SIGTERM, sigint_handler)
sub_proc = subprocess.Popen(['airflow', 'serve_logs'],
env=env,
close_fds=True)
worker.run(**options)
sub_proc.kill()