Skip to content

Commit

Permalink
fix(api): Drop StreamingResponse after issues on latest FastAPI
Browse files Browse the repository at this point in the history
  • Loading branch information
fschuch committed Jun 28, 2024
1 parent 874631a commit b806865
Show file tree
Hide file tree
Showing 5 changed files with 91 additions and 9 deletions.
2 changes: 1 addition & 1 deletion jobbergate-api/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
This file keeps track of all notable changes to jobbergate-api

## Unreleased

- Fixed issue when retrieving large files on get routes after upgrading to FastAPI 0.111

## 5.2.0a4 -- 2024-06-27
- Change pydantic.BaseSettings config to use `extra=ignore`
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from fastapi import APIRouter, BackgroundTasks, Body, Depends, File, HTTPException, Path, Query
from fastapi import Response as FastAPIResponse
from fastapi import UploadFile, status
from fastapi.responses import StreamingResponse
from fastapi_pagination import Page
from loguru import logger
from sqlalchemy.exc import IntegrityError
Expand Down Expand Up @@ -211,8 +210,8 @@ async def job_script_template_get_file(
logger.debug(f"Getting template file {file_name=} from job script template {typed_id_or_identifier=}")
job_script_template = await secure_services.crud.template.get(typed_id_or_identifier)
job_script_template_file = await secure_services.file.template.get(job_script_template.id, file_name)
return StreamingResponse(
content=await secure_services.file.template.stream_file_content(job_script_template_file),
return FastAPIResponse(
content=await secure_services.file.template.get_file_content(job_script_template_file),
media_type="text/plain",
headers={"filename": job_script_template_file.filename},
)
Expand Down Expand Up @@ -310,8 +309,8 @@ async def job_script_workflow_get_file(
logger.debug(f"Getting workflow file from job script template {typed_id_or_identifier=}")
job_script_template = await secure_services.crud.template.get(typed_id_or_identifier)
workflow_file = await secure_services.file.workflow.get(job_script_template.id, WORKFLOW_FILE_NAME)
return StreamingResponse(
content=await secure_services.file.workflow.stream_file_content(workflow_file),
return FastAPIResponse(
content=await secure_services.file.workflow.get_file_content(workflow_file),
media_type="text/plain",
headers={"filename": WORKFLOW_FILE_NAME},
)
Expand Down
5 changes: 2 additions & 3 deletions jobbergate-api/jobbergate_api/apps/job_scripts/routers.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
from fastapi import APIRouter, BackgroundTasks, Depends, File, HTTPException, Path, Query
from fastapi import Response as FastAPIResponse
from fastapi import UploadFile, status
from fastapi.responses import StreamingResponse
from fastapi_pagination import Page
from loguru import logger

Expand Down Expand Up @@ -288,8 +287,8 @@ async def job_script_get_file(
See https://fastapi.tiangolo.com/advanced/custom-response/#streamingresponse
"""
job_script_file = await secure_services.file.job_script.get(id, file_name)
return StreamingResponse(
content=await secure_services.file.job_script.stream_file_content(job_script_file),
return FastAPIResponse(
content=await secure_services.file.job_script.get_file_content(job_script_file),
media_type="text/plain",
headers={"filename": job_script_file.filename},
)
Expand Down
53 changes: 53 additions & 0 deletions jobbergate-api/tests/apps/job_script_templates/test_routers.py
Original file line number Diff line number Diff line change
Expand Up @@ -865,6 +865,36 @@ async def test_get__success(
assert response.status_code == status.HTTP_200_OK, f"Get failed: {response.text}"
assert response.content.decode() == "dummy file data"

async def test_get__large_file_success(
self,
client,
tester_email,
inject_security_header,
job_template_data,
synth_services,
):
"""
Ensure that large files can be retrieved with no problem.
This was created after strange issues were identified running on FastAPI 0.111.
"""
large_string = "print(1)\n" * 5000

parent_id = job_template_data.id
await synth_services.file.template.upsert(
parent_id=parent_id,
filename="test_template.py.j2",
upload_content=large_string,
file_type="ENTRYPOINT",
)

inject_security_header(tester_email, Permissions.JOB_TEMPLATES_READ)
response = await client.get(
f"jobbergate/job-script-templates/{job_template_data.id}/upload/template/test_template.py.j2"
)

assert response.status_code == status.HTTP_200_OK, f"Get failed: {response.text}"
assert response.content.decode() == large_string

@pytest.mark.parametrize(
"is_owner, permissions",
[
Expand Down Expand Up @@ -1171,6 +1201,29 @@ async def test_get__success(
assert response.status_code == status.HTTP_200_OK, f"Get failed: {response.text}"
assert response.content.decode() == "import this"

async def test_get__large_file_success(
self, client, tester_email, inject_security_header, job_template_data, synth_services
):
"""
Ensure that large files can be retrieved with no problem.
This was created after strange issues were identified running on FastAPI 0.111.
"""
large_string = "print(1)\n" * 5000

parent_id = job_template_data.id
await synth_services.file.workflow.upsert(
parent_id=parent_id,
filename=WORKFLOW_FILE_NAME,
upload_content=large_string,
runtime_config=dict(foo="bar"),
)

inject_security_header(tester_email, Permissions.JOB_TEMPLATES_READ)
response = await client.get(f"jobbergate/job-script-templates/{parent_id}/upload/workflow")

assert response.status_code == status.HTTP_200_OK, f"Get failed: {response.text}"
assert response.content.decode() == large_string

@pytest.mark.parametrize(
"is_owner, permissions",
[
Expand Down
31 changes: 31 additions & 0 deletions jobbergate-api/tests/apps/job_scripts/test_routers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1056,6 +1056,37 @@ async def test_get__success(
assert response.status_code == status.HTTP_200_OK
assert response.content.decode() == job_script_data_as_string

async def test_get__large_file_success(
self,
client,
tester_email,
inject_security_header,
job_script_data,
synth_services,
):
"""
Ensure that large files can be retrieved with no problem.
This was created after strange issues were identified running on FastAPI 0.111.
"""
large_string = "print(1)\n" * 5000

id = job_script_data.id
file_type = "ENTRYPOINT"
job_script_filename = "entrypoint.sh"

await synth_services.file.job_script.upsert(
parent_id=id,
filename=job_script_filename,
upload_content=large_string,
file_type=file_type,
)

inject_security_header(tester_email, Permissions.JOB_SCRIPTS_READ)
response = await client.get(f"jobbergate/job-scripts/{id}/upload/{job_script_filename}")

assert response.status_code == status.HTTP_200_OK, f"Get failed: {response.text}"
assert response.content.decode() == large_string

@pytest.mark.parametrize(
"is_owner, permissions",
[
Expand Down

0 comments on commit b806865

Please sign in to comment.