From db433cc2b2d925fe0c29481a20eb776621b01530 Mon Sep 17 00:00:00 2001 From: Ivan Date: Tue, 19 Nov 2024 12:34:18 -0800 Subject: [PATCH] Uploading kernel images via CLI --- frontend/src/gen/api.ts | 50 +++++++++++++++++++++++++++ frontend/src/hooks/api.tsx | 34 ++++++++++++++++++ store/app/crud/artifacts.py | 31 ----------------- store/app/routers/artifacts.py | 63 ++++++++++++++++++++++++++++++++++ store/settings/environment.py | 2 +- 5 files changed, 148 insertions(+), 32 deletions(-) diff --git a/frontend/src/gen/api.ts b/frontend/src/gen/api.ts index 6df43d17..6c5b9fc9 100644 --- a/frontend/src/gen/api.ts +++ b/frontend/src/gen/api.ts @@ -324,6 +324,23 @@ export interface paths { patch?: never; trace?: never; }; + "/artifacts/presigned/{listing_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Get Presigned Url */ + post: operations["get_presigned_url_artifacts_presigned__listing_id__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/keys/new": { parameters: { query?: never; @@ -2683,6 +2700,39 @@ export interface operations { }; }; }; + get_presigned_url_artifacts_presigned__listing_id__post: { + parameters: { + query: { + filename: string; + }; + header?: never; + path: { + listing_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": Record; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; new_key_keys_new_post: { parameters: { query?: never; diff --git a/frontend/src/hooks/api.tsx b/frontend/src/hooks/api.tsx index b92d0e5d..ae93302f 100644 --- a/frontend/src/hooks/api.tsx +++ b/frontend/src/hooks/api.tsx @@ -25,4 +25,38 @@ export default class api { }, }); } + + public async uploadKernel(file: File, listing_id: string) { + const { data } = await this.client.POST( + "/artifacts/presigned/{listing_id}", + { + params: { + path: { + listing_id, + }, + query: { + filename: file.name, + }, + }, + }, + ); + + if (!data?.upload_url) { + throw new Error("Failed to get upload URL"); + } + + const uploadResponse = await fetch(data.upload_url, { + method: "PUT", + body: file, + headers: { + "Content-Type": "application/octet-stream", + }, + }); + + if (!uploadResponse.ok) { + throw new Error(`Upload failed: ${uploadResponse.statusText}`); + } + + return data.artifact_id; + } } diff --git a/store/app/crud/artifacts.py b/store/app/crud/artifacts.py index 0ca2a52d..ad5e469b 100644 --- a/store/app/crud/artifacts.py +++ b/store/app/crud/artifacts.py @@ -260,35 +260,6 @@ async def _upload_and_store( ) return artifact - async def _upload_kernel( - self, - name: str, - file: UploadFile, - listing: Listing, - description: str | None = None, - ) -> Artifact: - artifact = Artifact.create( - user_id=listing.user_id, - listing_id=listing.id, - name=name, - artifact_type="kernel", - description=description, - ) - - file_data = io.BytesIO(await file.read()) - s3_filename = get_artifact_name(artifact=artifact) - - await asyncio.gather( - self._upload_to_s3( - data=file_data, - name=name, - filename=s3_filename, - content_type="application/octet-stream", - ), - self._add_item(artifact), - ) - return artifact - async def upload_artifact( self, name: str, @@ -300,8 +271,6 @@ async def upload_artifact( match artifact_type: case "image": return await self._upload_image(name, file, listing, description) - case "kernel": - return await self._upload_kernel(name, file, listing, description) case "stl" | "obj" | "ply" | "dae": return await self._upload_mesh(name, file, listing, artifact_type, description) case "urdf" | "mjcf": diff --git a/store/app/routers/artifacts.py b/store/app/routers/artifacts.py index bdbf1593..0aa61e73 100644 --- a/store/app/routers/artifacts.py +++ b/store/app/routers/artifacts.py @@ -16,11 +16,13 @@ Artifact, ArtifactSize, ArtifactType, + KernelArtifactType, Listing, User, can_write_artifact, can_write_listing, check_content_type, + get_artifact_name, get_artifact_type, get_artifact_urls, ) @@ -411,3 +413,64 @@ async def set_main_image( await crud.set_main_image(listing_id, artifact_id) return True + + +class PresignedUrlResponse(BaseModel): + upload_url: str + artifact_id: str + + +@router.post("/presigned/{listing_id}", response_model=PresignedUrlResponse) +async def get_presigned_url( + listing_id: str, + filename: str, + user: Annotated[User, Depends(get_session_user_with_write_permission)], + crud: Annotated[Crud, Depends(Crud.get)], +) -> PresignedUrlResponse: + if not filename: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Filename was not provided", + ) + + name, extension = os.path.splitext(filename) + if extension.lower() != ".img": + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="Only .img files are supported for kernel uploads" + ) + + artifact_type: KernelArtifactType = "kernel" + + listing = await crud.get_listing(listing_id) + if listing is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Could not find listing") + if not await can_write_listing(user, listing): + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="No permission to upload") + + artifact = Artifact.create( + user_id=listing.user_id, + listing_id=listing.id, + name=name, + artifact_type=artifact_type, + ) + await crud._add_item(artifact) + + try: + s3_filename = get_artifact_name(artifact=artifact) + + presigned_url = await crud.s3.meta.client.generate_presigned_url( + ClientMethod="put_object", + Params={ + "Bucket": settings.s3.bucket, + "Key": f"{settings.s3.prefix}{s3_filename}", + "ContentType": "application/octet-stream", + "ContentDisposition": f'attachment; filename="{name}"', + }, + ExpiresIn=3600, + ) + + return PresignedUrlResponse(upload_url=presigned_url, artifact_id=artifact.id) + + except Exception as e: + await crud._delete_item(artifact) + raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) diff --git a/store/settings/environment.py b/store/settings/environment.py index cbdb0b5f..17f3fe80 100644 --- a/store/settings/environment.py +++ b/store/settings/environment.py @@ -42,7 +42,7 @@ class ArtifactSettings: large_image_size: tuple[int, int] = field(default=(1536, 1536)) small_image_size: tuple[int, int] = field(default=(256, 256)) min_bytes: int = field(default=16) - max_bytes: int = field(default=2**30) + max_bytes: int = field(default=1536 * 1536 * 25) quality: int = field(default=80) max_concurrent_file_uploads: int = field(default=3)