utils/es.py

"""
All functions related to the ES client are placed here.
"""

from elasticsearch import Elasticsearch
from elasticsearch.client import IngestClient

import json
import os.path

from utils.tsdb import *


def get_client(elasticsearch_host, elasticsearch_ca_path, elasticsearch_user, elasticsearch_pwd, cloud_id, cloud_pwd):
    """
    Create ES client.
    If cloud values are provided, they will take priority over the local deployment.
    :param elasticsearch_host: ES host.
    :param elasticsearch_ca_path: Path to ES certificate.
    :param elasticsearch_user: Name of the ES user.
    :param elasticsearch_pwd: Password for ES.
    :param cloud_id: Cloud ID. Default is empty.
    :param cloud_pwd: Password for the elastic cloud. Default is empty.
    :return: ES client.
    """
    if cloud_id != "" and cloud_pwd != "":
        print("Client will connect to the cloud.")
        return Elasticsearch(
            cloud_id=cloud_id,
            basic_auth=("elastic", cloud_pwd)
        )
    return Elasticsearch(
        hosts=elasticsearch_host,
        ca_certs=elasticsearch_ca_path,
        basic_auth=(elasticsearch_user, elasticsearch_pwd)
    )


def add_doc_from_file(client: Elasticsearch, index_name: str, doc_path: str):
    """
    Given a JSON file, add the document to the index.
    This function does not check if the file exists, since that requirement was already
    checked before it was called.
    :param client: ES client.
    :param index_name: name of the index to place the document.
    :param doc_path: path to the document to add.
    """
    file = open(doc_path)
    content = json.load(file)
    file.close()
    client.index(index=index_name, document=content)


def place_documents(client: Elasticsearch, index_name: str, folder_docs: str):
    """
    Place all documents from folder to an index.
    :param client: ES client.
    :param index_name: name of the index to add the documents.
    :param folder_docs: path to the folder with the documents to add.
    """
    print("Placing documents on the index {name}...".format(name=index_name))
    if not client.indices.exists(index=index_name):
        print("Index {name} does not exist. Program will end.".format(name=index_name))
        exit(0)

    if not os.path.isdir(folder_docs):
        print("Folder {} does not exist. Documents cannot be placed. Program will end.".format(folder_docs))
        exit(0)

    for doc in os.listdir(folder_docs):
        doc_path = os.path.join(folder_docs, doc)
        if os.path.isfile(doc_path):
            add_doc_from_file(client, index_name, doc_path)

    # From Elastic docs: Use the refresh API to explicitly make all operations performed on one or more indices since
    # the last refresh available for search. If the request targets a data stream, it refreshes the stream’s backing
    # indices.
    client.indices.refresh(index=index_name)
    resp = client.search(index=index_name, query={"match_all": {}})
    n_docs = resp['hits']['total']['value']
    print("Successfully placed {} documents on the index {name}.\n".format(n_docs, name=index_name))


def create_index(client: Elasticsearch, index_name: str, mappings: {} = {}, settings: {} = {}):
    """
    Create new ES index. If the index already exists, it will be deleted and a new one created.
    :param client: ES client.
    :param index_name: name of the index.
    :param mappings: mappings to be used for the new index. If not specified, default ones will be used.
    :param settings: settings to be used for the new index. If not specified, default ones will be used.
    """
    if client.indices.exists(index=index_name):
        client.indices.delete(index=index_name)
    client.indices.create(index=index_name, mappings=mappings, settings=settings)
    print("Index {name} successfully created.\n".format(name=index_name))


def create_index_missing_for_docs(client: Elasticsearch):
    """
    Create an index to place all the documents that were updated at least one time.
    :param client: ES client.
    """
    create_index(client, overwritten_docs_index)
    pipelines = IngestClient(client)
    pipeline_name = 'get-missing-docs'
    pipelines.put_pipeline(id=pipeline_name, body={
        'description': "Drop all documents that were not overwritten.",
        "processors": [
            {
                "drop": {
                    "if": "ctx._version == 1"
                }
            }
        ]
    })
    dest = {
        "index": overwritten_docs_index,
        "version_type": "external",
        "pipeline": pipeline_name
    }
    client.reindex(source={"index": tsdb_index}, dest=dest, refresh=True)


def build_query(dimensions_exist: {}, dimensions_missing: []):
    """
    Build query to retrieve document based on the dimensions.
    :param dimensions_exist: Dictionary containing field and field value of the dimensions that exist in the document.
    :param dimensions_missing: Dictionary of the dimension fields missing in the document.
    :return: query.
    """
    query = {
        "bool": {
            "must": [],
            "must_not": []
        }
    }
    for field in dimensions_exist:
        term = {
            "term": {
                field: dimensions_exist[field]
            }
        }
        query["bool"]["must"].append(term)

    for field in dimensions_missing:
        exists = {
            "exists": {
                "field": field
            }
        }
        query["bool"]["must_not"].append(exists)

    return query


def get_and_place_documents(client: Elasticsearch, data_stream: str, dir_name: str, dimensions_values: {},
                            dimensions_missing: [], n: int, number_of_docs):
    """
    Given the dimensions, place the documents in the directory.
    :param client: ES client.
    :param data_stream: Name of the data stream.
    :param dir_name: Name of the parent directory.
    :param dimensions_values: Values for the dimension fields that exist in the document.
    :param dimensions_missing: List of dimension fields missing in the document.
    :param n: Number of the directory inside the parent directory. Example: 1 would create dir_name/1.
    :param number_of_docs: Number of documents to get with that set of dimensions.
    :param index_name: Index name in which documents are stored.
    """
    query = build_query(dimensions_values, dimensions_missing)

    res = client.search(index=data_stream, query=query, sort={"@timestamp": "asc"}, size=number_of_docs)

    dir_for_docs = os.path.join(dir_name, str(n))
    os.mkdir(dir_for_docs)

    for doc in res["hits"]["hits"]:
        name = doc["_id"] + ".json"
        with open(os.path.join(dir_for_docs, name), 'w') as file:
            json.dump(doc, file, indent=4)


def get_missing_docs_info(client: Elasticsearch, data_stream: str, display_docs: int, dir,
                          get_overlapping_files: bool, copy_docs_per_dimension: int):
    """
    Display the dimensions of the first @display_docs documents.
    If @get_overlapping_files is set to True, then @copy_docs_per_dimension documents will be placed in a directory
    (if the directory does not exist!).
    :param client: ES client.
    :param display_docs: number of documents to display.
    :param dir: name of the directory.
    :param get_overlapping_files: true if you want to place fields in the directory, false otherwise.
    :param copy_docs_per_dimension: number of documents to get for a set of dimensions.
    :param docs_index: name of the index with the documents.
    """
    if get_overlapping_files:
        if os.path.exists(dir):
            print("WARNING: The directory {} exists. Please delete it. Documents will not be placed.\n".format(dir))
            get_overlapping_files = False
        else:
            os.mkdir(dir)
            n = 1

    body = {'size': display_docs, 'query': {'match_all': {}}}
    res = client.search(index=overwritten_docs_index, body=body)
    dimensions = time_series_fields["dimension"]

    print("The timestamp and dimensions of the first {} overwritten documents are:".format(display_docs))
    for doc in res["hits"]["hits"]:
        if get_overlapping_files:
            dimensions_values = {"@timestamp": doc["_source"]["@timestamp"]}
            dimensions_missing = []

        print("- Timestamp {}:".format(doc["_source"]["@timestamp"]))
        for dimension in dimensions:
            el = doc["_source"]
            keys = dimension.split(".")
            for key in keys:
                if key not in el:
                    el = "(Missing value)"
                    break
                el = el[key]
            print("\t{}: {}".format(dimension, el))
            if get_overlapping_files:
                if el != "(Missing value)":
                    dimensions_values[dimension] = el
                else:
                    dimensions_missing.append(dimension)

        if get_overlapping_files:
            get_and_place_documents(client, data_stream, dir, dimensions_values, dimensions_missing, n, copy_docs_per_dimension)
            n += 1


def copy_docs_from_to(client: Elasticsearch, source_index: str, dest_index: str, max_docs: int):
    """
    Copy documents from one index to the other.
    :param client: ES client.
    :param source_index: source index with the documents to be copied to a new index.
    :param dest_index: destination index for the documents.
    :param max_docs: max number of documents to copy.
    :return: True if the number of documents is the same in the new index as it was in the old index.
    """
    print("Copying documents from {} to {}...".format(source_index, dest_index))
    if not client.indices.exists(index=source_index):
        print("Source index {name} does not exist. Program will end.".format(name=source_index))
        exit(0)

    if max_docs != -1:
        resp = client.reindex(source={"index": source_index}, dest={"index": dest_index}, refresh=True,
                              max_docs=max_docs)
    else:
        resp = client.reindex(source={"index": source_index}, dest={"index": dest_index}, refresh=True)
    if resp["updated"] > 0:
        print("WARNING: Out of {} documents from the index {}, {} of them were discarded.\n".format(resp["total"],
                                                                                                    source_index,
                                                                                                    resp[
                                                                                                        "updated"]))
        return False
    else:
        print(
            "All {} documents taken from index {} were successfully placed to index {}.\n".format(resp["total"],
                                                                                                  source_index,
                                                                                                  dest_index))
        return True


def get_tsdb_config(client: Elasticsearch, data_stream_name: str, docs_index: int, settings_mappings_index: int):
    """
    Get the index name where documents are placed, and mappings and settings for the new TSDB index.
    :param client: ES client.
    :param data_stream_name:
    :param docs_index: number of the index in the data stream with the documents to be moved to the TSDB index.
    :param settings_mappings_index: number of the index for the settings and mappings for the TSDB index.
    :return: documents index name, settings and mappings for the TSDB index.
    """
    data_stream = client.indices.get_data_stream(name=data_stream_name)
    n_indexes = len(data_stream["data_streams"][0]["indices"])

    # Get the index to use for document retrieval
    if docs_index == -1:
        docs_index = 0
    elif docs_index >= n_indexes:
        print("ERROR: Data stream {} has {} indexes. Documents index number {} is not valid.".format(data_stream_name,
                                                                                                     n_indexes,
                                                                                                     docs_index))
        exit(0)

    # Get index to use for settings/mappings
    if settings_mappings_index == -1:
        settings_mappings_index = n_indexes - 1
    elif settings_mappings_index >= n_indexes:
        print("ERROR: Data stream {} has {} indexes. Settings/mappings index number {} is not valid.".format(
            data_stream_name, n_indexes, settings_mappings_index))
        exit(0)

    docs_index_name = data_stream["data_streams"][0]["indices"][docs_index]["index_name"]
    settings_mappings_index_name = data_stream["data_streams"][0]["indices"][settings_mappings_index]["index_name"]

    print("Index being used for the documents is {}.".format(docs_index_name))
    print("Index being used for the settings and mappings is {}.".format(settings_mappings_index_name))
    print()

    mappings = client.indices.get_mapping(index=settings_mappings_index_name)[settings_mappings_index_name]["mappings"]
    settings = client.indices.get_settings(index=settings_mappings_index_name)[settings_mappings_index_name]["settings"]

    settings = get_tsdb_settings(mappings, settings)

    return docs_index_name, mappings, settings


def copy_from_data_stream(client: Elasticsearch, data_stream_name: str, docs_index: int,settings_mappings_index: int,
                          max_docs: int):
    """
    Given a data stream, it copies the documents retrieved from the given index and places them in a new
    index with TSDB enabled.
    :param client: ES client.
    :param data_stream_name: name of the data stream.
    :param docs_index: number of the index to use to retrieve the documents.
    :param settings_mappings_index: number of the index to use to get the mappings and settings for the TSDB index.
    :param max_docs: maximum documents to be reindexed.
    :return: True if the number of documents placed to the TSDB index remained the same. False otherwise.
    """
    print("Testing data stream {}.".format(data_stream_name))

    if not client.indices.exists(index=data_stream_name):
        print("\tData stream {} does not exist. Program will end.".format(data_stream_name))
        exit(0)

    source_index, mappings, settings = get_tsdb_config(client, data_stream_name, docs_index, settings_mappings_index)

    create_index(client, tsdb_index, mappings, settings)

    return copy_docs_from_to(client, source_index, tsdb_index, max_docs)