elasticsearch.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# @Version     : Python 3.11.4
# @Software    : Sublime Text 4
# @Author      : StudentCWZ
# @Email       : StudentCWZ@outlook.com
# @Date        : 2023/10/28 12:19
# @File        : elasticsearch.py
# @Description :
"""

from typing import Any

from elasticsearch import Elasticsearch
from flask import current_app, Flask
from loguru import logger

from application.common import ConfigHelper


class FlaskElasticsearch:
    def __init__(self, app=None):
        """
        Initialize a FlaskElasticsearch instance.

        :param app: Flask application instance for Elasticsearch initialization.
        """
        if app is not None:
            self.init_app(app)

    def init_app(self, app: Flask) -> None:
        """
        Initialize the app with Elasticsearch connection.

        :param app: Flask application instance for Elasticsearch initialization.
        """
        if hasattr(app, 'teardown_appcontext'):
            app.teardown_appcontext(self.teardown)
        else:
            app.teardown_request(self.teardown)

    @staticmethod
    def teardown(exception):
        """
        Clear the Elasticsearch connection after each request.

        :param exception: Exception instance, if any.
        """
        ctx = current_app._get_current_object()
        if hasattr(ctx, 'elasticsearch'):
            ctx.elasticsearch = None
        if exception is not None:
            raise RuntimeError(exception)

    def __getattr__(self, item: Any) -> Any:
        """
        Lazy initialization of Elasticsearch connection on first use.

        :param item: Attribute to get from the Elasticsearch connection.
        :return: Value of the requested attribute.
        """
        ctx = current_app._get_current_object()
        if ctx is not None:
            if not hasattr(ctx, 'elasticsearch'):
                cfg = self._get_config()
                ctx.elasticsearch = Elasticsearch(**cfg)
                if ctx.elasticsearch.ping():
                    logger.info('Connected to Elasticsearch')
                else:
                    logger.error('Can not connect to Elasticsearch')
                    raise ConnectionError('Can not connect to Elasticsearch')
            return getattr(ctx.elasticsearch, item)

    @staticmethod
    def _get_config():
        """Retrieves Elasticsearch configuration from the current Flask application context."""
        with current_app.app_context():
            if current_app:
                config_helper = ConfigHelper(current_app)
                cfg = config_helper.Elasticsearch
                host = cfg.Host or 'localhost'
                port = int(cfg.Port) or 9200
                user = cfg.User or None
                password = cfg.Password or None
                use_ssl = cfg.UseSsl == 'True' or False
                verify_certs = cfg.VerifyCerts == 'False'
                ca_certs = cfg.CaCerts or None

                options = dict(
                    hosts=[{'host': host, 'port': port}],
                    http_auth=None if not user else (user, password),
                    use_ssl=use_ssl,
                    verify_certs=verify_certs,
                    ca_certs=ca_certs
                )

                return options
            else:
                logger.error('Attempted to access application configuration outside of application context.')
                raise RuntimeError('Attempted to access application configuration outside of application context.')