energylink/public/app/plugins/datasource/elasticsearch/datasource.ts

///<reference path="../../../headers/common.d.ts" />

import angular from 'angular';
import _ from 'lodash';
import moment from 'moment';
import { ElasticQueryBuilder } from './query_builder';
import { IndexPattern } from './index_pattern';
import { ElasticResponse } from './elastic_response';

export class ElasticDatasource {
  basicAuth: string;
  withCredentials: boolean;
  url: string;
  name: string;
  index: string;
  timeField: string;
  esVersion: number;
  interval: string;
  maxConcurrentShardRequests: number;
  queryBuilder: ElasticQueryBuilder;
  indexPattern: IndexPattern;

  /** @ngInject */
  constructor(
    instanceSettings,
    private $q,
    private backendSrv,
    private templateSrv,
    private timeSrv
  ) {
    this.basicAuth = instanceSettings.basicAuth;
    this.withCredentials = instanceSettings.withCredentials;
    this.url = instanceSettings.url;
    this.name = instanceSettings.name;
    this.index = instanceSettings.index;
    this.timeField = instanceSettings.jsonData.timeField;
    this.esVersion = instanceSettings.jsonData.esVersion;
    this.indexPattern = new IndexPattern(
      instanceSettings.index,
      instanceSettings.jsonData.interval
    );
    this.interval = instanceSettings.jsonData.timeInterval;
    this.maxConcurrentShardRequests =
      instanceSettings.jsonData.maxConcurrentShardRequests;
    this.queryBuilder = new ElasticQueryBuilder({
      timeField: this.timeField,
      esVersion: this.esVersion,
    });
  }

  private request(method, url, data?) {
    var options: any = {
      url: this.url + '/' + url,
      method: method,
      data: data,
    };

    if (this.basicAuth || this.withCredentials) {
      options.withCredentials = true;
    }
    if (this.basicAuth) {
      options.headers = {
        Authorization: this.basicAuth,
      };
    }

    return this.backendSrv.datasourceRequest(options);
  }

  private get(url) {
    var range = this.timeSrv.timeRange();
    var index_list = this.indexPattern.getIndexList(
      range.from.valueOf(),
      range.to.valueOf()
    );
    if (_.isArray(index_list) && index_list.length) {
      return this.request('GET', index_list[0] + url).then(function(results) {
        results.data.$$config = results.config;
        return results.data;
      });
    } else {
      return this.request(
        'GET',
        this.indexPattern.getIndexForToday() + url
      ).then(function(results) {
        results.data.$$config = results.config;
        return results.data;
      });
    }
  }

  private post(url, data) {
    return this.request('POST', url, data)
      .then(function(results) {
        results.data.$$config = results.config;
        return results.data;
      })
      .catch(err => {
        if (err.data && err.data.error) {
          throw {
            message: 'Elasticsearch error: ' + err.data.error.reason,
            error: err.data.error,
          };
        }

        throw err;
      });
  }

  annotationQuery(options) {
    var annotation = options.annotation;
    var timeField = annotation.timeField || '@timestamp';
    var queryString = annotation.query || '*';
    var tagsField = annotation.tagsField || 'tags';
    var textField = annotation.textField || null;

    var range = {};
    range[timeField] = {
      from: options.range.from.valueOf(),
      to: options.range.to.valueOf(),
      format: 'epoch_millis',
    };

    var queryInterpolated = this.templateSrv.replace(queryString, {}, 'lucene');
    var query = {
      bool: {
        filter: [
          { range: range },
          {
            query_string: {
              query: queryInterpolated,
            },
          },
        ],
      },
    };

    var data = {
      query: query,
      size: 10000,
    };

    // fields field not supported on ES 5.x
    if (this.esVersion < 5) {
      data['fields'] = [timeField, '_source'];
    }

    var header: any = {
      search_type: 'query_then_fetch',
      ignore_unavailable: true,
    };

    // old elastic annotations had index specified on them
    if (annotation.index) {
      header.index = annotation.index;
    } else {
      header.index = this.indexPattern.getIndexList(
        options.range.from,
        options.range.to
      );
    }

    var payload = angular.toJson(header) + '\n' + angular.toJson(data) + '\n';

    return this.post('_msearch', payload).then(res => {
      var list = [];
      var hits = res.responses[0].hits.hits;

      var getFieldFromSource = function(source, fieldName) {
        if (!fieldName) {
          return;
        }

        var fieldNames = fieldName.split('.');
        var fieldValue = source;

        for (var i = 0; i < fieldNames.length; i++) {
          fieldValue = fieldValue[fieldNames[i]];
          if (!fieldValue) {
            console.log('could not find field in annotation: ', fieldName);
            return '';
          }
        }

        return fieldValue;
      };

      for (var i = 0; i < hits.length; i++) {
        var source = hits[i]._source;
        var time = source[timeField];
        if (typeof hits[i].fields !== 'undefined') {
          var fields = hits[i].fields;
          if (_.isString(fields[timeField]) || _.isNumber(fields[timeField])) {
            time = fields[timeField];
          }
        }

        var event = {
          annotation: annotation,
          time: moment.utc(time).valueOf(),
          text: getFieldFromSource(source, textField),
          tags: getFieldFromSource(source, tagsField),
        };

        // legacy support for title tield
        if (annotation.titleField) {
          const title = getFieldFromSource(source, annotation.titleField);
          if (title) {
            event.text = title + '\n' + event.text;
          }
        }

        if (typeof event.tags === 'string') {
          event.tags = event.tags.split(',');
        }

        list.push(event);
      }
      return list;
    });
  }

  testDatasource() {
    this.timeSrv.setTime({ from: 'now-1m', to: 'now' }, true);
    // validate that the index exist and has date field
    return this.getFields({ type: 'date' }).then(
      function(dateFields) {
        var timeField = _.find(dateFields, { text: this.timeField });
        if (!timeField) {
          return {
            status: 'error',
            message: 'No date field named ' + this.timeField + ' found',
          };
        }
        return { status: 'success', message: 'Index OK. Time field name OK.' };
      }.bind(this),
      function(err) {
        console.log(err);
        if (err.data && err.data.error) {
          var message = angular.toJson(err.data.error);
          if (err.data.error.reason) {
            message = err.data.error.reason;
          }
          return { status: 'error', message: message };
        } else {
          return { status: 'error', message: err.status };
        }
      }
    );
  }

  getQueryHeader(searchType, timeFrom, timeTo) {
    var query_header: any = {
      search_type: searchType,
      ignore_unavailable: true,
      index: this.indexPattern.getIndexList(timeFrom, timeTo),
    };
    if (this.esVersion >= 56) {
      query_header[
        'max_concurrent_shard_requests'
      ] = this.maxConcurrentShardRequests;
    }
    return angular.toJson(query_header);
  }

  query(options) {
    var payload = '';
    var target;
    var sentTargets = [];

    // add global adhoc filters to timeFilter
    var adhocFilters = this.templateSrv.getAdhocFilters(this.name);

    for (var i = 0; i < options.targets.length; i++) {
      target = options.targets[i];
      if (target.hide) {
        continue;
      }

      var queryString = this.templateSrv.replace(
        target.query || '*',
        options.scopedVars,
        'lucene'
      );
      var queryObj = this.queryBuilder.build(target, adhocFilters, queryString);
      var esQuery = angular.toJson(queryObj);

      var searchType =
        queryObj.size === 0 && this.esVersion < 5
          ? 'count'
          : 'query_then_fetch';
      var header = this.getQueryHeader(
        searchType,
        options.range.from,
        options.range.to
      );
      payload += header + '\n';

      payload += esQuery + '\n';
      sentTargets.push(target);
    }

    if (sentTargets.length === 0) {
      return this.$q.when([]);
    }

    payload = payload.replace(/\$timeFrom/g, options.range.from.valueOf());
    payload = payload.replace(/\$timeTo/g, options.range.to.valueOf());
    payload = this.templateSrv.replace(payload, options.scopedVars);

    return this.post('_msearch', payload).then(function(res) {
      return new ElasticResponse(sentTargets, res).getTimeSeries();
    });
  }

  getFields(query) {
    return this.get('/_mapping').then(function(result) {
      var typeMap = {
        float: 'number',
        double: 'number',
        integer: 'number',
        long: 'number',
        date: 'date',
        string: 'string',
        text: 'string',
        scaled_float: 'number',
        nested: 'nested',
      };

      function shouldAddField(obj, key, query) {
        if (key[0] === '_') {
          return false;
        }

        if (!query.type) {
          return true;
        }

        // equal query type filter, or via typemap translation
        return query.type === obj.type || query.type === typeMap[obj.type];
      }

      // Store subfield names: [system, process, cpu, total] -> system.process.cpu.total
      var fieldNameParts = [];
      var fields = {};

      function getFieldsRecursively(obj) {
        for (var key in obj) {
          var subObj = obj[key];

          // Check mapping field for nested fields
          if (_.isObject(subObj.properties)) {
            fieldNameParts.push(key);
            getFieldsRecursively(subObj.properties);
          }

          if (_.isObject(subObj.fields)) {
            fieldNameParts.push(key);
            getFieldsRecursively(subObj.fields);
          }

          if (_.isString(subObj.type)) {
            var fieldName = fieldNameParts.concat(key).join('.');

            // Hide meta-fields and check field type
            if (shouldAddField(subObj, key, query)) {
              fields[fieldName] = {
                text: fieldName,
                type: subObj.type,
              };
            }
          }
        }
        fieldNameParts.pop();
      }

      for (var indexName in result) {
        var index = result[indexName];
        if (index && index.mappings) {
          var mappings = index.mappings;
          for (var typeName in mappings) {
            var properties = mappings[typeName].properties;
            getFieldsRecursively(properties);
          }
        }
      }

      // transform to array
      return _.map(fields, function(value) {
        return value;
      });
    });
  }

  getTerms(queryDef) {
    var range = this.timeSrv.timeRange();
    var searchType = this.esVersion >= 5 ? 'query_then_fetch' : 'count';
    var header = this.getQueryHeader(searchType, range.from, range.to);
    var esQuery = angular.toJson(this.queryBuilder.getTermsQuery(queryDef));

    esQuery = esQuery.replace(/\$timeFrom/g, range.from.valueOf());
    esQuery = esQuery.replace(/\$timeTo/g, range.to.valueOf());
    esQuery = header + '\n' + esQuery + '\n';

    return this.post('_msearch?search_type=' + searchType, esQuery).then(
      function(res) {
        if (!res.responses[0].aggregations) {
          return [];
        }

        var buckets = res.responses[0].aggregations['1'].buckets;
        return _.map(buckets, function(bucket) {
          return {
            text: bucket.key_as_string || bucket.key,
            value: bucket.key,
          };
        });
      }
    );
  }

  metricFindQuery(query) {
    query = angular.fromJson(query);
    if (!query) {
      return this.$q.when([]);
    }

    if (query.find === 'fields') {
      query.field = this.templateSrv.replace(query.field, {}, 'lucene');
      return this.getFields(query);
    }

    if (query.find === 'terms') {
      query.query = this.templateSrv.replace(query.query || '*', {}, 'lucene');
      return this.getTerms(query);
    }
  }

  getTagKeys() {
    return this.getFields({});
  }

  getTagValues(options) {
    return this.getTerms({ field: options.key, query: '*' });
  }
}