Coverage for paperless_ai/db.py: 56%

1from __future__ import annotations 

2 

3from contextlib import contextmanager 

4 

5from django.db import connections 

6 

7 

8@contextmanager 

9def db_connection_released(): 

10 """ 

11 Return any checked-out DB connections to the pool for the duration of the 

12 wrapped block. 

13 

14 The AI endpoints run inside a synchronous web request (``ai_suggestions``) 

15 or a streaming response (``chat``). Django keeps the request's database 

16 connection checked out for the entire request/response, so a blocking LLM 

17 call - which can take many seconds - pins a pooled connection the whole 

18 time. With connection pooling enabled, enough concurrent AI requests check 

19 out every slot and all other requests then fail with 

20 ``psycopg_pool.PoolTimeout`` (see issue #12976). 

21 

22 No Django ORM access happens during the LLM call, so we hand the connection 

23 back to the pool first; Django transparently re-checks-out a connection on 

24 the next ORM use after the block. 

25 """ 

26 connections.close_all() 

27 try: 

28 yield 

29 finally: 

30 connections.close_all()