From f88cc9f095133715e9c63ab2e38f0c74aaec455f Mon Sep 17 00:00:00 2001
From: Xingheng Wang <xinghengwang@users.noreply.github.com>
Date: Wed, 4 Dec 2024 13:03:56 -0800
Subject: [PATCH] #188620653: wrap response to support trigger "close" (#73)

* wrap response to support trigger close

* bump version remove extra logging
---
 examples/flask/hello.py  | 16 ++++++++++++++--
 moesifwsgi/middleware.py | 26 +++++++++++++++++++++++---
 setup.py                 |  4 ++--
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/examples/flask/hello.py b/examples/flask/hello.py
index 11b6fa9..be16257 100644
--- a/examples/flask/hello.py
+++ b/examples/flask/hello.py
@@ -45,6 +45,7 @@ def mask_event(eventmodel):
     'GET_METADATA': get_metadata,
     'CAPTURE_OUTGOING_REQUESTS': False,
 }
+
 app.wsgi_app = MoesifMiddleware(app.wsgi_app, moesif_settings)
 
 @app.route('/')
@@ -53,7 +54,12 @@ def index():
 
 @app.route('/hello')
 def hello():
-    return 'Hello, world'
+      # Create a response object
+    response = Response('Hello, world')
+
+    # Attach a function to call_on_close
+    response.call_on_close(lambda: print("Response has been sent to the client."))
+    return response;
 
 @app.route('/user/<username>')
 def show_user_profile(username):
@@ -80,10 +86,16 @@ def show_post(post_id):
     }
 ]
 
+@app.route('/error', methods=['GET'])
+def error():
+    raise ValueError("This is a test error!")
+
 @app.route('/todo/api/v1.0/tasks', methods=['GET'])
 def get_tasks():
     json_body = json.dumps({'tasks': tasks})
-    return Response(response=json_body, status=201, mimetype='application/json')
+    response = Response(response=json_body, status=201, mimetype='application/json')
+    response.call_on_close(lambda: print("Response has been sent to the client."));
+    return response
 
 
 @app.route('/todo/api/v1.0/tasks/<int:task_id>', methods=['GET'])
diff --git a/moesifwsgi/middleware.py b/moesifwsgi/middleware.py
index 4b5b35a..ea953a5 100644
--- a/moesifwsgi/middleware.py
+++ b/moesifwsgi/middleware.py
@@ -104,7 +104,7 @@ def initialize_config(self):
             response_catcher = HttpResponseCatcher(self.DEBUG)
             self.api_client.http_call_back = response_catcher
         Configuration.BASE_URI = self.settings.get("BASE_URI") or self.settings.get("LOCAL_MOESIF_BASEURL", "https://api.moesif.net")
-        Configuration.version = "moesifwsgi-python/1.9.7"
+        Configuration.version = "moesifwsgi-python/1.9.8"
         if self.settings.get("CAPTURE_OUTGOING_REQUESTS", False):
             StartCapture().start_capture_outgoing(self.settings)
 
@@ -172,7 +172,8 @@ def _start_response(status, response_headers, *args):
           blocked_by = governed_response['blocked_by']
         else:
           # trigger next step in the process
-          response_chunks = event_info.finish_response(self.app(environ, _start_response))
+          original_response = self.app(environ, _start_response);
+          response_chunks = event_info.finish_response(original_response)
 
         # Add response chunks and response headers to the environ
         environ["moesif.response_body_chunks"] = response_chunks
@@ -186,10 +187,29 @@ def _start_response(status, response_headers, *args):
         self.add_user_and_metadata(event_info, environ, response_headers_mapping)
 
         try:
-            return response_chunks
+            return self.wrap_response(response_chunks, original_response)
         finally:
             self.process_and_add_event_if_required(event_info, environ, response_headers_mapping, blocked_by)
 
+    def wrap_response(self, response_chunks, original_response):
+        class WrappedResponse:
+            def __init__(self, chunks, original):
+                self.chunks = chunks
+                self.original = original
+
+            def __iter__(self):
+                return iter(self.chunks)
+
+            def close(self):
+                if hasattr(self.original, 'close'):
+                    self.original.close()
+
+            def __getattr__(self, name):
+                # Delegate attribute access to the original response
+                return getattr(self.original, name)
+
+        return WrappedResponse(response_chunks, original_response)
+
     def prepare_event_info(self, environ, start_response, request_time):
         event_info = DataHolder(
             self.settings.get("DISABLED_TRANSACTION_ID", False),
diff --git a/setup.py b/setup.py
index 21c2b5a..ab23178 100644
--- a/setup.py
+++ b/setup.py
@@ -28,7 +28,7 @@
     # Versions should comply with PEP440.  For a discussion on single-sourcing
     # the version across setup.py and the project code, see
     # https://packaging.python.org/en/latest/single_source_version.html
-    version='1.9.7',
+    version='1.9.8',
 
     description='Moesif Middleware for Python WSGI based platforms (Flask, Bottle & Others)',
     long_description=long_description,
@@ -39,7 +39,7 @@
 
     # Author details
     author='Moesif, Inc',
-    author_email='xing@moesif.com',
+    author_email='team@moesif.com',
 
     license='Apache Software License',