fix
This commit is contained in:
		@@ -1,3 +1,5 @@
 | 
				
			|||||||
 | 
					import asyncio
 | 
				
			||||||
 | 
					import collections
 | 
				
			||||||
import fastapi
 | 
					import fastapi
 | 
				
			||||||
import pydantic
 | 
					import pydantic
 | 
				
			||||||
import typing
 | 
					import typing
 | 
				
			||||||
@@ -5,7 +7,7 @@ import typing
 | 
				
			|||||||
from app.storage.mongo import tasks
 | 
					from app.storage.mongo import tasks
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
DEFAULT_RETRY_AFTER = 0.2
 | 
					locks = collections.defaultdict(asyncio.Lock)
 | 
				
			||||||
router = fastapi.APIRouter()
 | 
					router = fastapi.APIRouter()
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
@@ -19,13 +21,10 @@ class Response(pydantic.BaseModel):
 | 
				
			|||||||
    task: Task|None
 | 
					    task: Task|None
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
@router.get('/api/v1/take')
 | 
					@router.get('/api/v1/take', responses={404: {'description': 'Not found'}})
 | 
				
			||||||
async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
 | 
					async def execute(queue: typing.Annotated[str, fastapi.Header()]) -> Response:
 | 
				
			||||||
    try:
 | 
					    async with locks[queue]:
 | 
				
			||||||
        task = await tasks.take_task(queue)
 | 
					        task = await tasks.take_task(queue)
 | 
				
			||||||
        if not task:
 | 
					    if not task:
 | 
				
			||||||
            return Response(task=None)
 | 
					 | 
				
			||||||
        return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
 | 
					 | 
				
			||||||
    except Exception as e:
 | 
					 | 
				
			||||||
        print("GOT ERROR", e)
 | 
					 | 
				
			||||||
        return Response(task=None)
 | 
					        return Response(task=None)
 | 
				
			||||||
 | 
					    return Response(task=Task(id=str(task._id), attempt=task.attempts, payload=task.payload))
 | 
				
			||||||
 
 | 
				
			|||||||
		Reference in New Issue
	
	Block a user