<commit_before># pylint: disable=C,R,W
"""This module contains the 'Viz' objects

These objects represent the backend of all the visualizations that
Superset can render.
"""
from collections import defaultdict, OrderedDict
import copy
from datetime import datetime, timedelta
from functools import reduce
import hashlib
import inspect
from itertools import product
import logging
import math
import pickle as pkl
import re
import traceback
import uuid

from dateutil import relativedelta as rdelta
from flask import request
from flask_babel import lazy_gettext as _
import geohash
from geopy.point import Point
from markdown import markdown
import numpy as np
import pandas as pd
from pandas.tseries.frequencies import to_offset
from past.builtins import basestring
import polyline
import simplejson as json

from superset import app, cache, get_css_manifest_files
from superset.exceptions import NullValueException, SpatialException
from superset.utils import core as utils
from superset.utils.core import (
    DTTM_ALIAS,
    JS_MAX_INTEGER,
    merge_extra_filters,
    to_adhoc,
)


config = app.config
stats_logger = config.get('STATS_LOGGER')

METRIC_KEYS = [
    'metric', 'metrics', 'percent_metrics', 'metric_2', 'secondary_metric',
    'x', 'y', 'size',
]


class BaseViz(object):

    """All visualizations derive this base class"""

    viz_type = None
    verbose_name = 'Base Viz'
    credits = ''
    is_timeseries = False
    default_fillna = 0
    cache_type = 'df'
    enforce_numerical_metrics = True

    def __init__(self, datasource, form_data, force=False):
        if not datasource:
            raise Exception(_('Viz is missing a datasource'))

        self.datasource = datasource
        self.request = request
        self.viz_type = form_data.get('viz_type')
        self.form_data = form_data

        self.query = ''
        self.token = self.form_data.get(
            'token', 'token_' + uuid.uuid4().hex[:8])

        self.groupby = self.form_data.get('groupby') or []
        self.time_shift = timedelta()

        self.status = None
        self.error_message = None
        self.force = force

        # Keeping track of whether some data came from cache
        # this is useful to trigger the <CachedLabel /> when
        # in the cases where visualization have many queries
        # (FilterBox for instance)
        self._some_from_cache = False
        self._any_cache_key = None
        self._any_cached_dttm = None
        self._extra_chart_data = []

        self.process_metrics()

    def process_metrics(self):
        # metrics in TableViz is order sensitive, so metric_dict should be
        # OrderedDict
        self.metric_dict = OrderedDict()
        fd = self.form_data
        for mkey in METRIC_KEYS:
            val = fd.get(mkey)
            if val:
                if not isinstance(val, list):
                    val = [val]
                for o in val:
                    label = self.get_metric_label(o)
                    if isinstance(o, dict):
                        o['label'] = label
                    self.metric_dict[label] = o

        # Cast to list needed to return serializable object in py3
        self.all_metrics = list(self.metric_dict.values())
        self.metric_labels = list(self.metric_dict.keys())

    def get_metric_label(self, metric):
        if isinstance(metric, str):
            return metric

        if isinstance(metric, dict):
            metric = metric.get('label')

        if self.datasource.type == 'table':
            db_engine_spec = self.datasource.database.db_engine_spec
            metric = db_engine_spec.mutate_expression_label(metric)
        return metric

    @staticmethod
    def handle_js_int_overflow(data):
        for d in data.get('records', dict()):
            for k, v in list(d.items()):
                if isinstance(v, int):
                    # if an int is too big for Java Script to handle
                    # convert it to a string
                    if abs(v) > JS_MAX_INTEGER:
                        d[k] = str(v)
        return data

    def run_extra_queries(self):
        """Lifecycle method to use when more than one query is needed

        In rare-ish cases, a visualization may need to execute multiple
        queries. That is the case for FilterBox or for time comparison
        in Line chart for instance.

        In those cases, we need to make sure these queries run before the
        main `get_payload` method gets called, so that the overall caching
        metadata can be right. The way it works here is that if any of
        the previous `get_df_payload` calls hit the cache, the main
        payload's metadata will reflect that.

        The multi-query support may need more work to become a first class
        use case in the framework, and for the UI to reflect the subtleties
        (show that only some of the queries were served from cache for
        instance). In the meantime, since multi-query is rare, we treat
        it with a bit of a hack. Note that the hack became necessary
        when moving from caching the visualization's data itself, to caching
        the underlying query(ies).
        """
        pass

    def handle_nulls(self, df):
        fillna = self.get_fillna_for_columns(df.columns)
        return df.fillna(fillna)

    def get_fillna_for_col(self, col):
        """Returns the value to use as filler for a specific Column.type"""
        if col:
            if col.is_string:
                return ' NULL'
        return self.default_fillna

    def get_fillna_for_columns(self, columns=None):
        """Returns a dict or scalar that can be passed to DataFrame.fillna"""
        if columns is None:
            return self.default_fillna
        columns_dict = {col.column_name: col for col in self.datasource.columns}
        fillna = {
            c: self.get_fillna_for_col(columns_dict.get(c))
            for c in columns
        }
        return fillna

    def get_samples(self):
        query_obj = self.query_obj()
        query_obj.update({
            'groupby': [],
            'metrics': [],
            'row_limit': 1000,
            'columns': [o.column_name for o in self.datasource.columns],
        })
        df = self.get_df(query_obj)
        return df.to_dict(orient='records')

    def get_df(self, query_obj=None):
        """Returns a pandas dataframe based on the query object"""
        if not query_obj:
            query_obj = self.query_obj()
        if not query_obj:
            return None

        self.error_msg = ''

        timestamp_format = None
        if self.datasource.type == 'table':
            dttm_col = self.datasource.get_col(query_obj['granularity'])
            if dttm_col:
                timestamp_format = dttm_col.python_date_format

        # The datasource here can be different backend but the interface is common
        self.results = self.datasource.query(query_obj)
        self.query = self.results.query
        self.status = self.results.status
        self.error_message = self.results.error_message

        df = self.results.df
        # Transform the timestamp we received from database to pandas supported
        # datetime format. If no python_date_format is specified, the pattern will
        # be considered as the default ISO date format
        # If the datetime format is unix, the parse will use the corresponding
        # parsing logic.
        if df is not None and not df.empty:
            if DTTM_ALIAS in df.columns:
                if timestamp_format in ('epoch_s', 'epoch_ms'):
                    # Column has already been formatted as a timestamp.
                    df[DTTM_ALIAS] = df[DTTM_ALIAS].apply(pd.Timestamp)
                else:
                    df[DTTM_ALIAS] = pd.to_datetime(
                        df[DTTM_ALIAS], utc=False, format=timestamp_format)
                if self.datasource.offset:
                    df[DTTM_ALIAS] += timedelta(hours=self.datasource.offset)
                df[DTTM_ALIAS] += self.time_shift

            if self.enforce_numerical_metrics:
                self.df_metrics_to_num(df)

            df.replace([np.inf, -np.inf], np.nan)
            df = self.handle_nulls(df)
        return df

    def df_metrics_to_num(self, df):
        """Converting metrics to numeric when pandas.read_sql cannot"""
        metrics = self.metric_labels
        for col, dtype in df.dtypes.items():
            if dtype.type == np.object_ and col in metrics:
                df[col] = pd.to_numeric(df[col], errors='coerce')

    def process_query_filters(self):
        utils.convert_legacy_filters_into_adhoc(self.form_data)
        merge_extra_filters(self.form_data)
        utils.split_adhoc_filters_into_base_filters(self.form_data)

    def query_obj(self):
        """Building a query object"""
        form_data = self.form_data
        self.process_query_filters()
        gb = form_data.get('groupby') or []
        metrics = self.all_metrics or []
        columns = form_data.get('columns') or []
        groupby = []
        for o in gb + columns:
            if o not in groupby:
                groupby.append(o)

        is_timeseries = self.is_timeseries
        if DTTM_ALIAS in groupby:
            groupby.remove(DTTM_ALIAS)
            is_timeseries = True

        granularity = (
            form_data.get('granularity') or
            form_data.get('granularity_sqla')
        )
        limit = int(form_data.get('limit') or 0)
        timeseries_limit_metric = form_data.get('timeseries_limit_metric')
        row_limit = int(form_data.get('row_limit') or config.get('ROW_LIMIT'))

        # default order direction
        order_desc = form_data.get('order_desc', True)

        since, until = utils.get_since_until(form_data.get('time_range'),
                                             form_data.get('since'),
                                             form_data.get('until'))
        time_shift = form_data.get('time_shift', '')
        self.time_shift = utils.parse_human_timedelta(time_shift)
        from_dttm = None if since is None else (since - self.time_shift)
        to_dttm = None if until is None else (until - self.time_shift)
        if from_dttm and to_dttm and from_dttm > to_dttm:
            raise Exception(_('From date cannot be larger than to date'))

        self.from_dttm = from_dttm
        self.to_dttm = to_dttm

        # extras are used to query elements specific to a datasource type
        # for instance the extra where clause that applies only to Tables
        extras = {
            'where': form_data.get('where', ''),
            'having': form_data.get('having', ''),
            'having_druid': form_data.get('having_filters', []),
            'time_grain_sqla': form_data.get('time_grain_sqla', ''),
            'druid_time_origin': form_data.get('druid_time_origin', ''),
        }

        d = {
            'granularity': granularity,
            'from_dttm': from_dttm,
            'to_dttm': to_dttm,
            'is_timeseries': is_timeseries,
            'groupby': groupby,
            'metrics': metrics,
            'row_limit': row_limit,
            'filter': self.form_data.get('filters', []),
            'timeseries_limit': limit,
            'extras': extras,
            'timeseries_limit_metric': timeseries_limit_metric,
            'order_desc': order_desc,
            'prequeries': [],
            'is_prequery': False,
        }
        return d

    @property
    def cache_timeout(self):
        if self.form_data.get('cache_timeout') is not None:
            return int(self.form_data.get('cache_timeout'))
        if self.datasource.cache_timeout is not None:
            return self.datasource.cache_timeout
        if (
                hasattr(self.datasource, 'database') and
                self.datasource.database.cache_timeout) is not None:
            return self.datasource.database.cache_timeout
        return config.get('CACHE_DEFAULT_TIMEOUT')

    def get_json(self):
        return json.dumps(
            self.get_payload(),
            default=utils.json_int_dttm_ser, ignore_nan=True)

    def cache_key(self, query_obj, **extra):
        """
        The cache key is made out of the key/values in `query_obj`, plus any
        other key/values in `extra`.

        We remove datetime bounds that are hard values, and replace them with
        the use-provided inputs to bounds, which may be time-relative (as in
        "5 days ago" or "now").

        The `extra` arguments are currently used by time shift queries, since
        different time shifts wil differ only in the `from_dttm` and `to_dttm`
        values which are stripped.
        """
        cache_dict = copy.copy(query_obj)
        cache_dict.update(extra)

        for k in ['from_dttm', 'to_dttm']:
            del cache_dict[k]

        cache_dict['time_range'] = self.form_data.get('time_range')
        cache_dict['datasource'] = self.datasource.uid
        json_data = self.json_dumps(cache_dict, sort_keys=True)
        return hashlib.md5(json_data.encode('utf-8')).hexdigest()

    def get_payload(self, query_obj=None):
        """Returns a payload of metadata and data"""
        self.run_extra_queries()
        payload = self.get_df_payload(query_obj)

        df = payload.get('df')
        if self.status != utils.QueryStatus.FAILED:
            if df is not None and df.empty:
                payload['error'] = 'No data'
            else:
                payload['data'] = self.get_data(df)
        if 'df' in payload:
            del payload['df']
        return payload

    def get_df_payload(self, query_obj=None, **kwargs):
        """Handles caching around the df payload retrieval"""
        if not query_obj:
            query_obj = self.query_obj()
        cache_key = self.cache_key(query_obj, **kwargs) if query_obj else None
        logging.info('Cache key: {}'.format(cache_key))
        is_loaded = False
        stacktrace = None
        df = None
        cached_dttm = datetime.utcnow().isoformat().split('.')[0]
        if cache_key and cache and not self.force:
            cache_value = cache.get(cache_key)
            if cache_value:
                stats_logger.incr('loaded_from_cache')
                try:
                    cache_value = pkl.loads(cache_value)
                    df = cache_value['df']
                    self.query = cache_value['query']
                    self._any_cached_dttm = cache_value['dttm']
                    self._any_cache_key = cache_key
                    self.status = utils.QueryStatus.SUCCESS
                    is_loaded = True
                except Exception as e:
                    logging.exception(e)
                    logging.error('Error reading cache: ' +
                                  utils.error_msg_from_exception(e))
                logging.info('Serving from cache')

        if query_obj and not is_loaded:
            try:
                df = self.get_df(query_obj)
                if self.status != utils.QueryStatus.FAILED:
                    stats_logger.incr('loaded_from_source')
                    is_loaded = True
            except Exception as e:
                logging.exception(e)
                if not self.error_message:
                    self.error_message = '{}'.format(e)
                self.status = utils.QueryStatus.FAILED
                stacktrace = traceback.format_exc()

            if (
                    is_loaded and
                    cache_key and
                    cache and
                    self.status != utils.QueryStatus.FAILED):
                try:
                    cache_value = dict(
                        dttm=cached_dttm,
                        df=df if df is not None else None,
                        query=self.query,
                    )
                    cache_value = pkl.dumps(
                        cache_value, protocol=pkl.HIGHEST_PROTOCOL)

                    logging.info('Caching {} chars at key {}'.format(
                        len(cache_value), cache_key))

                    stats_logger.incr('set_cache_key')
                    cache.set(
                        cache_key,
                        cache_value,
                        timeout=self.cache_timeout)
                except Exception as e:
                    # cache.set call can fail if the backend is down or if
                    # the key is too large or whatever other reasons
                    logging.warning('Could not cache key {}'.format(cache_key))
                    logging.exception(e)
                    cache.delete(cache_key)
        return {
            'cache_key': self._any_cache_key,
            'cached_dttm': self._any_cached_dttm,
            'cache_timeout': self.cache_timeout,
            'df': df,
            'error': self.error_message,
            'form_data': self.form_data,
            'is_cached': self._any_cache_key is not None,
            'query': self.query,
            'status': self.status,
            'stacktrace': stacktrace,
            'rowcount': len(df.index) if df is not None else 0,
        }

    def json_dumps(self, obj, sort_keys=False):
        return json.dumps(
            obj,
            default=utils.json_int_dttm_ser,
            ignore_nan=True,
            sort_keys=sort_keys,
        )

    def payload_json_and_has_error(self, payload):
        has_error = payload.get('status') == utils.QueryStatus.FAILED or \
            payload.get('error') is not None
        return self.json_dumps(payload), has_error

    @property
    def data(self):
        """This is the data object serialized to the js layer"""
        content = {
            'form_data': self.form_data,
            'token': self.token,
            'viz_name': self.viz_type,
            'filter_select_enabled': self.datasource.filter_select_enabled,
        }
        return content

    def get_csv(self):
        df = self.get_df()
        include_index = not isinstance(df.index, pd.RangeIndex)
        return df.to_csv(index=include_index, **config.get('CSV_EXPORT'))

    def get_data(self, df):
        return df.to_dict(orient='records')

    @property
    def json_data(self):
        return json.dumps(self.data)


class TableViz(BaseViz):

    """A basic html table that is sortable and searchable"""

    viz_type = 'table'
    verbose_name = _('Table View')
    credits = 'a <a href="https://github.com/airbnb/superset">Superset</a> original'
    is_timeseries = False
    enforce_numerical_metrics = False

    def should_be_timeseries(self):
        fd = self.form_data
        # TODO handle datasource-type-specific code in datasource
        conditions_met = (
            (fd.get('granularity') and fd.get('granularity') != 'all') or
            (fd.get('granularity_sqla') and fd.get('time_grain_sqla'))
        )
        if fd.get('include_time') and not conditions_met:
            raise Exception(_(
                'Pick a granularity in the Time section or '
                "uncheck 'Include Time'"))
        return fd.get('include_time')

    def query_obj(self):
        d = super(TableViz, self).query_obj()
        fd = self.form_data

        if fd.get('all_columns') and (fd.get('groupby') or fd.get('metrics')):
            raise Exception(_(
                'Choose either fields to [Group By] and [Metrics] or '
                '[Columns], not both'))

        sort_by = fd.get('timeseries_limit_metric')
        if fd.get('all_columns'):
            d['columns'] = fd.get('all_columns')
            d['groupby'] = []
            order_by_cols = fd.get('order_by_cols') or []
            d['orderby'] = [json.loads(t) for t in order_by_cols]
        elif sort_by:
            sort_by_label = utils.get_metric_name(sort_by)
            if sort_by_label not in utils.get_metric_names(d['metrics']):
                d['metrics'] += [sort_by]
            d['orderby'] = [(sort_by, not fd.get('order_desc', True))]

        # Add all percent metrics that are not already in the list
        if 'percent_metrics' in fd:
            d['metrics'] = d['metrics'] + list(filter(
                lambda m: m not in d['metrics'],
                fd['percent_metrics'] or [],
            ))

        d['is_timeseries'] = self.should_be_timeseries()
        return d

    def get_data(self, df):
        fd = self.form_data
        if (
                not self.should_be_timeseries() and
                df is not None and
                DTTM_ALIAS in df
        ):
            del df[DTTM_ALIAS]

        # Sum up and compute percentages for all percent metrics
        percent_metrics = fd.get('percent_metrics') or []
        percent_metrics = [self.get_metric_label(m) for m in percent_metrics]

        if len(percent_metrics):
            percent_metrics = list(filter(lambda m: m in df, percent_metrics))
            metric_sums = {
                m: reduce(lambda a, b: a + b, df[m])
                for m in percent_metrics
            }
            metric_percents = {
                m: list(map(
                    lambda a: None if metric_sums[m] == 0 else a / metric_sums[m], df[m]))
                for m in percent_metrics
            }
            for m in percent_metrics:
                m_name = '%' + m
                df[m_name] = pd.Series(metric_percents[m], name=m_name)
            # Remove metrics that are not in the main metrics list
            metrics = fd.get('metrics') or []
            metrics = [self.get_metric_label(m) for m in metrics]
            for m in filter(
                lambda m: m not in metrics and m in df.columns,
                percent_metrics,
            ):
                del df[m]

        data = self.handle_js_int_overflow(
            dict(
                records=df.to_dict(orient='records'),
                columns=list(df.columns),
            ))

        return data

    def json_dumps(self, obj, sort_keys=False):
        if self.form_data.get('all_columns'):
            return json.dumps(
                obj,
                default=utils.json_iso_dttm_ser,
                sort_keys=sort_keys,
                ignore_nan=True)
        else:
            return super(TableViz, self).json_dumps(obj)


class TimeTableViz(BaseViz):

    """A data table with rich time-series related columns"""

    viz_type = 'time_table'
    verbose_name = _('Time Table View')
    credits = 'a <a href="https://github.com/airbnb/superset">Superset</a> original'
    is_timeseries = True

    def query_obj(self):<commit_msg>json_iso_dttm_ser use for TableViz (#6563)<commit_after># pylint: disable=C,R,W
"""This module contains the 'Viz' objects

These objects represent the backend of all the visualizations that
Superset can render.
"""
from collections import defaultdict, OrderedDict
import copy
from datetime import datetime, timedelta
from functools import reduce
import hashlib
import inspect
from itertools import product
import logging
import math
import pickle as pkl
import re
import traceback
import uuid

from dateutil import relativedelta as rdelta
from flask import request
from flask_babel import lazy_gettext as _
import geohash
from geopy.point import Point
from markdown import markdown
import numpy as np
import pandas as pd
from pandas.tseries.frequencies import to_offset
from past.builtins import basestring
import polyline
import simplejson as json

from superset import app, cache, get_css_manifest_files
from superset.exceptions import NullValueException, SpatialException
from superset.utils import core as utils
from superset.utils.core import (
    DTTM_ALIAS,
    JS_MAX_INTEGER,
    merge_extra_filters,
    to_adhoc,
)


config = app.config
stats_logger = config.get('STATS_LOGGER')

METRIC_KEYS = [
    'metric', 'metrics', 'percent_metrics', 'metric_2', 'secondary_metric',
    'x', 'y', 'size',
]


class BaseViz(object):

    """All visualizations derive this base class"""

    viz_type = None
    verbose_name = 'Base Viz'
    credits = ''
    is_timeseries = False
    default_fillna = 0
    cache_type = 'df'
    enforce_numerical_metrics = True

    def __init__(self, datasource, form_data, force=False):
        if not datasource:
            raise Exception(_('Viz is missing a datasource'))

        self.datasource = datasource
        self.request = request
        self.viz_type = form_data.get('viz_type')
        self.form_data = form_data

        self.query = ''
        self.token = self.form_data.get(
            'token', 'token_' + uuid.uuid4().hex[:8])

        self.groupby = self.form_data.get('groupby') or []
        self.time_shift = timedelta()

        self.status = None
        self.error_message = None
        self.force = force

        # Keeping track of whether some data came from cache
        # this is useful to trigger the <CachedLabel /> when
        # in the cases where visualization have many queries
        # (FilterBox for instance)
        self._some_from_cache = False
        self._any_cache_key = None
        self._any_cached_dttm = None
        self._extra_chart_data = []

        self.process_metrics()

    def process_metrics(self):
        # metrics in TableViz is order sensitive, so metric_dict should be
        # OrderedDict
        self.metric_dict = OrderedDict()
        fd = self.form_data
        for mkey in METRIC_KEYS:
            val = fd.get(mkey)
            if val:
                if not isinstance(val, list):
                    val = [val]
                for o in val:
                    label = self.get_metric_label(o)
                    if isinstance(o, dict):
                        o['label'] = label
                    self.metric_dict[label] = o

        # Cast to list needed to return serializable object in py3
        self.all_metrics = list(self.metric_dict.values())
        self.metric_labels = list(self.metric_dict.keys())

    def get_metric_label(self, metric):
        if isinstance(metric, str):
            return metric

        if isinstance(metric, dict):
            metric = metric.get('label')

        if self.datasource.type == 'table':
            db_engine_spec = self.datasource.database.db_engine_spec
            metric = db_engine_spec.mutate_expression_label(metric)
        return metric

    @staticmethod
    def handle_js_int_overflow(data):
        for d in data.get('records', dict()):
            for k, v in list(d.items()):
                if isinstance(v, int):
                    # if an int is too big for Java Script to handle
                    # convert it to a string
                    if abs(v) > JS_MAX_INTEGER:
                        d[k] = str(v)
        return data

    def run_extra_queries(self):
        """Lifecycle method to use when more than one query is needed

        In rare-ish cases, a visualization may need to execute multiple
        queries. That is the case for FilterBox or for time comparison
        in Line chart for instance.

        In those cases, we need to make sure these queries run before the
        main `get_payload` method gets called, so that the overall caching
        metadata can be right. The way it works here is that if any of
        the previous `get_df_payload` calls hit the cache, the main
        payload's metadata will reflect that.

        The multi-query support may need more work to become a first class
        use case in the framework, and for the UI to reflect the subtleties
        (show that only some of the queries were served from cache for
        instance). In the meantime, since multi-query is rare, we treat
        it with a bit of a hack. Note that the hack became necessary
        when moving from caching the visualization's data itself, to caching
        the underlying query(ies).
        """
        pass

    def handle_nulls(self, df):
        fillna = self.get_fillna_for_columns(df.columns)
        return df.fillna(fillna)

    def get_fillna_for_col(self, col):
        """Returns the value to use as filler for a specific Column.type"""
        if col:
            if col.is_string:
                return ' NULL'
        return self.default_fillna

    def get_fillna_for_columns(self, columns=None):
        """Returns a dict or scalar that can be passed to DataFrame.fillna"""
        if columns is None:
            return self.default_fillna
        columns_dict = {col.column_name: col for col in self.datasource.columns}
        fillna = {
            c: self.get_fillna_for_col(columns_dict.get(c))
            for c in columns
        }
        return fillna

    def get_samples(self):
        query_obj = self.query_obj()
        query_obj.update({
            'groupby': [],
            'metrics': [],
            'row_limit': 1000,
            'columns': [o.column_name for o in self.datasource.columns],
        })
        df = self.get_df(query_obj)
        return df.to_dict(orient='records')

    def get_df(self, query_obj=None):
        """Returns a pandas dataframe based on the query object"""
        if not query_obj:
            query_obj = self.query_obj()
        if not query_obj:
            return None

        self.error_msg = ''

        timestamp_format = None
        if self.datasource.type == 'table':
            dttm_col = self.datasource.get_col(query_obj['granularity'])
            if dttm_col:
                timestamp_format = dttm_col.python_date_format

        # The datasource here can be different backend but the interface is common
        self.results = self.datasource.query(query_obj)
        self.query = self.results.query
        self.status = self.results.status
        self.error_message = self.results.error_message

        df = self.results.df
        # Transform the timestamp we received from database to pandas supported
        # datetime format. If no python_date_format is specified, the pattern will
        # be considered as the default ISO date format
        # If the datetime format is unix, the parse will use the corresponding
        # parsing logic.
        if df is not None and not df.empty:
            if DTTM_ALIAS in df.columns:
                if timestamp_format in ('epoch_s', 'epoch_ms'):
                    # Column has already been formatted as a timestamp.
                    df[DTTM_ALIAS] = df[DTTM_ALIAS].apply(pd.Timestamp)
                else:
                    df[DTTM_ALIAS] = pd.to_datetime(
                        df[DTTM_ALIAS], utc=False, format=timestamp_format)
                if self.datasource.offset:
                    df[DTTM_ALIAS] += timedelta(hours=self.datasource.offset)
                df[DTTM_ALIAS] += self.time_shift

            if self.enforce_numerical_metrics:
                self.df_metrics_to_num(df)

            df.replace([np.inf, -np.inf], np.nan)
            df = self.handle_nulls(df)
        return df

    def df_metrics_to_num(self, df):
        """Converting metrics to numeric when pandas.read_sql cannot"""
        metrics = self.metric_labels
        for col, dtype in df.dtypes.items():
            if dtype.type == np.object_ and col in metrics:
                df[col] = pd.to_numeric(df[col], errors='coerce')

    def process_query_filters(self):
        utils.convert_legacy_filters_into_adhoc(self.form_data)
        merge_extra_filters(self.form_data)
        utils.split_adhoc_filters_into_base_filters(self.form_data)

    def query_obj(self):
        """Building a query object"""
        form_data = self.form_data
        self.process_query_filters()
        gb = form_data.get('groupby') or []
        metrics = self.all_metrics or []
        columns = form_data.get('columns') or []
        groupby = []
        for o in gb + columns:
            if o not in groupby:
                groupby.append(o)

        is_timeseries = self.is_timeseries
        if DTTM_ALIAS in groupby:
            groupby.remove(DTTM_ALIAS)
            is_timeseries = True

        granularity = (
            form_data.get('granularity') or
            form_data.get('granularity_sqla')
        )
        limit = int(form_data.get('limit') or 0)
        timeseries_limit_metric = form_data.get('timeseries_limit_metric')
        row_limit = int(form_data.get('row_limit') or config.get('ROW_LIMIT'))

        # default order direction
        order_desc = form_data.get('order_desc', True)

        since, until = utils.get_since_until(form_data.get('time_range'),
                                             form_data.get('since'),
                                             form_data.get('until'))
        time_shift = form_data.get('time_shift', '')
        self.time_shift = utils.parse_human_timedelta(time_shift)
        from_dttm = None if since is None else (since - self.time_shift)
        to_dttm = None if until is None else (until - self.time_shift)
        if from_dttm and to_dttm and from_dttm > to_dttm:
            raise Exception(_('From date cannot be larger than to date'))

        self.from_dttm = from_dttm
        self.to_dttm = to_dttm

        # extras are used to query elements specific to a datasource type
        # for instance the extra where clause that applies only to Tables
        extras = {
            'where': form_data.get('where', ''),
            'having': form_data.get('having', ''),
            'having_druid': form_data.get('having_filters', []),
            'time_grain_sqla': form_data.get('time_grain_sqla', ''),
            'druid_time_origin': form_data.get('druid_time_origin', ''),
        }

        d = {
            'granularity': granularity,
            'from_dttm': from_dttm,
            'to_dttm': to_dttm,
            'is_timeseries': is_timeseries,
            'groupby': groupby,
            'metrics': metrics,
            'row_limit': row_limit,
            'filter': self.form_data.get('filters', []),
            'timeseries_limit': limit,
            'extras': extras,
            'timeseries_limit_metric': timeseries_limit_metric,
            'order_desc': order_desc,
            'prequeries': [],
            'is_prequery': False,
        }
        return d

    @property
    def cache_timeout(self):
        if self.form_data.get('cache_timeout') is not None:
            return int(self.form_data.get('cache_timeout'))
        if self.datasource.cache_timeout is not None:
            return self.datasource.cache_timeout
        if (
                hasattr(self.datasource, 'database') and
                self.datasource.database.cache_timeout) is not None:
            return self.datasource.database.cache_timeout
        return config.get('CACHE_DEFAULT_TIMEOUT')

    def get_json(self):
        return json.dumps(
            self.get_payload(),
            default=utils.json_int_dttm_ser, ignore_nan=True)

    def cache_key(self, query_obj, **extra):
        """
        The cache key is made out of the key/values in `query_obj`, plus any
        other key/values in `extra`.

        We remove datetime bounds that are hard values, and replace them with
        the use-provided inputs to bounds, which may be time-relative (as in
        "5 days ago" or "now").

        The `extra` arguments are currently used by time shift queries, since
        different time shifts wil differ only in the `from_dttm` and `to_dttm`
        values which are stripped.
        """
        cache_dict = copy.copy(query_obj)
        cache_dict.update(extra)

        for k in ['from_dttm', 'to_dttm']:
            del cache_dict[k]

        cache_dict['time_range'] = self.form_data.get('time_range')
        cache_dict['datasource'] = self.datasource.uid
        json_data = self.json_dumps(cache_dict, sort_keys=True)
        return hashlib.md5(json_data.encode('utf-8')).hexdigest()

    def get_payload(self, query_obj=None):
        """Returns a payload of metadata and data"""
        self.run_extra_queries()
        payload = self.get_df_payload(query_obj)

        df = payload.get('df')
        if self.status != utils.QueryStatus.FAILED:
            if df is not None and df.empty:
                payload['error'] = 'No data'
            else:
                payload['data'] = self.get_data(df)
        if 'df' in payload:
            del payload['df']
        return payload

    def get_df_payload(self, query_obj=None, **kwargs):
        """Handles caching around the df payload retrieval"""
        if not query_obj:
            query_obj = self.query_obj()
        cache_key = self.cache_key(query_obj, **kwargs) if query_obj else None
        logging.info('Cache key: {}'.format(cache_key))
        is_loaded = False
        stacktrace = None
        df = None
        cached_dttm = datetime.utcnow().isoformat().split('.')[0]
        if cache_key and cache and not self.force:
            cache_value = cache.get(cache_key)
            if cache_value:
                stats_logger.incr('loaded_from_cache')
                try:
                    cache_value = pkl.loads(cache_value)
                    df = cache_value['df']
                    self.query = cache_value['query']
                    self._any_cached_dttm = cache_value['dttm']
                    self._any_cache_key = cache_key
                    self.status = utils.QueryStatus.SUCCESS
                    is_loaded = True
                except Exception as e:
                    logging.exception(e)
                    logging.error('Error reading cache: ' +
                                  utils.error_msg_from_exception(e))
                logging.info('Serving from cache')

        if query_obj and not is_loaded:
            try:
                df = self.get_df(query_obj)
                if self.status != utils.QueryStatus.FAILED:
                    stats_logger.incr('loaded_from_source')
                    is_loaded = True
            except Exception as e:
                logging.exception(e)
                if not self.error_message:
                    self.error_message = '{}'.format(e)
                self.status = utils.QueryStatus.FAILED
                stacktrace = traceback.format_exc()

            if (
                    is_loaded and
                    cache_key and
                    cache and
                    self.status != utils.QueryStatus.FAILED):
                try:
                    cache_value = dict(
                        dttm=cached_dttm,
                        df=df if df is not None else None,
                        query=self.query,
                    )
                    cache_value = pkl.dumps(
                        cache_value, protocol=pkl.HIGHEST_PROTOCOL)

                    logging.info('Caching {} chars at key {}'.format(
                        len(cache_value), cache_key))

                    stats_logger.incr('set_cache_key')
                    cache.set(
                        cache_key,
                        cache_value,
                        timeout=self.cache_timeout)
                except Exception as e:
                    # cache.set call can fail if the backend is down or if
                    # the key is too large or whatever other reasons
                    logging.warning('Could not cache key {}'.format(cache_key))
                    logging.exception(e)
                    cache.delete(cache_key)
        return {
            'cache_key': self._any_cache_key,
            'cached_dttm': self._any_cached_dttm,
            'cache_timeout': self.cache_timeout,
            'df': df,
            'error': self.error_message,
            'form_data': self.form_data,
            'is_cached': self._any_cache_key is not None,
            'query': self.query,
            'status': self.status,
            'stacktrace': stacktrace,
            'rowcount': len(df.index) if df is not None else 0,
        }

    def json_dumps(self, obj, sort_keys=False):
        return json.dumps(
            obj,
            default=utils.json_int_dttm_ser,
            ignore_nan=True,
            sort_keys=sort_keys,
        )

    def payload_json_and_has_error(self, payload):
        has_error = payload.get('status') == utils.QueryStatus.FAILED or \
            payload.get('error') is not None
        return self.json_dumps(payload), has_error

    @property
    def data(self):
        """This is the data object serialized to the js layer"""
        content = {
            'form_data': self.form_data,
            'token': self.token,
            'viz_name': self.viz_type,
            'filter_select_enabled': self.datasource.filter_select_enabled,
        }
        return content

    def get_csv(self):
        df = self.get_df()
        include_index = not isinstance(df.index, pd.RangeIndex)
        return df.to_csv(index=include_index, **config.get('CSV_EXPORT'))

    def get_data(self, df):
        return df.to_dict(orient='records')

    @property
    def json_data(self):
        return json.dumps(self.data)


class TableViz(BaseViz):

    """A basic html table that is sortable and searchable"""

    viz_type = 'table'
    verbose_name = _('Table View')
    credits = 'a <a href="https://github.com/airbnb/superset">Superset</a> original'
    is_timeseries = False
    enforce_numerical_metrics = False

    def should_be_timeseries(self):
        fd = self.form_data
        # TODO handle datasource-type-specific code in datasource
        conditions_met = (
            (fd.get('granularity') and fd.get('granularity') != 'all') or
            (fd.get('granularity_sqla') and fd.get('time_grain_sqla'))
        )
        if fd.get('include_time') and not conditions_met:
            raise Exception(_(
                'Pick a granularity in the Time section or '
                "uncheck 'Include Time'"))
        return fd.get('include_time')

    def query_obj(self):
        d = super(TableViz, self).query_obj()
        fd = self.form_data

        if fd.get('all_columns') and (fd.get('groupby') or fd.get('metrics')):
            raise Exception(_(
                'Choose either fields to [Group By] and [Metrics] or '
                '[Columns], not both'))

        sort_by = fd.get('timeseries_limit_metric')
        if fd.get('all_columns'):
            d['columns'] = fd.get('all_columns')
            d['groupby'] = []
            order_by_cols = fd.get('order_by_cols') or []
            d['orderby'] = [json.loads(t) for t in order_by_cols]
        elif sort_by:
            sort_by_label = utils.get_metric_name(sort_by)
            if sort_by_label not in utils.get_metric_names(d['metrics']):
                d['metrics'] += [sort_by]
            d['orderby'] = [(sort_by, not fd.get('order_desc', True))]

        # Add all percent metrics that are not already in the list
        if 'percent_metrics' in fd:
            d['metrics'] = d['metrics'] + list(filter(
                lambda m: m not in d['metrics'],
                fd['percent_metrics'] or [],
            ))

        d['is_timeseries'] = self.should_be_timeseries()
        return d

    def get_data(self, df):
        fd = self.form_data
        if (
                not self.should_be_timeseries() and
                df is not None and
                DTTM_ALIAS in df
        ):
            del df[DTTM_ALIAS]

        # Sum up and compute percentages for all percent metrics
        percent_metrics = fd.get('percent_metrics') or []
        percent_metrics = [self.get_metric_label(m) for m in percent_metrics]

        if len(percent_metrics):
            percent_metrics = list(filter(lambda m: m in df, percent_metrics))
            metric_sums = {
                m: reduce(lambda a, b: a + b, df[m])
                for m in percent_metrics
            }
            metric_percents = {
                m: list(map(
                    lambda a: None if metric_sums[m] == 0 else a / metric_sums[m], df[m]))
                for m in percent_metrics
            }
            for m in percent_metrics:
                m_name = '%' + m
                df[m_name] = pd.Series(metric_percents[m], name=m_name)
            # Remove metrics that are not in the main metrics list
            metrics = fd.get('metrics') or []
            metrics = [self.get_metric_label(m) for m in metrics]
            for m in filter(
                lambda m: m not in metrics and m in df.columns,
                percent_metrics,
            ):
                del df[m]

        data = self.handle_js_int_overflow(
            dict(
                records=df.to_dict(orient='records'),
                columns=list(df.columns),
            ))

        return data

    def json_dumps(self, obj, sort_keys=False):
        return json.dumps(
            obj,
            default=utils.json_iso_dttm_ser,
            sort_keys=sort_keys,
            ignore_nan=True)


class TimeTableViz(BaseViz):

    """A data table with rich time-series related columns"""

    viz_type = 'time_table'
    verbose_name = _('Time Table View')
    credits = 'a <a href="https://github.com/airbnb/superset">Superset</a> original'
    is_timeseries = True

    def query_obj(self):
{'fraction_non_alphanumeric': 0.08936671706839812, 'fraction_numerical': 0.0008004980876990127, 'mean_word_length': 1.928310977991926, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 17, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3549589', 'n_tokens_mistral': 13056, 'n_tokens_neox': 12587, 'n_words': 3437}