tools/plugins/elastic.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
    common elasticsearch database setup
    also adds defaults for most methods
"""

import sys
import logging
import certifi
from . import ponymailconfig

try:
    from elasticsearch import Elasticsearch, helpers, AsyncElasticsearch
    from elasticsearch import VERSION as ES_VERSION
    from elasticsearch import ConnectionError as ES_ConnectionError
except ImportError as e:
    sys.exit(
        "Sorry, you need to install the elasticsearch module from pip first. (%s)"
        % str(e)
    )


class Elastic:
    db_mbox:            str
    db_source:          str
    db_attachment:      str
    db_account:         str
    db_session:         str
    db_notification:    str
    db_auditlog:        str
    dbname:             str

    def __init__(self, logger_level=None, trace_level=None, is_async=False):
        # Fetch config
        config = ponymailconfig.PonymailConfig()

        # Set default names for all indices we use
        dbname = config.get('elasticsearch', 'dbname', fallback='ponymail')
        self.dbname = dbname
        self.db_mbox = dbname + '-mbox'
        self.db_source = dbname + '-source'
        self.db_account = dbname + '-account'
        self.db_attachment = dbname + '-attachment'
        self.db_session = dbname + '-session'
        self.db_notification = dbname + '-notification'
        self.db_auditlog = dbname + '-auditlog'
        self.db_version = 0
        self.is_async = is_async

        dburl = config.get('elasticsearch', 'dburl', fallback=None)
        if not dburl:
            ssl = config.get('elasticsearch', 'ssl', fallback=False)
            uri = config.get('elasticsearch', 'uri', fallback='')
            auth = None
            if config.has_option('elasticsearch', 'user'):
                auth = (
                    config.get('elasticsearch', 'user'),
                    config.get('elasticsearch', 'password')
                )
            dburl = {
                "host": config.get('elasticsearch', 'hostname', fallback='localhost'),
                "port": config.get('elasticsearch', 'port', fallback=9200),
                "use_ssl": ssl,
                "url_prefix": uri,
                "auth": auth,
                "ca_certs": certifi.where(),
            }

        # Always allow this to be set; will be replaced as necessary by wait_for_active_shards
        self.consistency = config.get("elasticsearch", "write", fallback="quorum")

        if logger_level:
            eslog = logging.getLogger("elasticsearch")
            eslog.setLevel(logger_level)
            eslog.addHandler(logging.StreamHandler())
        else:
            # elasticsearch logs lots of warnings on retries/connection failure
            logging.getLogger("elasticsearch").setLevel(logging.ERROR)

        if trace_level:
            trace = logging.getLogger("elasticsearch.trace")
            trace.setLevel(trace_level)
            trace.addHandler(logging.StreamHandler())
        if self.is_async:
            self.es = AsyncElasticsearch(
                [
                dburl
                ],
                max_retries=5,
                retry_on_timeout=True,
            )
        else:
            self.es = Elasticsearch(
                [
                    dburl
                ],
                max_retries=5,
                retry_on_timeout=True,
            )
            # This won't work with async, so for now we'll ignore it there...
            es_engine_major = self.engineMajor()
            if es_engine_major in [7, 8]:
                self.wait_for_active_shards = config.get("elasticsearch", "wait", fallback=1)
            else:
                raise Exception("Unexpected elasticsearch version ", es_engine_major)



        # Mimic ES hierarchy: es.indices.xyz()
        self.indices = _indices_wrap(self)

    # convert index type to index name
    def index_name(self, index):
        return self.dbname + "-" + index

    @staticmethod
    def libraryVersion():
        return ES_VERSION

    @staticmethod
    def libraryMajor():
        return ES_VERSION[0]

    def engineVersion(self):
        if not self.db_version:
            try:
                self.db_version = self.es.info()["version"]["number"]
            except ES_ConnectionError:
                # default if cannot connect; allows retry
                return "0.0.0"
        return self.db_version

    def engineMajor(self):
        return int(self.engineVersion().split(".")[0])

    def search(self, **kwargs):
        return self.es.search(**kwargs)

    def index(self, **kwargs):
        kwargs["wait_for_active_shards"] = self.wait_for_active_shards
        kwargs["doc_type"] = "_doc"
        return self.es.index(**kwargs)

    def create(self, **kwargs):
        return self.es.create(**kwargs)

    def info(self, **kwargs):
        return self.es.info(**kwargs)

    def update(self, **kwargs):
        return self.es.update(**kwargs)

    # TODO: is this used? Does it make sense for ES7 ?
    def scan(self, scroll="3m", size=100, **kwargs):
        return self.es.search(
            search_type="scan", size=size, scroll=scroll, **kwargs
        )

    def get(self, **kwargs):
        return self.es.get(**kwargs)

    def scroll(self, **kwargs):
        return self.es.scroll(**kwargs)

    def bulk(self, actions, **kwargs):
        return helpers.bulk(self.es, actions, **kwargs)

    def streaming_bulk(self, actions, **kwargs):
        return helpers.streaming_bulk(self.es, actions, **kwargs)

    def clear_scroll(self, *args, **kwargs):
        """
            Call this to release the scroll id and its resources

            It looks like the Python library already releases the SID
            if the caller scrolls to the end of the results, so only need to call this
            when terminating scrolling early.
        """
        return self.es.clear_scroll(*args, **kwargs)
class _indices_wrap:
    """
        Wrapper for the ES indices methods we use
    """

    def __init__(self, parent):
        self.es = parent.es

    def exists(self, *args, **kwargs):
        return self.es.indices.exists(*args, **kwargs)

    def create(self, *args, **kwargs):
        return self.es.indices.create(*args, **kwargs)

    def get_mapping(self, **kwargs):
        return self.es.indices.get_mapping(**kwargs)

    def put_mapping(self, **kwargs):
        return self.es.indices.put_mapping(**kwargs)