Update project

application/libs/flask_elasticsearch/elasticsearch.py

@@ -14,46 +14,47 @@ import time
 from typing import Any
 
 from elasticsearch import Elasticsearch
-from flask import current_app, Flask
+from flask import Flask
+from flask_elasticsearch import FlaskElasticsearch as BaseFlaskElasticsearch
 from loguru import logger
 
 from application.common import ConfigHelper
 
 
-class FlaskElasticsearch:
-    def __init__(self, app=None):
-        """
-        Initialize a FlaskElasticsearch instance.
+# Find the stack on which we want to store the database connection.
+# Starting with Flask 0.9, the _app_ctx_stack is the correct one,
+# before that we need to use the _request_ctx_stack.
+try:
+    from flask import _app_ctx_stack as stack_context
+except ImportError:
+    from flask import _request_ctx_stack as stack_context
 
-        :param app: Flask application instance for Elasticsearch initialization.
+
+class FlaskElasticsearch(BaseFlaskElasticsearch):
+    def __init__(self, app=None, **kwargs):
+        """
+        Initializes a FlaskElasticsearch instance.
+        :param app: Flask app instance for Elasticsearch initialization.
         """
-        if app is not None:
-            self.init_app(app)
+        self.extra_options = None
+        super().__init__(app, **kwargs)
 
-    def init_app(self, app: Flask) -> None:
+    def init_app(self, app: Flask, **kwargs) -> None:
         """
         Initialize the app with Elasticsearch connection.
 
         :param app: Flask application instance for Elasticsearch initialization.
         """
-        if hasattr(app, 'teardown_appcontext'):
+
+        self.extra_options = kwargs
+
+        # Use the new_style teardown_appcontext if it's available,
+        # otherwise fall back to the request context
+        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.
@@ -61,11 +62,11 @@ class FlaskElasticsearch:
         :param item: Attribute to get from the Elasticsearch connection.
         :return: Value of the requested attribute.
         """
-        ctx = current_app._get_current_object()
+        ctx = stack_context.top
         if ctx is not None:
             if not hasattr(ctx, 'elasticsearch'):
-                cfg = self._get_config()
-                ctx.elasticsearch = Elasticsearch(**cfg)
+                cfg = self._get_config(ctx.app)
+                ctx.elasticsearch = Elasticsearch(**cfg, **self.extra_options)
                 # Retry connection on failure
                 for i in range(5):  # Retry up to 5 times
                     if ctx.elasticsearch.ping():
@@ -80,22 +81,18 @@ class FlaskElasticsearch:
             return getattr(ctx.elasticsearch, item)
 
     @staticmethod
-    def get_es_config() -> dict:
+    def get_es_config(app: Flask) -> dict:
         """Get Elasticsearch configuration from the current Flask application context."""
-        with current_app.app_context():
-            if current_app:
-                config_helper = ConfigHelper(current_app)
-                return config_helper.Elasticsearch
-            else:
-                logger.error('Attempted to access application configuration outside of application context.')
-                raise RuntimeError('Attempted to access application configuration outside of application context.')
+        config_helper = ConfigHelper(app)
+        return config_helper.Elasticsearch
 
     @staticmethod
     def build_es_options(cfg) -> dict:
         """
-        Build Elasticsearch connection options from the configuration.
+        Builds Elasticsearch connection options from the configuration.
 
         :param cfg: Elasticsearch's configuration from the Flask application context.
+        :return: Elasticsearch's connection options dictionary.
         """
         host = cfg.get('Host', 'localhost')
         port = int(cfg.get('Port', 9200))
@@ -115,7 +112,11 @@ class FlaskElasticsearch:
 
         return options
 
-    def _get_config(self) -> dict:
-        """Combines getting Elasticsearch configuration and building options."""
-        cfg = self.get_es_config()
+    def _get_config(self, app: Flask) -> dict:
+        """
+        Combines getting Elasticsearch configuration and building options.
+        :param app: Flask app instance.
+        :return: Elasticsearch's connection options dictionary.
+        """
+        cfg = self.get_es_config(app)
         return self.build_es_options(cfg)

application/views/user/user.py

@@ -54,13 +54,13 @@ def create_user():
 
 @user_api.route('/tests', methods=['POST'])
 def test_user():
-    # data = request.get_json()
-    # start = data.get('start')
-    # end = data.get('end')
-    # dsl = ElasticsearchUtil.dsl(start, end)
-    # index = current_app.config.Elasticsearch.Index
-    # res = ElasticsearchUtil.search(index, dsl)
-    # print(len(res))
+    data = request.get_json()
+    start = data.get('start')
+    end = data.get('end')
+    dsl = ElasticsearchUtil.dsl(start, end)
+    index = current_app.config.Elasticsearch.Index
+    res = ElasticsearchUtil.search(index, dsl)
+    print(len(res))
     data = {
         "user_name": "libai",
         "user_age": 18,

requirements.txt

@@ -13,10 +13,11 @@ dynaconf==3.2.3
 elastic-transport==8.10.0
 elasticsearch==7.13.0
 email-validator==2.1.0.post1
-Flask==3.0.0
+Flask==2.3.2
 flask-apispec==0.11.4
 Flask-Bcrypt==1.0.1
 Flask-Cors==4.0.0
+Flask-Elasticsearch==0.2.5
 flask-marshmallow==0.15.0
 Flask-Migrate==4.0.5
 Flask-RESTful==0.3.10