[Python] How to generate a tar.gz stream to be returned as a Streaming response in...

I'm making an FastAPI/Starlette server which requests to another server (well, S3) large files. These chunks feed a tarfile.TarFile object to produce a .tar.gz stream. This stream should be sent on the fly to a StreamingResponse.

S3 server --files chunks--> My server --tar.gz chunks--> User


There is a big issue with the tarfile.TarFile implementation as:

it writes chunks in file-like objects (with a write method),


This is not easy to be interfaced with the stralette.StreamingResponse which expects a chunk generator.

My idea was to rewrite the StreamingResponse.stream_response method (cf here). Something like this:

"""A script to test the tar.gz streaming."""

import os
import tarfile
from pathlib import Path
from typing import Mapping

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.testclient import TestClient
from starlette.background import BackgroundTask
from starlette.types import Send

TAR_FILE_PATH = Path("archive.tar.gz")

CHUNK_SIZE = 1024

app = FastAPI()


class FileStreamingResponse(StreamingResponse):

def __init__(
self,
files_to_tar: list[Path],
status_code: int = 200,
headers: Mapping[str, str] | None = None,
media_type: str | None = None,
background: BackgroundTask | None = None,
) -> None:
self.files_to_tar = files_to_tar
self.status_code = status_code
self.media_type = self.media_type if media_type is None else media_type
self.background = background
self.init_headers(headers)

async def stream_response(self, send: Send) -> None:
await send(
{
"type": "http.response.start",
"status": self.status_code,
"headers": self.raw_headers,
}
)

class DumpWriter:
async def write(buffer):
print(f"Really sending {len(buffer)} bytes")
await send(
{"type": "http.response.body", "body": buffer, "more_body": True}
)

async with tarfile.open(
mode="w|gz", fileobj=DumpWriter(), bufsize=CHUNK_SIZE
) as file:
for input_file in self.files_to_tar:
await file.add(input_file.open("rb"))

await send({"type": "http.response.body", "body": b"", "more_body": False})


FILES_TO_TAR = Path("src").iterdir()


@app.get("/")
def send_tar() -> StreamingResponse:
"""Send a tar file."""

return FileStreamingResponse(
files_to_tar=FILES_TO_TAR,
media_type="application/tar+gzip",
headers={"Content-Disposition": f'attachment; filename="{TAR_FILE_PATH.name}"'},
)


#
# TESTS
#

TEST_TAR_FILE_PATH = Path("archive.tar.gz")
client = TestClient(app)


def test_main():
if TEST_TAR_FILE_PATH.exists():
os.remove(TEST_TAR_FILE_PATH)

response = client.get("/")
response.raise_for_status()

with TEST_TAR_FILE_PATH.open("wb") as file:
for chunk in response.iter_bytes(CHUNK_SIZE):
file.write(chunk)

with tarfile.open(TEST_TAR_FILE_PATH, "r:gz") as tar_file:
files = [tarinfo.name for tarinfo in tar_file.getmembers()]

src_files = [file.name for file in FILES_TO_TAR]
assert set(files) == set(src_files)


but the issue is that...

it writes them synchronuosly


Ok, then to write it correctly, I need to run the DumpWriter.write method synchronously? I tried to do

class DumpWriter:
def write(buffer):
print(f"Really sending {len(buffer)} bytes")
asyncio.run_coroutine_threadsafe(
send(
{
"type": "http.response.body",
"body": buffer,
"more_body": True,
}
),
asyncio.get_running_loop(),
)


but the data chunks were not transmitted.

I also thought to rewrite the TarFile class to be async, but this is a nightmare!

How can I achieve that?

Continue reading...