nucliadb-utils 5.2.1.post1904__py3-none-any.whl → 5.2.1.post1912__py3-none-any.whl

nucliadb_utils/audit/stream.py

@@ -30,8 +30,9 @@ import nats
 from fastapi import Request
 from google.protobuf.timestamp_pb2 import Timestamp
 from opentelemetry.trace import format_trace_id, get_current_span
+from starlette.background import BackgroundTask
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
-from starlette.responses import Response, StreamingResponse
+from starlette.responses import Response
 from starlette.types import ASGIApp
 
 from nucliadb_protos.audit_pb2 import AuditField, AuditRequest, ChatContext, ClientType, RetrievedContext
@@ -98,11 +99,15 @@ class AuditMiddleware(BaseHTTPMiddleware):
 
         response = await call_next(request)
 
-        if isinstance(response, StreamingResponse):
-            response = self.wrap_streaming_response(response, context)
-        else:
-            self.enqueue_pending(context)
+        # This task will run when the response finishes streaming
+        # When dealing with streaming responses, AND if we depend on any state that only will be available once
+        # the request is fully finished, the response we have after the dispatch call_next is not enough, as
+        # there, no iteration of the streaming response has been done yet.
+        response.background = BackgroundTask(self.enqueue_pending, context)
 
+        # It is safe to reset the context here since the asyncio task for generating the streaming response is
+        # already running. If we want to spawn a different task during streaming and we want that task be able
+        # to read the context_var, we need to manually pass the context into that task.
         request_context_var.reset(token)
 
         return response
@@ -116,29 +121,6 @@ class AuditMiddleware(BaseHTTPMiddleware):
             if self.audit_utility is not None:
                 self.audit_utility.send(context.audit_request)
 
-    def wrap_streaming_response(
-        self, response: StreamingResponse, context: RequestContext
-    ) -> StreamingResponse:
-        """
-        When dealing with streaming responses, AND if we depend on any state that only will be available once
-        the request is fully finished, the response we have after the dispatch call_next is not enough, as
-        there, no iteration of the streaming response has been done yet.
-
-        This is why we need to rewrap to be able to to the auditing at the _real_ request end without losing
-        any audit bits.
-        """
-        original_body_iterator = response.body_iterator
-
-        async def custom_body_iterator():
-            try:
-                async for chunk in original_body_iterator:
-                    yield chunk
-            finally:
-                self.enqueue_pending(context)
-
-        response.body_iterator = custom_body_iterator()
-        return response
-
 
 KB_USAGE_STREAM_SUBJECT = "kb-usage.nuclia_db"
 

{nucliadb_utils-5.2.1.post1904.dist-info → nucliadb_utils-5.2.1.post1912.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb_utils
-Version: 5.2.1.post1904
+Version: 5.2.1.post1912
 Home-page: https://nuclia.com
 License: BSD
 Classifier: Development Status :: 4 - Beta
@@ -24,8 +24,8 @@ Requires-Dist: PyNaCl
 Requires-Dist: pyjwt>=2.4.0
 Requires-Dist: memorylru>=1.1.2
 Requires-Dist: mrflagly>=0.2.9
-Requires-Dist: nucliadb-protos>=5.2.1.post1904
-Requires-Dist: nucliadb-telemetry>=5.2.1.post1904
+Requires-Dist: nucliadb-protos>=5.2.1.post1912
+Requires-Dist: nucliadb-telemetry>=5.2.1.post1912
 Provides-Extra: cache
 Requires-Dist: redis>=4.3.4; extra == "cache"
 Requires-Dist: orjson>=3.6.7; extra == "cache"

{nucliadb_utils-5.2.1.post1904.dist-info → nucliadb_utils-5.2.1.post1912.dist-info}/RECORD

@@ -24,7 +24,7 @@ nucliadb_utils/aiopynecone/models.py,sha256=XkNIZx4bxdbVo9zYVn8IRp70q4DWUMWN79yb
 nucliadb_utils/audit/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb_utils/audit/audit.py,sha256=hH1L5XeRjTO896wMyINqiI5dkgGD_h8kbMsRs2HZ8Ts,3326
 nucliadb_utils/audit/basic.py,sha256=qKA7YAesuhRfW3iQXInI_GVm32az-WI6UPdS8QKQ_Fw,3974
-nucliadb_utils/audit/stream.py,sha256=vouDT75vCZwQEio0MRUm8JObK_25eknSotvVlaz3ppY,15162
+nucliadb_utils/audit/stream.py,sha256=j2C-jrm1lLnHlojrUn2nnf3Rk9LCItw7IQbwkfSfw9E,14828
 nucliadb_utils/cache/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb_utils/cache/exceptions.py,sha256=Zu-O_-0-yctOEgoDGI92gPzWfBMRrpiAyESA62ld6MA,975
 nucliadb_utils/cache/nats.py,sha256=-AjCfkFgKVdJUlGR0hT9JDSNkPVFg4S6w9eW-ZIcXPM,7037
@@ -59,8 +59,8 @@ nucliadb_utils/tests/indexing.py,sha256=YW2QhkhO9Q_8A4kKWJaWSvXvyQ_AiAwY1VylcfVQ
 nucliadb_utils/tests/local.py,sha256=7nuP8EFUAiA8ZH50R1iPV9EUXBySQxOanVm3Zht_e0g,1835
 nucliadb_utils/tests/nats.py,sha256=xqpww4jZjTKY9oPGlJdDJG67L3FIBQsa9qDHxILR8r8,7687
 nucliadb_utils/tests/s3.py,sha256=IdMxK_cNdSHLvO1u8BwsKFzD87Hk1MVPDZ57zx6h-rA,3656
-nucliadb_utils-5.2.1.post1904.dist-info/METADATA,sha256=gyjajvmdK3t-uM1BP0E9n13yuDqwNujNzEC1ea6_KvA,2055
-nucliadb_utils-5.2.1.post1904.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-nucliadb_utils-5.2.1.post1904.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
-nucliadb_utils-5.2.1.post1904.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb_utils-5.2.1.post1904.dist-info/RECORD,,
+nucliadb_utils-5.2.1.post1912.dist-info/METADATA,sha256=Gb_v9Cp3mabD74kyLbFswHqT0jDUFLo8AZdfokbZK5w,2055
+nucliadb_utils-5.2.1.post1912.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+nucliadb_utils-5.2.1.post1912.dist-info/top_level.txt,sha256=fE3vJtALTfgh7bcAWcNhcfXkNPp_eVVpbKK-2IYua3E,15
+nucliadb_utils-5.2.1.post1912.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb_utils-5.2.1.post1912.dist-info/RECORD,,