#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software 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.

import json
import time
from urllib.parse import urlencode

import flask_appbuilder.models.sqla.filters as fab_sqlafilters
import markdown
import sqlalchemy as sqla
from flask import Markup, Response, request, url_for
from flask_appbuilder.forms import DateTimeField, FieldConverter
from flask_appbuilder.models.sqla.interface import SQLAInterface
from pygments import highlight, lexers
from pygments.formatters import HtmlFormatter

from airflow.configuration import conf
from airflow.models.baseoperator import BaseOperator
from airflow.models.dag import DagBag, DagModel
from airflow.operators.subdag_operator import SubDagOperator
from airflow.utils import timezone
from airflow.utils.code_utils import get_python_source
from airflow.utils.json import AirflowJsonEncoder
from airflow.utils.state import State
from airflow.www.widgets import AirflowDateTimePickerWidget

DEFAULT_SENSITIVE_VARIABLE_FIELDS = (
    'password',
    'secret',
    'passwd',
    'authorization',
    'api_key',
    'apikey',
    'access_token',
)


def should_hide_value_for_key(key_name):
    # It is possible via importing variables from file that a key is empty.
    if key_name:
        config_set = conf.getboolean('admin',
                                     'hide_sensitive_variable_fields')
        field_comp = any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS)
        return config_set and field_comp
    return False


def get_params(**kwargs):
    hide_paused_dags_by_default = conf.getboolean('webserver',
                                                  'hide_paused_dags_by_default')
    if 'showPaused' in kwargs:
        show_paused_dags_url_param = kwargs['showPaused']
        if _should_remove_show_paused_from_url_params(
            show_paused_dags_url_param,
            hide_paused_dags_by_default
        ):
            kwargs.pop('showPaused')
    return urlencode({d: v if v is not None else '' for d, v in kwargs.items()})


def _should_remove_show_paused_from_url_params(show_paused_dags_url_param,
                                               hide_paused_dags_by_default):
    return any([
        show_paused_dags_url_param != hide_paused_dags_by_default,
        show_paused_dags_url_param is None
    ])


def generate_pages(current_page, num_of_pages,
                   search=None, showPaused=None, window=7):
    """
    Generates the HTML for a paging component using a similar logic to the paging
    auto-generated by Flask managed views. The paging component defines a number of
    pages visible in the pager (window) and once the user goes to a page beyond the
    largest visible, it would scroll to the right the page numbers and keeps the
    current one in the middle of the pager component. When in the last pages,
    the pages won't scroll and just keep moving until the last page. Pager also contains
    <first, previous, ..., next, last> pages.
    This component takes into account custom parameters such as search and showPaused,
    which could be added to the pages link in order to maintain the state between
    client and server. It also allows to make a bookmark on a specific paging state.

    :param current_page: the current page number, 0-indexed
    :param num_of_pages: the total number of pages
    :param search: the search query string, if any
    :param showPaused: false if paused dags will be hidden, otherwise true to show them
    :param window: the number of pages to be shown in the paging component (7 default)
    :return: the HTML string of the paging component
    """

    void_link = 'javascript:void(0)'
    first_node = Markup("""<li class="paginate_button {disabled}" id="dags_first">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="0" tabindex="0">&laquo;</a>
</li>""")

    previous_node = Markup("""<li class="paginate_button previous {disabled}" id="dags_previous">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="0" tabindex="0">&lsaquo;</a>
</li>""")

    next_node = Markup("""<li class="paginate_button next {disabled}" id="dags_next">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="3" tabindex="0">&rsaquo;</a>
</li>""")

    last_node = Markup("""<li class="paginate_button {disabled}" id="dags_last">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="3" tabindex="0">&raquo;</a>
</li>""")

    page_node = Markup("""<li class="paginate_button {is_active}">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="2" tabindex="0">{page_num}</a>
</li>""")

    output = [Markup('<ul class="pagination" style="margin-top:0px;">')]

    is_disabled = 'disabled' if current_page <= 0 else ''
    output.append(first_node.format(href_link="?{}"
                                    .format(get_params(page=0,
                                                       search=search,
                                                       showPaused=showPaused)),
                                    disabled=is_disabled))

    page_link = void_link
    if current_page > 0:
        page_link = '?{}'.format(get_params(page=(current_page - 1),
                                            search=search,
                                            showPaused=showPaused))

    output.append(previous_node.format(href_link=page_link,
                                       disabled=is_disabled))

    mid = int(window / 2)
    last_page = num_of_pages - 1

    if current_page <= mid or num_of_pages < window:
        pages = [i for i in range(0, min(num_of_pages, window))]
    elif mid < current_page < last_page - mid:
        pages = [i for i in range(current_page - mid, current_page + mid + 1)]
    else:
        pages = [i for i in range(num_of_pages - window, last_page + 1)]

    def is_current(current, page):
        return page == current

    for page in pages:
        vals = {
            'is_active': 'active' if is_current(current_page, page) else '',
            'href_link': void_link if is_current(current_page, page)
                         else '?{}'.format(get_params(page=page,
                                                      search=search,
                                                      showPaused=showPaused)),
            'page_num': page + 1
        }
        output.append(page_node.format(**vals))

    is_disabled = 'disabled' if current_page >= num_of_pages - 1 else ''

    page_link = (void_link if current_page >= num_of_pages - 1
                 else '?{}'.format(get_params(page=current_page + 1,
                                              search=search,
                                              showPaused=showPaused)))

    output.append(next_node.format(href_link=page_link, disabled=is_disabled))
    output.append(last_node.format(href_link="?{}"
                                   .format(get_params(page=last_page,
                                                      search=search,
                                                      showPaused=showPaused)),
                                   disabled=is_disabled))

    output.append(Markup('</ul>'))

    return Markup('\n'.join(output))


def epoch(dttm):
    """Returns an epoch-type date"""
    return int(time.mktime(dttm.timetuple())) * 1000,


def json_response(obj):
    """
    returns a json response from a json serializable python object
    """
    return Response(
        response=json.dumps(
            obj, indent=4, cls=AirflowJsonEncoder),
        status=200,
        mimetype="application/json")


def make_cache_key(*args, **kwargs):
    """
    Used by cache to get a unique key per URL
    """
    path = request.path
    args = str(hash(frozenset(request.args.items())))
    return (path + args).encode('ascii', 'ignore')


def task_instance_link(attr):
    dag_id = attr.get('dag_id')
    task_id = attr.get('task_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.task',
        dag_id=dag_id,
        task_id=task_id,
        execution_date=execution_date.isoformat())
    url_root = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        root=task_id,
        execution_date=execution_date.isoformat())
    return Markup(
        """
        <span style="white-space: nowrap;">
        <a href="{url}">{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(url=url, task_id=task_id, url_root=url_root)


def state_token(state):
    color = State.color(state)
    return Markup(
        '<span class="label" style="background-color:{color};">'
        '{state}</span>').format(color=color, state=state)


def state_f(attr):
    state = attr.get('state')
    return state_token(state)


def nobr_f(attr_name):
    def nobr(attr):
        f = attr.get(attr_name)
        return Markup("<nobr>{}</nobr>").format(f)
    return nobr


def datetime_f(attr_name):
    def dt(attr):
        f = attr.get(attr_name)
        as_iso = f.isoformat() if f else ''
        if not as_iso:
            return Markup('')
        f = as_iso
        if timezone.utcnow().isoformat()[:4] == f[:4]:
            f = f[5:]
        # The empty title will be replaced in JS code when non-UTC dates are displayed
        return Markup('<nobr><time title="" datetime="{}">{}</time></nobr>').format(as_iso, f)
    return dt


def dag_link(attr):
    dag_id = attr.get('dag_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        execution_date=execution_date)
    return Markup(
        '<a href="{}">{}</a>').format(url, dag_id)


def dag_run_link(attr):
    dag_id = attr.get('dag_id')
    run_id = attr.get('run_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        run_id=run_id,
        execution_date=execution_date)
    return Markup(
        '<a href="{url}">{run_id}</a>').format(url=url, run_id=run_id)


def pygment_html_render(s, lexer=lexers.TextLexer):
    return highlight(
        s,
        lexer(),
        HtmlFormatter(linenos=True),
    )


def render(obj, lexer):
    out = ""
    if isinstance(obj, str):
        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>"
        if s is not None
        else None
    )


def get_attr_renderer():
    return {
        '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(get_python_source(x), lexers.PythonLexer),
    }


def recurse_tasks(tasks, task_ids, dag_ids, task_id_to_dag):
    if isinstance(tasks, list):
        for task in tasks:
            recurse_tasks(task, task_ids, dag_ids, task_id_to_dag)
        return
    if isinstance(tasks, SubDagOperator):
        subtasks = tasks.subdag.tasks
        dag_ids.append(tasks.subdag.dag_id)
        for subtask in subtasks:
            if subtask.task_id not in task_ids:
                task_ids.append(subtask.task_id)
                task_id_to_dag[subtask.task_id] = tasks.subdag
        recurse_tasks(subtasks, task_ids, dag_ids, task_id_to_dag)
    if isinstance(tasks, BaseOperator):
        task_id_to_dag[tasks.task_id] = tasks.dag


def get_chart_height(dag):
    """
    TODO(aoen): See [AIRFLOW-1263] We use the number of tasks in the DAG as a heuristic to
    approximate the size of generated chart (otherwise the charts are tiny and unreadable
    when DAGs have a large number of tasks). Ideally nvd3 should allow for dynamic-height
    charts, that is charts that take up space based on the size of the components within.
    """
    return 600 + len(dag.tasks) * 10


class UtcAwareFilterMixin:
    def apply(self, query, value):
        value = timezone.parse(value, timezone=timezone.utc)

        return super().apply(query, value)


class UtcAwareFilterEqual(UtcAwareFilterMixin, fab_sqlafilters.FilterEqual):
    pass


class UtcAwareFilterGreater(UtcAwareFilterMixin, fab_sqlafilters.FilterGreater):
    pass


class UtcAwareFilterSmaller(UtcAwareFilterMixin, fab_sqlafilters.FilterSmaller):
    pass


class UtcAwareFilterNotEqual(UtcAwareFilterMixin, fab_sqlafilters.FilterNotEqual):
    pass


class UtcAwareFilterConverter(fab_sqlafilters.SQLAFilterConverter):

    conversion_table = (
        (('is_utcdatetime', [UtcAwareFilterEqual,
                             UtcAwareFilterGreater,
                             UtcAwareFilterSmaller,
                             UtcAwareFilterNotEqual]),) +
        fab_sqlafilters.SQLAFilterConverter.conversion_table
    )


class CustomSQLAInterface(SQLAInterface):
    """
    FAB does not know how to handle columns with leading underscores because
    they are not supported by WTForm. This hack will remove the leading
    '_' from the key to lookup the column names.

    """
    def __init__(self, obj):
        super().__init__(obj)

        def clean_column_names():
            if self.list_properties:
                self.list_properties = {
                    k.lstrip('_'): v for k, v in self.list_properties.items()}
            if self.list_columns:
                self.list_columns = {
                    k.lstrip('_'): v for k, v in self.list_columns.items()}

        clean_column_names()

    def is_utcdatetime(self, col_name):
        from airflow.utils.sqlalchemy import UtcDateTime

        if col_name in self.list_columns:
            obj = self.list_columns[col_name].type
            return isinstance(obj, UtcDateTime) or \
                isinstance(obj, sqla.types.TypeDecorator) and \
                isinstance(obj.impl, UtcDateTime)
        return False

    filter_converter_class = UtcAwareFilterConverter


# This class is used directly (i.e. we cant tell Fab to use a different
# subclass) so we have no other option than to edit the converstion table in
# place
FieldConverter.conversion_table = (
    (('is_utcdatetime', DateTimeField, AirflowDateTimePickerWidget),) +
    FieldConverter.conversion_table
)


def get_dag(orm_dag: DagModel, store_serialized_dags=False):
    """Creates a dagbag to load and return a DAG.

    Calling it from UI should set store_serialized_dags = STORE_SERIALIZED_DAGS.
    There may be a delay for scheduler to write serialized DAG into database,
    loads from file in this case.
    FIXME: remove it when webserver does not access to DAG folder in future.
    """
    dag = DagBag(
        dag_folder=orm_dag.fileloc, store_serialized_dags=store_serialized_dags).get_dag(orm_dag.dag_id)
    if store_serialized_dags and dag is None:
        dag = DagBag(
            dag_folder=orm_dag.fileloc, store_serialized_dags=False).get_dag(orm_dag.dag_id)
    return dag

airflow/www/utils.py
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software 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.

import json
import time
from urllib.parse import urlencode

import flask_appbuilder.models.sqla.filters as fab_sqlafilters
import markdown
import sqlalchemy as sqla
from flask import Markup, Response, request, url_for
from flask_appbuilder.forms import DateTimeField, FieldConverter
from flask_appbuilder.models.sqla.interface import SQLAInterface
from pygments import highlight, lexers
from pygments.formatters import HtmlFormatter

from airflow.configuration import conf
from airflow.models.baseoperator import BaseOperator
from airflow.models.dag import DagBag, DagModel
from airflow.operators.subdag_operator import SubDagOperator
from airflow.utils import timezone
from airflow.utils.code_utils import get_python_source
from airflow.utils.json import AirflowJsonEncoder
from airflow.utils.state import State
from airflow.www.widgets import AirflowDateTimePickerWidget

DEFAULT_SENSITIVE_VARIABLE_FIELDS = (
    'password',
    'secret',
    'passwd',
    'authorization',
    'api_key',
    'apikey',
    'access_token',
)


def should_hide_value_for_key(key_name):
    # It is possible via importing variables from file that a key is empty.
    if key_name:
        config_set = conf.getboolean('admin',
                                     'hide_sensitive_variable_fields')
        field_comp = any(s in key_name.lower() for s in DEFAULT_SENSITIVE_VARIABLE_FIELDS)
        return config_set and field_comp
    return False


def get_params(**kwargs):
    hide_paused_dags_by_default = conf.getboolean('webserver',
                                                  'hide_paused_dags_by_default')
    if 'showPaused' in kwargs:
        show_paused_dags_url_param = kwargs['showPaused']
        if _should_remove_show_paused_from_url_params(
            show_paused_dags_url_param,
            hide_paused_dags_by_default
        ):
            kwargs.pop('showPaused')
    return urlencode({d: v if v is not None else '' for d, v in kwargs.items()})


def _should_remove_show_paused_from_url_params(show_paused_dags_url_param,
                                               hide_paused_dags_by_default):
    return any([
        show_paused_dags_url_param != hide_paused_dags_by_default,
        show_paused_dags_url_param is None
    ])


def generate_pages(current_page, num_of_pages,
                   search=None, showPaused=None, window=7):
    """
    Generates the HTML for a paging component using a similar logic to the paging
    auto-generated by Flask managed views. The paging component defines a number of
    pages visible in the pager (window) and once the user goes to a page beyond the
    largest visible, it would scroll to the right the page numbers and keeps the
    current one in the middle of the pager component. When in the last pages,
    the pages won't scroll and just keep moving until the last page. Pager also contains
    <first, previous, ..., next, last> pages.
    This component takes into account custom parameters such as search and showPaused,
    which could be added to the pages link in order to maintain the state between
    client and server. It also allows to make a bookmark on a specific paging state.

    :param current_page: the current page number, 0-indexed
    :param num_of_pages: the total number of pages
    :param search: the search query string, if any
    :param showPaused: false if paused dags will be hidden, otherwise true to show them
    :param window: the number of pages to be shown in the paging component (7 default)
    :return: the HTML string of the paging component
    """

    void_link = 'javascript:void(0)'
    first_node = Markup("""<li class="paginate_button {disabled}" id="dags_first">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="0" tabindex="0">&laquo;</a>
</li>""")

    previous_node = Markup("""<li class="paginate_button previous {disabled}" id="dags_previous">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="0" tabindex="0">&lsaquo;</a>
</li>""")

    next_node = Markup("""<li class="paginate_button next {disabled}" id="dags_next">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="3" tabindex="0">&rsaquo;</a>
</li>""")

    last_node = Markup("""<li class="paginate_button {disabled}" id="dags_last">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="3" tabindex="0">&raquo;</a>
</li>""")

    page_node = Markup("""<li class="paginate_button {is_active}">
    <a href="{href_link}" aria-controls="dags" data-dt-idx="2" tabindex="0">{page_num}</a>
</li>""")

    output = [Markup('<ul class="pagination" style="margin-top:0px;">')]

    is_disabled = 'disabled' if current_page <= 0 else ''
    output.append(first_node.format(href_link="?{}"
                                    .format(get_params(page=0,
                                                       search=search,
                                                       showPaused=showPaused)),
                                    disabled=is_disabled))

    page_link = void_link
    if current_page > 0:
        page_link = '?{}'.format(get_params(page=(current_page - 1),
                                            search=search,
                                            showPaused=showPaused))

    output.append(previous_node.format(href_link=page_link,
                                       disabled=is_disabled))

    mid = int(window / 2)
    last_page = num_of_pages - 1

    if current_page <= mid or num_of_pages < window:
        pages = [i for i in range(0, min(num_of_pages, window))]
    elif mid < current_page < last_page - mid:
        pages = [i for i in range(current_page - mid, current_page + mid + 1)]
    else:
        pages = [i for i in range(num_of_pages - window, last_page + 1)]

    def is_current(current, page):
        return page == current

    for page in pages:
        vals = {
            'is_active': 'active' if is_current(current_page, page) else '',
            'href_link': void_link if is_current(current_page, page)
                         else '?{}'.format(get_params(page=page,
                                                      search=search,
                                                      showPaused=showPaused)),
            'page_num': page + 1
        }
        output.append(page_node.format(**vals))

    is_disabled = 'disabled' if current_page >= num_of_pages - 1 else ''

    page_link = (void_link if current_page >= num_of_pages - 1
                 else '?{}'.format(get_params(page=current_page + 1,
                                              search=search,
                                              showPaused=showPaused)))

    output.append(next_node.format(href_link=page_link, disabled=is_disabled))
    output.append(last_node.format(href_link="?{}"
                                   .format(get_params(page=last_page,
                                                      search=search,
                                                      showPaused=showPaused)),
                                   disabled=is_disabled))

    output.append(Markup('</ul>'))

    return Markup('\n'.join(output))


def epoch(dttm):
    """Returns an epoch-type date"""
    return int(time.mktime(dttm.timetuple())) * 1000,


def json_response(obj):
    """
    returns a json response from a json serializable python object
    """
    return Response(
        response=json.dumps(
            obj, indent=4, cls=AirflowJsonEncoder),
        status=200,
        mimetype="application/json")


def make_cache_key(*args, **kwargs):
    """
    Used by cache to get a unique key per URL
    """
    path = request.path
    args = str(hash(frozenset(request.args.items())))
    return (path + args).encode('ascii', 'ignore')


def task_instance_link(attr):
    dag_id = attr.get('dag_id')
    task_id = attr.get('task_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.task',
        dag_id=dag_id,
        task_id=task_id,
        execution_date=execution_date.isoformat())
    url_root = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        root=task_id,
        execution_date=execution_date.isoformat())
    return Markup(
        """
        <span style="white-space: nowrap;">
        <a href="{url}">{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(url=url, task_id=task_id, url_root=url_root)


def state_token(state):
    color = State.color(state)
    return Markup(
        '<span class="label" style="background-color:{color};">'
        '{state}</span>').format(color=color, state=state)


def state_f(attr):
    state = attr.get('state')
    return state_token(state)


def nobr_f(attr_name):
    def nobr(attr):
        f = attr.get(attr_name)
        return Markup("<nobr>{}</nobr>").format(f)
    return nobr


def datetime_f(attr_name):
    def dt(attr):
        f = attr.get(attr_name)
        as_iso = f.isoformat() if f else ''
        if not as_iso:
            return Markup('')
        f = as_iso
        if timezone.utcnow().isoformat()[:4] == f[:4]:
            f = f[5:]
        # The empty title will be replaced in JS code when non-UTC dates are displayed
        return Markup('<nobr><time title="" datetime="{}">{}</time></nobr>').format(as_iso, f)
    return dt


def dag_link(attr):
    dag_id = attr.get('dag_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        execution_date=execution_date)
    return Markup(
        '<a href="{}">{}</a>').format(url, dag_id)


def dag_run_link(attr):
    dag_id = attr.get('dag_id')
    run_id = attr.get('run_id')
    execution_date = attr.get('execution_date')
    url = url_for(
        'Airflow.graph',
        dag_id=dag_id,
        run_id=run_id,
        execution_date=execution_date)
    return Markup(
        '<a href="{url}">{run_id}</a>').format(url=url, run_id=run_id)


def pygment_html_render(s, lexer=lexers.TextLexer):
    return highlight(
        s,
        lexer(),
        HtmlFormatter(linenos=True),
    )


def render(obj, lexer):
    out = ""
    if isinstance(obj, str):
        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>"
        if s is not None
        else None
    )


def get_attr_renderer():
    return {
        '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(get_python_source(x), lexers.PythonLexer),
    }


def recurse_tasks(tasks, task_ids, dag_ids, task_id_to_dag):
    if isinstance(tasks, list):
        for task in tasks:
            recurse_tasks(task, task_ids, dag_ids, task_id_to_dag)
        return
    if isinstance(tasks, SubDagOperator):
        subtasks = tasks.subdag.tasks
        dag_ids.append(tasks.subdag.dag_id)
        for subtask in subtasks:
            if subtask.task_id not in task_ids:
                task_ids.append(subtask.task_id)
                task_id_to_dag[subtask.task_id] = tasks.subdag
        recurse_tasks(subtasks, task_ids, dag_ids, task_id_to_dag)
    if isinstance(tasks, BaseOperator):
        task_id_to_dag[tasks.task_id] = tasks.dag


def get_chart_height(dag):
    """
    TODO(aoen): See [AIRFLOW-1263] We use the number of tasks in the DAG as a heuristic to
    approximate the size of generated chart (otherwise the charts are tiny and unreadable
    when DAGs have a large number of tasks). Ideally nvd3 should allow for dynamic-height
    charts, that is charts that take up space based on the size of the components within.
    """
    return 600 + len(dag.tasks) * 10


class UtcAwareFilterMixin:
    def apply(self, query, value):
        value = timezone.parse(value, timezone=timezone.utc)

        return super().apply(query, value)


class UtcAwareFilterEqual(UtcAwareFilterMixin, fab_sqlafilters.FilterEqual):
    pass


class UtcAwareFilterGreater(UtcAwareFilterMixin, fab_sqlafilters.FilterGreater):
    pass


class UtcAwareFilterSmaller(UtcAwareFilterMixin, fab_sqlafilters.FilterSmaller):
    pass


class UtcAwareFilterNotEqual(UtcAwareFilterMixin, fab_sqlafilters.FilterNotEqual):
    pass


class UtcAwareFilterConverter(fab_sqlafilters.SQLAFilterConverter):

    conversion_table = (
        (('is_utcdatetime', [UtcAwareFilterEqual,
                             UtcAwareFilterGreater,
                             UtcAwareFilterSmaller,
                             UtcAwareFilterNotEqual]),) +
        fab_sqlafilters.SQLAFilterConverter.conversion_table
    )


class CustomSQLAInterface(SQLAInterface):
    """
    FAB does not know how to handle columns with leading underscores because
    they are not supported by WTForm. This hack will remove the leading
    '_' from the key to lookup the column names.

    """
    def __init__(self, obj, session=None):
        super().__init__(obj, session=session)

        def clean_column_names():
            if self.list_properties:
                self.list_properties = {
                    k.lstrip('_'): v for k, v in self.list_properties.items()}
            if self.list_columns:
                self.list_columns = {
                    k.lstrip('_'): v for k, v in self.list_columns.items()}

        clean_column_names()

    def is_utcdatetime(self, col_name):
        from airflow.utils.sqlalchemy import UtcDateTime

        if col_name in self.list_columns:
            obj = self.list_columns[col_name].type
            return isinstance(obj, UtcDateTime) or \
                isinstance(obj, sqla.types.TypeDecorator) and \
                isinstance(obj.impl, UtcDateTime)
        return False

    filter_converter_class = UtcAwareFilterConverter


# This class is used directly (i.e. we cant tell Fab to use a different
# subclass) so we have no other option than to edit the converstion table in
# place
FieldConverter.conversion_table = (
    (('is_utcdatetime', DateTimeField, AirflowDateTimePickerWidget),) +
    FieldConverter.conversion_table
)


def get_dag(orm_dag: DagModel, store_serialized_dags=False):
    """Creates a dagbag to load and return a DAG.

    Calling it from UI should set store_serialized_dags = STORE_SERIALIZED_DAGS.
    There may be a delay for scheduler to write serialized DAG into database,
    loads from file in this case.
    FIXME: remove it when webserver does not access to DAG folder in future.
    """
    dag = DagBag(
        dag_folder=orm_dag.fileloc, store_serialized_dags=store_serialized_dags).get_dag(orm_dag.dag_id)
    if store_serialized_dags and dag is None:
        dag = DagBag(
            dag_folder=orm_dag.fileloc, store_serialized_dags=False).get_dag(orm_dag.dag_id)
    return dag

airflow/www/utils.py
diff --git a/airflow/www/utils.py b/airflow/www/utils.py
index a14233f242232..98ed570b6236d 100644
--- a/airflow/www/utils.py
+++ b/airflow/www/utils.py
@@ -406,8 +406,8 @@ class CustomSQLAInterface(SQLAInterface):
     '_' from the key to lookup the column names.
 
     """
-    def __init__(self, obj):
-        super().__init__(obj)
+    def __init__(self, obj, session=None):
+        super().__init__(obj, session=session)
 
         def clean_column_names():
             if self.list_properties:

Security Views broken
Thanks, @KostyaEsmukov for finding this bug. I was able to replicate this bug.


**Apache Airflow version**: 1.10.10.rc3


**Environment**:

- **OS** (e.g. from /etc/os-release): MacOs
- **Kernel** (e.g. `uname -a`): `Darwin MacBook-Pro.local 19.3.0 Darwin Kernel Version 19.3.0: Thu Jan  9 20:58:23 PST 2020; root:xnu-6153.81.5~1/RELEASE_X86_64 x86_64`

**What happened**:

Most of the tabs under **Security** dropdown are broken:

![image](https://user-images.githubusercontent.com/8811558/78505788-a7e27c00-776d-11ea-9204-ac0b60810e3e.png)

For example when clicking on **List Users** above: it shows the following error:

```

                          ____/ (  (    )   )  \___
                         /( (  (  )   _    ))  )   )\
                       ((     (   )(    )  )   (   )  )
                     ((/  ( _(   )   (   _) ) (  () )  )
                    ( (  ( (_)   ((    (   )  .((_ ) .  )_
                   ( (  )    (      (  )    )   ) . ) (   )
                  (  (   (  (   ) (  _  ( _) ).  ) . ) ) ( )
                  ( (  (   ) (  )   (  ))     ) _)(   )  )  )
                 ( (  ( \ ) (    (_  ( ) ( )  )   ) )  )) ( )
                  (  (   (  (   (_ ( ) ( _    )  ) (  )  )   )
                 ( (  ( (  (  )     (_  )  ) )  _)   ) _( ( )
                  ((  (   )(    (     _    )   _) _(_ (  (_ )
                   (_((__(_(__(( ( ( |  ) ) ) )_))__))_)___)
                   ((__)        \\||lll|l||///          \_))
                            (   /(/ (  )  ) )\   )
                          (    ( ( ( | | ) ) )\   )
                           (   /(| / ( )) ) ) )) )
                         (     ( ((((_(|)_)))))     )
                          (      ||\(|(|)|/||     )
                        (        |(||(||)||||        )
                          (     //|/l|||)|\\ \     )
                        (/ / //  /|//||||\\  \ \  \ _)
-------------------------------------------------------------------------------
Node: 1.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.ip6.arpa
-------------------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/app.py", line 2447, in wsgi_app
    response = self.full_dispatch_request()
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/app.py", line 1952, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/app.py", line 1821, in handle_user_exception
    reraise(exc_type, exc_value, tb)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/_compat.py", line 39, in reraise
    raise value
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/app.py", line 1950, in full_dispatch_request
    rv = self.dispatch_request()
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/app.py", line 1936, in dispatch_request
    return self.view_functions[rule.endpoint](**req.view_args)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/security/decorators.py", line 109, in wraps
    return f(self, *args, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/views.py", line 553, in list
    self.list_template, title=self.list_title, widgets=widgets
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/baseviews.py", line 281, in render_template
    template, **dict(list(kwargs.items()) + list(self.extra_args.items()))
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/templating.py", line 140, in render_template
    ctx.app,
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask/templating.py", line 120, in _render
    rv = template.render(context)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/asyncsupport.py", line 76, in render
    return original_render(self, *args, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/environment.py", line 1008, in render
    return self.environment.handle_exception(exc_info, True)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/environment.py", line 780, in handle_exception
    reraise(exc_type, exc_value, tb)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/_compat.py", line 37, in reraise
    raise value.with_traceback(tb)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/model/list.html", line 2, in top-level template code
    {% import 'appbuilder/general/lib.html' as lib %}
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/environment.py", line 1005, in render
    return concat(self.root_render_func(self.new_context(vars)))
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/model/list.html", line 17, in root
    {% endblock %}
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/base.html", line 15, in root
  File "/Users/kaxilnaik/Documents/Github/apache/airflow/airflow/www_rbac/templates/airflow/master.html", line 16, in root
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/baselayout.html", line 17, in root
    {% include 'appbuilder/flash.html' %}
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/init.html", line 32, in root
    <link href="{{url_for('appbuilder.static',filename='select2/select2.css')}}" rel="stylesheet">
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/baselayout.html", line 37, in block_body
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/model/list.html", line 29, in block_content
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/model/list.html", line 56, in block_list_search
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 262, in call
    return __obj(*args, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 570, in __call__
    return self._invoke(arguments, autoescape)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/asyncsupport.py", line 110, in _invoke
    return original_invoke(self, arguments, autoescape)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 574, in _invoke
    rv = self._func(*arguments)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/lib.html", line 783, in macro
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 262, in call
    return __obj(*args, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 570, in __call__
    return self._invoke(arguments, autoescape)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/asyncsupport.py", line 110, in _invoke
    return original_invoke(self, arguments, autoescape)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 574, in _invoke
    rv = self._func(*arguments)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/templates/appbuilder/general/model/list.html", line 51, in macro
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/jinja2/runtime.py", line 262, in call
    return __obj(*args, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/widgets.py", line 115, in __call__
    form_fields[col] = self.template_args["form"][col]()
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/wtforms/fields/core.py", line 155, in __call__
    return self.meta.render_field(self, kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/wtforms/meta.py", line 56, in render_field
    return field.widget(field, **render_kw)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/fieldwidgets.py", line 176, in __call__
    return super(Select2ManyWidget, self).__call__(field, **kwargs)
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/wtforms/widgets/core.py", line 323, in __call__
    for val, label, selected in field.iter_choices():
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/fields.py", line 208, in iter_choices
    for pk, obj in self._get_object_list():
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/fields.py", line 128, in _get_object_list
    objs = self.query_func()
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/forms.py", line 139, in <lambda>
    return lambda: self.datamodel.get_related_interface(col_name).query()[1]
  File "/Users/kaxilnaik/.virtualenvs/airflow_test_upgrade/lib/python3.7/site-packages/flask_appbuilder/models/sqla/interface.py", line 536, in get_related_interface
    return self.__class__(self.get_related_model(col_name), self.session)
TypeError: __init__() takes 2 positional arguments but 3 were given
```




