Merge pull request #11 from lsst-sqre/tickets/DM-35350

[DM-35350] Use per-thread Google storage clients

src/crawlspace/dependencies/gcs.py

@@ -2,12 +2,14 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from contextvars import ContextVar
 
 from google.cloud import storage
 
 from ..config import config
 
+_GCS_CLIENT: ContextVar[storage.Client] = ContextVar("_GCS_CLIENT")
+
 __all__ = [
     "GCSClientDependency",
     "gcs_client_dependency",
@@ -17,14 +19,9 @@
 class GCSClientDependency:
     """Provides a `google.cloud.storage.Client` as a dependency."""
 
-    def __init__(self) -> None:
-        self.gcs: Optional[storage.Client] = None
-
     async def __call__(self) -> storage.client.Client:
         """Return the cached `google.cloud.storage.Client`."""
-        if not self.gcs:
-            self.gcs = storage.Client(project=config.gcs_project)
-        return self.gcs
+        return storage.Client(project=config.gcs_project)
 
 
 gcs_client_dependency = GCSClientDependency()

src/crawlspace/handlers/external.py

@@ -3,7 +3,7 @@
 from typing import List
 
 from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
-from fastapi.responses import RedirectResponse, StreamingResponse
+from fastapi.responses import RedirectResponse
 from google.cloud import storage
 from safir.dependencies.logger import logger_dependency
 from structlog.stdlib import BoundLogger
@@ -40,13 +40,15 @@ def get_file(
     etags: List[str] = Depends(etag_validation_dependency),
     logger: BoundLogger = Depends(logger_dependency),
 ) -> Response:
+    logger.debug("File request", path=path)
     if path == "":
         path = "index.html"
 
     file_service = FileService(gcs)
     try:
         crawlspace_file = file_service.get_file(path)
     except GCSFileNotFoundError:
+        logger.debug("File not found", path=path)
         raise HTTPException(status_code=404, detail="File not found")
     except Exception as e:
         logger.exception(f"Failed to retrieve {path}", error=str(e))
@@ -55,15 +57,19 @@ def get_file(
         )
 
     if crawlspace_file.blob.etag in etags:
+        logger.debug("File unchanged", path=path)
         return Response(
             status_code=304,
             content="",
             headers=crawlspace_file.headers,
             media_type=crawlspace_file.media_type,
         )
     else:
-        return StreamingResponse(
-            crawlspace_file.stream(),
+        logger.debug("Returning file", path=path)
+        data = crawlspace_file.download_as_bytes()
+        return Response(
+            status_code=200,
+            content=data,
             media_type=crawlspace_file.media_type,
             headers=crawlspace_file.headers,
         )
@@ -82,20 +88,23 @@ def head_file(
     gcs: storage.Client = Depends(gcs_client_dependency),
     logger: BoundLogger = Depends(logger_dependency),
 ) -> Response:
+    logger.debug("Head request", path=path)
     if path == "":
         path = "index.html"
 
     file_service = FileService(gcs)
     try:
         crawlspace_file = file_service.get_file(path)
     except GCSFileNotFoundError:
+        logger.debug("File not found for head request", path=path)
         raise HTTPException(status_code=404, detail="File not found")
     except Exception as e:
         logger.exception(f"Failed to retrieve {path}", error=str(e))
         raise HTTPException(
             status_code=500, detail="Failed to retrieve file from GCS"
         )
 
+    logger.debug("Returning file metadata", path=path)
     return Response(
         status_code=200,
         content="",

src/crawlspace/services/file.py

@@ -52,6 +52,10 @@ def from_blob(cls, path: str, blob: storage.Blob) -> CrawlspaceFile:
             media_type = guessed_type if guessed_type else "text/plain"
         return cls(blob=blob, headers=headers, media_type=media_type)
 
+    def download_as_bytes(self) -> bytes:
+        """Download the content from GCS."""
+        return self.blob.download_as_bytes()
+
     def stream(self) -> Iterator[bytes]:
         """Stream the content from GCS."""
         with self.blob.open("rb") as content:

tests/support/gcs.py

@@ -30,6 +30,9 @@ def __init__(self, name: str) -> None:
             self.updated = datetime.fromtimestamp(mtime, tz=timezone.utc)
             self.etag = str(self._path.stat().st_ino)
 
+    def download_as_bytes(self) -> bytes:
+        return self._path.read_bytes()
+
     def exists(self) -> bool:
         return self._exists