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: '*'});
  }
}