#!/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 class FlaskElasticsearch: def __init__(self, app=None): if app is not None: self.init_app(app) def init_app(self, app: Flask) -> None: """Initialize the app""" if hasattr(app, 'teardown_appcontext'): app.teardown_appcontext(self.teardown) else: app.teardown_request(self.teardown) @staticmethod def teardown(exception): """Clears the Elasticsearch connection after each request.""" 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.""" 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('Connected to Elasticsearch') raise ConnectionError('Connected 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: # Retrieve configuration from current_app.config and return it host = current_app.config.Elasticsearch.Host if current_app.config.get( 'Elasticsearch') is not None else 'localhost' port = int(current_app.config.Elasticsearch.Port) if current_app.config.get( 'Elasticsearch') is not None else 9200 user = current_app.config.Elasticsearch.User if current_app.config.get( 'Elasticsearch') is not None else 'user' password = current_app.config.Elasticsearch.Password if current_app.config.get( 'Elasticsearch') is not None else 'password' use_ssl = current_app.config.Elasticsearch.UseSsl == 'True' if current_app.config.get( 'Elasticsearch') is not None else False verify_certs = current_app.config.Elasticsearch.VerifyCerts == 'False' if current_app.config.get( 'Elasticsearch') is not None else True ca_certs = current_app.config.Elasticsearch.CaCerts if current_app.config.get( 'Elasticsearch') is not None else None es_config = 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 es_config else: logger.error('Attempted to access application configuration outside of application context.') raise RuntimeError('Attempted to access application configuration outside of application context.')