Add presigned URL support for direct S3 downloads (#48)
This commit is contained in:
Mondo Diaz
@@ -1,9 +1,9 @@
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Form, Request, Query, Header, Response
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.responses import StreamingResponse, RedirectResponse
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy import or_, func
|
||||
from typing import List, Optional
|
||||
from typing import List, Optional, Literal
|
||||
import math
|
||||
import re
|
||||
import io
|
||||
@@ -29,8 +29,10 @@ from .schemas import (
|
||||
ResumableUploadCompleteResponse,
|
||||
ResumableUploadStatusResponse,
|
||||
GlobalSearchResponse, SearchResultProject, SearchResultPackage, SearchResultArtifact,
|
||||
PresignedUrlResponse,
|
||||
)
|
||||
from .metadata import extract_metadata
|
||||
from .config import get_settings
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@@ -844,27 +846,13 @@ def get_upload_status(
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
# Download artifact with range request support
|
||||
@router.get("/api/v1/project/{project_name}/{package_name}/+/{ref}")
|
||||
def download_artifact(
|
||||
project_name: str,
|
||||
package_name: str,
|
||||
# Helper function to resolve artifact reference
|
||||
def _resolve_artifact_ref(
|
||||
ref: str,
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
storage: S3Storage = Depends(get_storage),
|
||||
range: Optional[str] = Header(None),
|
||||
):
|
||||
# Get project and package
|
||||
project = db.query(Project).filter(Project.name == project_name).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
package = db.query(Package).filter(Package.project_id == project.id, Package.name == package_name).first()
|
||||
if not package:
|
||||
raise HTTPException(status_code=404, detail="Package not found")
|
||||
|
||||
# Resolve reference to artifact
|
||||
package: Package,
|
||||
db: Session,
|
||||
) -> Optional[Artifact]:
|
||||
"""Resolve a reference (tag name, artifact:hash, tag:name) to an artifact"""
|
||||
artifact = None
|
||||
|
||||
# Check for explicit prefixes
|
||||
@@ -885,11 +873,76 @@ def download_artifact(
|
||||
# Try as direct artifact ID
|
||||
artifact = db.query(Artifact).filter(Artifact.id == ref).first()
|
||||
|
||||
return artifact
|
||||
|
||||
|
||||
# Download artifact with range request support and download modes
|
||||
@router.get("/api/v1/project/{project_name}/{package_name}/+/{ref}")
|
||||
def download_artifact(
|
||||
project_name: str,
|
||||
package_name: str,
|
||||
ref: str,
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
storage: S3Storage = Depends(get_storage),
|
||||
range: Optional[str] = Header(None),
|
||||
mode: Optional[Literal["proxy", "redirect", "presigned"]] = Query(
|
||||
default=None,
|
||||
description="Download mode: proxy (stream through backend), redirect (302 to presigned URL), presigned (return JSON with URL)"
|
||||
),
|
||||
):
|
||||
settings = get_settings()
|
||||
|
||||
# Get project and package
|
||||
project = db.query(Project).filter(Project.name == project_name).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
package = db.query(Package).filter(Package.project_id == project.id, Package.name == package_name).first()
|
||||
if not package:
|
||||
raise HTTPException(status_code=404, detail="Package not found")
|
||||
|
||||
# Resolve reference to artifact
|
||||
artifact = _resolve_artifact_ref(ref, package, db)
|
||||
if not artifact:
|
||||
raise HTTPException(status_code=404, detail="Artifact not found")
|
||||
|
||||
filename = artifact.original_name or f"{artifact.id}"
|
||||
|
||||
# Determine download mode (query param overrides server default)
|
||||
download_mode = mode or settings.download_mode
|
||||
|
||||
# Handle presigned mode - return JSON with presigned URL
|
||||
if download_mode == "presigned":
|
||||
presigned_url = storage.generate_presigned_url(
|
||||
artifact.s3_key,
|
||||
response_content_type=artifact.content_type,
|
||||
response_content_disposition=f'attachment; filename="{filename}"',
|
||||
)
|
||||
expires_at = datetime.now(timezone.utc) + timedelta(seconds=settings.presigned_url_expiry)
|
||||
|
||||
return PresignedUrlResponse(
|
||||
url=presigned_url,
|
||||
expires_at=expires_at,
|
||||
method="GET",
|
||||
artifact_id=artifact.id,
|
||||
size=artifact.size,
|
||||
content_type=artifact.content_type,
|
||||
original_name=artifact.original_name,
|
||||
checksum_sha256=artifact.id,
|
||||
checksum_md5=artifact.checksum_md5,
|
||||
)
|
||||
|
||||
# Handle redirect mode - return 302 redirect to presigned URL
|
||||
if download_mode == "redirect":
|
||||
presigned_url = storage.generate_presigned_url(
|
||||
artifact.s3_key,
|
||||
response_content_type=artifact.content_type,
|
||||
response_content_disposition=f'attachment; filename="{filename}"',
|
||||
)
|
||||
return RedirectResponse(url=presigned_url, status_code=302)
|
||||
|
||||
# Proxy mode (default fallback) - stream through backend
|
||||
# Handle range requests
|
||||
if range:
|
||||
stream, content_length, content_range = storage.get_stream(artifact.s3_key, range)
|
||||
@@ -923,6 +976,63 @@ def download_artifact(
|
||||
)
|
||||
|
||||
|
||||
# Get presigned URL endpoint (explicit endpoint for getting URL without redirect)
|
||||
@router.get("/api/v1/project/{project_name}/{package_name}/+/{ref}/url", response_model=PresignedUrlResponse)
|
||||
def get_artifact_url(
|
||||
project_name: str,
|
||||
package_name: str,
|
||||
ref: str,
|
||||
db: Session = Depends(get_db),
|
||||
storage: S3Storage = Depends(get_storage),
|
||||
expiry: Optional[int] = Query(
|
||||
default=None,
|
||||
description="Custom expiry time in seconds (defaults to server setting)"
|
||||
),
|
||||
):
|
||||
"""
|
||||
Get a presigned URL for direct S3 download.
|
||||
This endpoint always returns a presigned URL regardless of server download mode.
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
# Get project and package
|
||||
project = db.query(Project).filter(Project.name == project_name).first()
|
||||
if not project:
|
||||
raise HTTPException(status_code=404, detail="Project not found")
|
||||
|
||||
package = db.query(Package).filter(Package.project_id == project.id, Package.name == package_name).first()
|
||||
if not package:
|
||||
raise HTTPException(status_code=404, detail="Package not found")
|
||||
|
||||
# Resolve reference to artifact
|
||||
artifact = _resolve_artifact_ref(ref, package, db)
|
||||
if not artifact:
|
||||
raise HTTPException(status_code=404, detail="Artifact not found")
|
||||
|
||||
filename = artifact.original_name or f"{artifact.id}"
|
||||
url_expiry = expiry or settings.presigned_url_expiry
|
||||
|
||||
presigned_url = storage.generate_presigned_url(
|
||||
artifact.s3_key,
|
||||
expiry=url_expiry,
|
||||
response_content_type=artifact.content_type,
|
||||
response_content_disposition=f'attachment; filename="{filename}"',
|
||||
)
|
||||
expires_at = datetime.now(timezone.utc) + timedelta(seconds=url_expiry)
|
||||
|
||||
return PresignedUrlResponse(
|
||||
url=presigned_url,
|
||||
expires_at=expires_at,
|
||||
method="GET",
|
||||
artifact_id=artifact.id,
|
||||
size=artifact.size,
|
||||
content_type=artifact.content_type,
|
||||
original_name=artifact.original_name,
|
||||
checksum_sha256=artifact.id,
|
||||
checksum_md5=artifact.checksum_md5,
|
||||
)
|
||||
|
||||
|
||||
# HEAD request for download (to check file info without downloading)
|
||||
@router.head("/api/v1/project/{project_name}/{package_name}/+/{ref}")
|
||||
def head_artifact(
|
||||
@@ -941,23 +1051,8 @@ def head_artifact(
|
||||
if not package:
|
||||
raise HTTPException(status_code=404, detail="Package not found")
|
||||
|
||||
# Resolve reference to artifact (same logic as download)
|
||||
artifact = None
|
||||
if ref.startswith("artifact:"):
|
||||
artifact_id = ref[9:]
|
||||
artifact = db.query(Artifact).filter(Artifact.id == artifact_id).first()
|
||||
elif ref.startswith("tag:") or ref.startswith("version:"):
|
||||
tag_name = ref.split(":", 1)[1]
|
||||
tag = db.query(Tag).filter(Tag.package_id == package.id, Tag.name == tag_name).first()
|
||||
if tag:
|
||||
artifact = db.query(Artifact).filter(Artifact.id == tag.artifact_id).first()
|
||||
else:
|
||||
tag = db.query(Tag).filter(Tag.package_id == package.id, Tag.name == ref).first()
|
||||
if tag:
|
||||
artifact = db.query(Artifact).filter(Artifact.id == tag.artifact_id).first()
|
||||
else:
|
||||
artifact = db.query(Artifact).filter(Artifact.id == ref).first()
|
||||
|
||||
# Resolve reference to artifact
|
||||
artifact = _resolve_artifact_ref(ref, package, db)
|
||||
if not artifact:
|
||||
raise HTTPException(status_code=404, detail="Artifact not found")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user