How to close a socket inside an fastapi StreamingResponse generator?

Question

The code below doesn't reach out of the with. Is there any way I can close the socket when the client disconnects?

import fastapi
import uvicorn
import socket

def gen_video(port):
    with socket.socket() as s:
        s.connect(("127.0.0.1", port))
        while True:
            yield s.recv(1024)


app = fastapi.FastAPI()

@app.get("/stream")
async def stream():
    return fastapi.responses.StreamingResponse(gen_video(1234), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)

Related answer can be found [here](https://stackoverflow.com/a/73283272/17865804) as well. — Chris, Nov 08 '22 at 16:26
@Chris Doing the BackgroundTask method like you mentioned in the first question can resolve my problem. I will write an answer here using it. Thank you — jvx8ss, Nov 08 '22 at 16:41

score 0 · Accepted Answer · answered Nov 08 '22 at 16:44

An approach using fastapi BackgroundTasks to close the socket

import fastapi
import uvicorn
import socket

def read_socket(s):
    while True:
        yield s.recv(1024)

app = fastapi.FastAPI()

@app.get("/stream")
async def stream(background_tasks: fastapi.background.BackgroundTasks):
    s = socket.socket()
    s.connect(("127.0.0.1", 1234))
    background_tasks.add_task(lambda f: f.close(), s)
    return fastapi.responses.StreamingResponse(read_socket(s), media_type="video/ogg")

if __name__ == "__main__":
    uvicorn.run(app, port=8000)

How to close a socket inside an fastapi StreamingResponse generator?

1 Answers1