Coverage for airflow.www.views : 72%

import sys 

import os 

import socket 

from functools import wraps 

from datetime import datetime, timedelta 

import dateutil.parser 

import copy 

from itertools import chain, product 

from past.utils import old_div 

from past.builtins import basestring 

import inspect 

import traceback 

import sqlalchemy as sqla 

from sqlalchemy import or_ 

from flask import redirect, url_for, request, Markup, Response, current_app, render_template 

from flask.ext.admin import BaseView, expose, AdminIndexView 

from flask.ext.admin.contrib.sqla import ModelView 

from flask.ext.admin.actions import action 

from flask.ext.login import current_user, flash, logout_user, login_required 

from flask._compat import PY2 

import jinja2 

import markdown 

import json 

from wtforms import ( 

    widgets, 

    Form, SelectField, TextAreaField, PasswordField, StringField) 

from pygments import highlight, lexers 

from pygments.formatters import HtmlFormatter 

import airflow 

from airflow import models 

from airflow.settings import Session 

from airflow import configuration 

from airflow import login 

from airflow import utils 

from airflow.utils import AirflowException 

from airflow.www import utils as wwwutils 

from airflow import settings 

from airflow.models import State 

from airflow.www.forms import DateTimeForm, TreeForm 

QUERY_LIMIT = 100000 

CHART_LIMIT = 200000 

dagbag = models.DagBag(os.path.expanduser(configuration.get('core', 'DAGS_FOLDER'))) 

login_required = airflow.login.login_required 

current_user = airflow.login.current_user 

logout_user = airflow.login.logout_user 

FILTER_BY_OWNER = False 

if configuration.getboolean('webserver', 'FILTER_BY_OWNER'): 

    # filter_by_owner if authentication is enabled and filter_by_owner is true 

    FILTER_BY_OWNER = not current_app.config['LOGIN_DISABLED'] 

def dag_link(v, c, m, p): 

    url = url_for( 

        'airflow.graph', 

        dag_id=m.dag_id) 

    return Markup( 

        '<a href="{url}">{m.dag_id}</a>'.format(**locals())) 

def log_link(v, c, m, p): 

    url = url_for( 

        'airflow.log', 

        dag_id=m.dag_id, 

        task_id=m.task_id, 

        execution_date=m.execution_date.isoformat()) 

    return Markup( 

        '<a href="{url}">' 

        '    <span class="glyphicon glyphicon-book" aria-hidden="true">' 

        '</span></a>').format(**locals()) 

def task_instance_link(v, c, m, p): 

    url = url_for( 

        'airflow.task', 

        dag_id=m.dag_id, 

        task_id=m.task_id, 

        execution_date=m.execution_date.isoformat()) 

    url_root = url_for( 

        'airflow.graph', 

        dag_id=m.dag_id, 

        root=m.task_id, 

        execution_date=m.execution_date.isoformat()) 

    return Markup( 

        """ 

        <span style="white-space: nowrap;"> 

        <a href="{url}">{m.task_id}</a> 

        <a href="{url_root}" title="Filter on this task and upstream"> 

        <span class="glyphicon glyphicon-filter" style="margin-left: 0px;" 

            aria-hidden="true"></span> 

        </a> 

        </span> 

        """.format(**locals())) 

def state_token(state): 

    color = State.color(state) 

    return Markup( 

        '<span class="label" style="background-color:{color};">' 

        '{state}</span>'.format(**locals())) 

def state_f(v, c, m, p): 

    return state_token(m.state) 

def duration_f(v, c, m, p): 

    if m.end_date and m.duration: 

        return timedelta(seconds=m.duration) 

def datetime_f(v, c, m, p): 

    attr = getattr(m, p) 

    dttm = attr.isoformat() if attr else '' 

    if datetime.now().isoformat()[:4] == dttm[:4]: 

        dttm = dttm[5:] 

    return Markup("<nobr>{}</nobr>".format(dttm)) 

def nobr_f(v, c, m, p): 

    return Markup("<nobr>{}</nobr>".format(getattr(m, p))) 

def label_link(v, c, m, p): 

    try: 

        default_params = eval(m.default_params) 

    except: 

        default_params = {} 

    url = url_for( 

        'airflow.chart', chart_id=m.id, iteration_no=m.iteration_no, 

        **default_params) 

    return Markup("<a href='{url}'>{m.label}</a>".format(**locals())) 

def pool_link(v, c, m, p): 

    url = '/admin/taskinstance/?flt1_pool_equals=' + m.pool 

    return Markup("<a href='{url}'>{m.pool}</a>".format(**locals())) 

def pygment_html_render(s, lexer=lexers.TextLexer): 

    return highlight( 

        s, 

        lexer(), 

        HtmlFormatter(linenos=True), 

    ) 

def render(obj, lexer): 

    out = "" 

    if isinstance(obj, basestring): 

        out += pygment_html_render(obj, lexer) 

    elif isinstance(obj, (tuple, list)): 

        for i, s in enumerate(obj): 

            out += "<div>List item #{}</div>".format(i) 

            out += "<div>" + pygment_html_render(s, lexer) + "</div>" 

    elif isinstance(obj, dict): 

        for k, v in obj.items(): 

            out += '<div>Dict item "{}"</div>'.format(k) 

            out += "<div>" + pygment_html_render(v, lexer) + "</div>" 

    return out 

def wrapped_markdown(s): 

    return '<div class="rich_doc">' + markdown.markdown(s) + "</div>" 

attr_renderer = { 

    'bash_command': lambda x: render(x, lexers.BashLexer), 

    'hql': lambda x: render(x, lexers.SqlLexer), 

    'sql': lambda x: render(x, lexers.SqlLexer), 

    'doc': lambda x: render(x, lexers.TextLexer), 

    'doc_json': lambda x: render(x, lexers.JsonLexer), 

    'doc_rst': lambda x: render(x, lexers.RstLexer), 

    'doc_yaml': lambda x: render(x, lexers.YamlLexer), 

    'doc_md': wrapped_markdown, 

    'python_callable': lambda x: render( 

        inspect.getsource(x), lexers.PythonLexer), 

} 

def data_profiling_required(f): 

    ''' 

    Decorator for views requiring data profiling access 

    ''' 

    @wraps(f) 

    def decorated_function(*args, **kwargs): 

        if ( 

                    current_app.config['LOGIN_DISABLED'] or 

                    (not current_user.is_anonymous() and current_user.data_profiling()) 

        ): 

            return f(*args, **kwargs) 

        else: 

            flash("This page requires data profiling privileges", "error") 

            return redirect(url_for('admin.index')) 

    return decorated_function 

def fused_slots(v, c, m, p): 

    url = ( 

        '/admin/taskinstance/' + 

        '?flt1_pool_equals=' + m.pool + 

        '&flt2_state_equals=running') 

    return Markup("<a href='{0}'>{1}</a>".format(url, m.used_slots())) 

def fqueued_slots(v, c, m, p): 

    url = ( 

        '/admin/taskinstance/' + 

        '?flt1_pool_equals=' + m.pool + 

        '&flt2_state_equals=queued&sort=10&desc=1') 

    return Markup("<a href='{0}'>{1}</a>".format(url, m.queued_slots())) 

class VisiblePasswordInput(widgets.PasswordInput): 

    def __init__(self, hide_value=False): 

        self.hide_value = hide_value 

class VisiblePasswordField(PasswordField): 

    widget = VisiblePasswordInput() 

class Airflow(BaseView): 

    def is_visible(self): 

        return False 

    @expose('/') 

    @login_required 

    def index(self): 

        return self.render('airflow/dags.html') 

    @expose('/chart_data') 

    @data_profiling_required 

    @wwwutils.gzipped 

    # @cache.cached(timeout=3600, key_prefix=wwwutils.make_cache_key) 

    def chart_data(self): 

        session = settings.Session() 

        chart_id = request.args.get('chart_id') 

        csv = request.args.get('csv') == "true" 

        chart = session.query(models.Chart).filter_by(id=chart_id).first() 

        db = session.query( 

            models.Connection).filter_by(conn_id=chart.conn_id).first() 

        session.expunge_all() 

        session.commit() 

        session.close() 

        payload = {} 

        payload['state'] = 'ERROR' 

        payload['error'] = '' 

        # Processing templated fields 

        try: 

            args = eval(chart.default_params) 

            if type(args) is not type(dict()): 

                raise AirflowException('Not a dict') 

        except: 

            args = {} 

            payload['error'] += ( 

                "Default params is not valid, string has to evaluate as " 

                "a Python dictionary. ") 

        request_dict = {k: request.args.get(k) for k in request.args} 

        from airflow import macros 

        args.update(request_dict) 

        args['macros'] = macros 

        sql = jinja2.Template(chart.sql).render(**args) 

        label = jinja2.Template(chart.label).render(**args) 

        payload['sql_html'] = Markup(highlight( 

            sql, 

            lexers.SqlLexer(),  # Lexer call 

            HtmlFormatter(noclasses=True)) 

        ) 

        payload['label'] = label 

        import pandas as pd 

        pd.set_option('display.max_colwidth', 100) 

        hook = db.get_hook() 

        try: 

            df = hook.get_pandas_df(wwwutils.limit_sql(sql, CHART_LIMIT, conn_type=db.conn_type)) 

            df = df.fillna(0) 

        except Exception as e: 

            payload['error'] += "SQL execution failed. Details: " + str(e) 

        if csv: 

            return Response( 

                response=df.to_csv(index=False), 

                status=200, 

                mimetype="application/text") 

        if not payload['error'] and len(df) == CHART_LIMIT: 

            payload['warning'] = ( 

                "Data has been truncated to {0}" 

                " rows. Expect incomplete results.").format(CHART_LIMIT) 

        if not payload['error'] and len(df) == 0: 

            payload['error'] += "Empty result set. " 

        elif ( 

                            not payload['error'] and 

                                chart.sql_layout == 'series' and 

                            chart.chart_type != "datatable" and 

                        len(df.columns) < 3): 

            payload['error'] += "SQL needs to return at least 3 columns. " 

        elif ( 

                        not payload['error'] and 

                            chart.sql_layout == 'columns'and 

                        len(df.columns) < 2): 

            payload['error'] += "SQL needs to return at least 2 columns. " 

        elif not payload['error']: 

            import numpy as np 

            chart_type = chart.chart_type 

            data = None 

            if chart_type == "datatable": 

                chart.show_datatable = True 

            if chart.show_datatable: 

                data = df.to_dict(orient="split") 

                data['columns'] = [{'title': c} for c in data['columns']] 

            # Trying to convert time to something Highcharts likes 

            x_col = 1 if chart.sql_layout == 'series' else 0 

            if chart.x_is_date: 

                try: 

                    # From string to datetime 

                    df[df.columns[x_col]] = pd.to_datetime( 

                        df[df.columns[x_col]]) 

                except Exception as e: 

                    raise AirflowException(str(e)) 

                df[df.columns[x_col]] = df[df.columns[x_col]].apply( 

                    lambda x: int(x.strftime("%s")) * 1000) 

            series = [] 

            colorAxis = None 

            if chart_type == 'datatable': 

                payload['data'] = data 

                payload['state'] = 'SUCCESS' 

                return Response( 

                    response=json.dumps( 

                        payload, indent=4, cls=utils.AirflowJsonEncoder), 

                    status=200, 

                    mimetype="application/json") 

            elif chart_type == 'para': 

                df.rename(columns={ 

                    df.columns[0]: 'name', 

                    df.columns[1]: 'group', 

                }, inplace=True) 

                return Response( 

                    response=df.to_csv(index=False), 

                    status=200, 

                    mimetype="application/text") 

            elif chart_type == 'heatmap': 

                color_perc_lbound = float( 

                    request.args.get('color_perc_lbound', 0)) 

                color_perc_rbound = float( 

                    request.args.get('color_perc_rbound', 1)) 

                color_scheme = request.args.get('color_scheme', 'blue_red') 

                if color_scheme == 'blue_red': 

                    stops = [ 

                        [color_perc_lbound, '#00D1C1'], 

                        [ 

                            color_perc_lbound + 

                            ((color_perc_rbound - color_perc_lbound)/2), 

                            '#FFFFCC' 

                        ], 

                        [color_perc_rbound, '#FF5A5F'] 

                    ] 

                elif color_scheme == 'blue_scale': 

                    stops = [ 

                        [color_perc_lbound, '#FFFFFF'], 

                        [color_perc_rbound, '#2222FF'] 

                    ] 

                elif color_scheme == 'fire': 

                    diff = float(color_perc_rbound - color_perc_lbound) 

                    stops = [ 

                        [color_perc_lbound, '#FFFFFF'], 

                        [color_perc_lbound + 0.33*diff, '#FFFF00'], 

                        [color_perc_lbound + 0.66*diff, '#FF0000'], 

                        [color_perc_rbound, '#000000'] 

                    ] 

                else: 

                    stops = [ 

                        [color_perc_lbound, '#FFFFFF'], 

                        [ 

                            color_perc_lbound + 

                            ((color_perc_rbound - color_perc_lbound)/2), 

                            '#888888' 

                        ], 

                        [color_perc_rbound, '#000000'], 

                    ] 

                xaxis_label = df.columns[1] 

                yaxis_label = df.columns[2] 

                data = [] 

                for row in df.itertuples(): 

                    data.append({ 

                        'x': row[2], 

                        'y': row[3], 

                        'value': row[4], 

                    }) 

                x_format = '{point.x:%Y-%m-%d}' \ 

                    if chart.x_is_date else '{point.x}' 

                series.append({ 

                    'data': data, 

                    'borderWidth': 0, 

                    'colsize': 24 * 36e5, 

                    'turboThreshold': sys.float_info.max, 

                    'tooltip': { 

                        'headerFormat': '', 

                        'pointFormat': ( 

                            df.columns[1] + ': ' + x_format + '<br/>' + 

                            df.columns[2] + ': {point.y}<br/>' + 

                            df.columns[3] + ': <b>{point.value}</b>' 

                        ), 

                    }, 

                }) 

                colorAxis = { 

                    'stops': stops, 

                    'minColor': '#FFFFFF', 

                    'maxColor': '#000000', 

                    'min': 50, 

                    'max': 2200, 

                } 

            else: 

                if chart.sql_layout == 'series': 

                    # User provides columns (series, x, y) 

                    xaxis_label = df.columns[1] 

                    yaxis_label = df.columns[2] 

                    df[df.columns[2]] = df[df.columns[2]].astype(np.float) 

                    df = df.pivot_table( 

                        index=df.columns[1], 

                        columns=df.columns[0], 

                        values=df.columns[2], aggfunc=np.sum) 

                else: 

                    # User provides columns (x, y, metric1, metric2, ...) 

                    xaxis_label = df.columns[0] 

                    yaxis_label = 'y' 

                    df.index = df[df.columns[0]] 

                    df = df.sort(df.columns[0]) 

                    del df[df.columns[0]] 

                    for col in df.columns: 

                        df[col] = df[col].astype(np.float) 

                for col in df.columns: 

                    series.append({ 

                        'name': col, 

                        'data': [ 

                            (k, df[col][k]) 

                            for k in df[col].keys() 

                            if not np.isnan(df[col][k])] 

                    }) 

                series = [serie for serie in sorted( 

                    series, key=lambda s: s['data'][0][1], reverse=True)] 

            if chart_type == "stacked_area": 

                stacking = "normal" 

                chart_type = 'area' 

            elif chart_type == "percent_area": 

                stacking = "percent" 

                chart_type = 'area' 

            else: 

                stacking = None 

            hc = { 

                'chart': { 

                    'type': chart_type 

                }, 

                'plotOptions': { 

                    'series': { 

                        'marker': { 

                            'enabled': False 

                        } 

                    }, 

                    'area': {'stacking': stacking}, 

                }, 

                'title': {'text': ''}, 

                'xAxis': { 

                    'title': {'text': xaxis_label}, 

                    'type': 'datetime' if chart.x_is_date else None, 

                }, 

                'yAxis': { 

                    'title': {'text': yaxis_label}, 

                }, 

                'colorAxis': colorAxis, 

                'tooltip': { 

                    'useHTML': True, 

                    'backgroundColor': None, 

                    'borderWidth': 0, 

                }, 

                'series': series, 

            } 

            if chart.y_log_scale: 

                hc['yAxis']['type'] = 'logarithmic' 

                hc['yAxis']['minorTickInterval'] = 0.1 

                if 'min' in hc['yAxis']: 

                    del hc['yAxis']['min'] 

            payload['state'] = 'SUCCESS' 

            payload['hc'] = hc 

            payload['data'] = data 

            payload['request_dict'] = request_dict 

        return Response( 

            response=json.dumps( 

                payload, indent=4, cls=utils.AirflowJsonEncoder),