xref: /external/autotest/client/common_lib/cros/graphite/es_utils.py

# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# This file defines helper functions for putting entries into elasticsearch.

"""Utils for sending metadata to elasticsearch

Elasticsearch is a key-value store NOSQL database.
Source is here: https://github.com/elasticsearch/elasticsearch
We will be using es to store our metadata.

For example, if we wanted to store the following metadata:

metadata = {
    'host_id': 1
    'job_id': 20
    'time_start': 100000
    'time_recorded': 100006
}

The following call will send metadata to the default es server.
    es_utils.ESMetadata().post(index, metadata)
We can also specify which port and host to use.

Using for testing: Sometimes, when we choose a single index
to put entries into, we want to clear that index of all
entries before running our tests. Use clear_index function.
(see es_utils_functionaltest.py for an example)

This file also contains methods for sending queries to es. Currently,
the query (json dict) we send to es is quite complicated (but flexible).
We've included several methods that composes queries that would be useful.
These methods are all named create_*_query()

For example, the below query returns job_id, host_id, and job_start
for all job_ids in [0, 99999] and host_id matching 10.

range_eq_query = {
    'fields': ['job_id', 'host_id', 'job_start'],
    'query': {
        'filtered': {
            'query': {
                'match': {
                    'host_id': 10,
                }
            }
            'filter': {
                'range': {
                    'job_id': {
                        'gte': 0,
                        'lte': 99999,
                    }
                }
            }
        }
    }
}

To send a query once it is created, call execute_query() to send it to the
intended elasticsearch server.

"""

import collections
import json
import logging
import socket
import time

try:
    import elasticsearch
    from elasticsearch import helpers as elasticsearch_helpers
except ImportError:
    logging.debug('Failed to import elasticsearch. Mock classes will be used '
                  'and calls to Elasticsearch server will be no-op. Test run '
                  'is not affected by the missing elasticsearch module.')
    import elasticsearch_mock as elasticsearch
    elasticsearch_helpers = elasticsearch.Elasticsearch()


DEFAULT_TIMEOUT = 30


class EsUtilException(Exception):
    """Exception raised when functions here fail. """
    pass


QueryResult = collections.namedtuple('QueryResult', ['total', 'hits'])


class ESMetadata(object):
    """Class handling es connection for metadata."""

    @property
    def es(self):
        """Read only property, lazily initialized"""
        if not self._es:
            self._es = elasticsearch.Elasticsearch(host=self.host,
                                                   port=self.port,
                                                   timeout=self.timeout)
        return self._es


    def __init__(self, use_http, host, port, index, udp_port,
                 timeout=DEFAULT_TIMEOUT):
        """Initialize ESMetadata object.

        @param use_http: Whether to send data to ES using HTTP.
        @param host: Elasticsearch host.
        @param port: Elasticsearch port.
        @param index: What index the metadata is stored in.
        @param udp_port: What port to use for UDP data.
        @param timeout: How long to wait while connecting to es.
        """
        self.use_http = use_http
        self.host = host
        self.port = port
        self.index = index
        self.udp_port = udp_port
        self.timeout = timeout
        self._es = None


    def _send_data_http(self, type_str, metadata):
        """Sends data to insert into elasticsearch using HTTP.

        @param type_str: sets the _type field in elasticsearch db.
        @param metadata: dictionary object containing metadata
        """
        try:
            self.es.index(index=self.index, doc_type=type_str, body=metadata)
        except elasticsearch.ElasticsearchException as e:
            # Mute exceptions from metadata reporting to prevent meta data
            # reporting errors from killing test.
            logging.error(e)


    def _send_data_udp(self, type_str, metadata):
        """Sends data to insert into elasticsearch using UDP.

        @param type_str: sets the _type field in elasticsearch db.
        @param metadata: dictionary object containing metadata
        """
        try:
            # Header.
            message = json.dumps(
                    {'index': {'_index': self.index, '_type': type_str}},
                    separators=(', ', ' : '))
            message += '\n'
            # Metadata.
            message += json.dumps(metadata, separators=(', ', ' : '))
            message += '\n'

            sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            sock.sendto(message, (self.host, self.udp_port))
        except socket.error as e:
            logging.warn(e)


    def post(self, type_str, metadata, log_time_recorded=True, **kwargs):
        """Wraps call of send_data, inserts entry into elasticsearch.

        @param type_str: Sets the _type field in elasticsearch db.
        @param metadata: Dictionary object containing metadata
        @param log_time_recorded: Whether to automatically record the time
                                  this metadata is recorded. Default is True.
        @param kwargs: Additional metadata fields

        @return: True if post action succeeded. Otherwise return False.

        """
        if not metadata:
            return True

        metadata = metadata.copy()
        metadata.update(kwargs)
        # metadata should not contain anything with key '_type'
        if '_type' in metadata:
            type_str = metadata['_type']
            del metadata['_type']
        if log_time_recorded:
            metadata['time_recorded'] = time.time()
        try:
            if self.use_http:
                self._send_data_http(type_str, metadata)
            else:
                self._send_data_udp(type_str, metadata)
            return True
        except elasticsearch.ElasticsearchException as e:
            logging.error(e)
            return False


    def bulk_post(self, data_list, log_time_recorded=True, **kwargs):
        """Wraps call of send_data, inserts entry into elasticsearch.

        @param data_list: A list of dictionary objects containing metadata.
        @param log_time_recorded: Whether to automatically record the time
                                  this metadata is recorded. Default is True.
        @param kwargs: Additional metadata fields

        @return: True if post action succeeded. Otherwise return False.

        """
        if not data_list:
            return True

        actions = []
        for metadata in data_list:
            metadata = metadata.copy()
            metadata.update(kwargs)
            if log_time_recorded and not 'time_recorded' in metadata:
                metadata['time_recorded'] = time.time()
            metadata['_index'] = self.index
            actions.append(metadata)

        try:
            elasticsearch_helpers.bulk(self.es, actions)
            return True
        except elasticsearch.ElasticsearchException as e:
            logging.error(e)
            return False


    def _compose_query(self, equality_constraints=[], fields_returned=None,
                       range_constraints=[], size=1000000, sort_specs=None,
                       regex_constraints=[], batch_constraints=[]):
        """Creates a dict. representing multple range and/or equality queries.

        Example input:
        _compose_query(
            fields_returned = ['time_recorded', 'hostname',
                               'status', 'dbg_str'],
            equality_constraints = [
                ('_type', 'host_history'),
                ('hostname', '172.22.169.106'),
            ],
            range_constraints = [
                ('time_recorded', 1405628341.904379, 1405700341.904379)
            ],
            size=20,
            sort_specs=[
                'hostname',
                {'time_recorded': 'asc'},
            ]
        )

        Output:
        {
            'fields': ['time_recorded', 'hostname', 'status', 'dbg_str'],
            'query': {
                'bool': {
                    'minimum_should_match': 3,
                    'should': [
                        {
                            'term':  {
                                '_type': 'host_history'
                            }
                        },
                        {
                            'term': {
                                'hostname': '172.22.169.106'
                            }
                        },
                        {
                            'range': {
                                'time_recorded': {
                                    'gte': 1405628341.904379,
                                    'lte': 1405700341.904379
                                }
                            }
                        }
                    ]
                },
            },
            'size': 20
            'sort': [
                'hostname',
                { 'time_recorded': 'asc'},
            ]
        }

        @param equality_constraints: list of tuples of (field, value) pairs
            representing what each field should equal to in the query.
            e.g. [ ('field1', 1), ('field2', 'value') ]
        @param fields_returned: list of fields that we should return when
            the query is executed. Set it to None to return all fields. Note
            that the key/vals will be stored in _source key of the hit object,
            if fields_returned is set to None.
        @param range_constraints: list of tuples of (field, low, high) pairs
            representing what each field should be between (inclusive).
            e.g. [ ('field1', 2, 10), ('field2', -1, 20) ]
            If you want one side to be unbounded, you can use None.
            e.g. [ ('field1', 2, None) ] means value of field1 >= 2.
        @param size: max number of entries to return. Default is 1000000.
        @param sort_specs: A list of fields to sort on, tiebreakers will be
            broken by the next field(s).
        @param regex_constraints: A list of regex constraints of tuples of
            (field, value) pairs, e.g., [('filed1', '.*value.*')].
        @param batch_constraints: list of tuples of (field, list) pairs
            representing each field should be equal to one of the values
            in the list.
            e.g., [ ('job_id', [10, 11, 12, 13]) ]
        @returns: dictionary object that represents query to es.
                  This will return None if there are no equality constraints
                  and no range constraints.
        """
        if not equality_constraints and not range_constraints:
            raise EsUtilException('No range or equality constraints specified.')

        # Creates list of range dictionaries to put in the 'should' list.
        range_list = []
        if range_constraints:
            for key, low, high in range_constraints:
                if low is None and high is None:
                    continue
                temp_dict = {}
                if low is not None:
                    temp_dict['gte'] = low
                if high is not None:
                    temp_dict['lte'] = high
                range_list.append( {'range': {key: temp_dict}})

        # Creates the list of term dictionaries to put in the 'should' list.
        eq_list = [{'term': {k: v}} for k, v in equality_constraints if k]
        batch_list = [{'terms': {k: v}} for k, v in batch_constraints if k]
        regex_list = [{'regexp': {k: v}} for k, v in regex_constraints if k]
        constraints = eq_list + batch_list + range_list + regex_list
        query = {
            'query': {
                'bool': {
                    'must': constraints,
                }
            },
        }
        if fields_returned:
            query['fields'] = fields_returned
        query['size'] = size
        if sort_specs:
            query['sort'] = sort_specs
        return query


    def execute_query(self, query):
        """Makes a query on the given index.

        @param query: query dictionary (see _compose_query)
        @returns: A QueryResult instance describing the result.

        Example output:
        {
            "took" : 5,
            "timed_out" : false,
            "_shards" : {
                "total" : 16,
                "successful" : 16,
                "failed" : 0
            },
            "hits" : {
                "total" : 4,
                "max_score" : 1.0,
                "hits" : [ {
                    "_index" : "graphite_metrics2",
                    "_type" : "metric",
                    "_id" : "rtntrjgdsafdsfdsfdsfdsfdssssssss",
                    "_score" : 1.0,
                    "_source":{"target_type": "timer",
                               "host_id": 1,
                               "job_id": 22,
                               "time_start": 400}
                }, {
                    "_index" : "graphite_metrics2",
                    "_type" : "metric",
                    "_id" : "dfgfddddddddddddddddddddddhhh",
                    "_score" : 1.0,
                    "_source":{"target_type": "timer",
                        "host_id": 2,
                        "job_id": 23,
                        "time_start": 405}
                }, {
                "_index" : "graphite_metrics2",
                "_type" : "metric",
                "_id" : "erwerwerwewtrewgfednvfngfngfrhfd",
                "_score" : 1.0,
                "_source":{"target_type": "timer",
                           "host_id": 3,
                           "job_id": 24,
                           "time_start": 4098}
                }, {
                    "_index" : "graphite_metrics2",
                    "_type" : "metric",
                    "_id" : "dfherjgwetfrsupbretowegoegheorgsa",
                    "_score" : 1.0,
                    "_source":{"target_type": "timer",
                               "host_id": 22,
                               "job_id": 25,
                               "time_start": 4200}
                } ]
            }
        }

        """
        if not self.es.indices.exists(index=self.index):
            logging.error('Index (%s) does not exist on %s:%s',
                          self.index, self.host, self.port)
            return None
        result = self.es.search(index=self.index, body=query,
                                timeout=DEFAULT_TIMEOUT,
                                request_timeout=DEFAULT_TIMEOUT)
        # Check if all matched records are returned. It could be the size is
        # set too small. Special case for size set to 1, as that means that
        # the query cares about the first matched entry.
        # TODO: Use pagination in Elasticsearch. This needs major change on how
        #       query results are iterated.
        size = query.get('size', 1)
        return_count = len(result['hits']['hits'])
        total_match = result['hits']['total']
        if total_match > return_count and size != 1:
            logging.error('There are %d matched records, only %d entries are '
                          'returned. Query size is set to %d.', total_match,
                          return_count, size)

        # Extract the actual results from the query.
        output = QueryResult(total_match, [])
        for hit in result['hits']['hits']:
            converted = {}
            if 'fields' in hit:
                for key, value in hit['fields'].items():
                    converted[key] = value[0] if len(value)==1 else value
            else:
                converted = hit['_source'].copy()
            output.hits.append(converted)
        return output


    def query(self, *args, **kwargs):
        """The arguments to this function are the same as _compose_query."""
        query = self._compose_query(*args, **kwargs)
        return self.execute_query(query)