Encapsulate common GCS operations in Storage class

pkalita-lbl · pkalita-lbl · commit eba99ee2c436 · 2025-07-18T10:26:41.000-07:00
diff --git a/nmdc_server/cli.py b/nmdc_server/cli.py
@@ -10,11 +10,12 @@
 import click
 import requests
 
-from nmdc_server import jobs, storage
+from nmdc_server import jobs
 from nmdc_server.config import settings
 from nmdc_server.database import SessionLocalIngest
 from nmdc_server.ingest import errors
 from nmdc_server.static_files import generate_submission_schema_files, initialize_static_directory
+from nmdc_server.storage import BucketName, storage
 
 
 def send_slack_message(text: str) -> bool:
@@ -350,17 +351,19 @@ def generate_static_files(remove_existing):
 @cli.command()
 def ensure_storage_buckets():
     """Ensure that the storage buckets exist."""
-    for bucket_name in storage.Bucket:
+    for bucket_name in BucketName:
         click.echo(f"Ensuring bucket '{bucket_name}' exists")
-        try:
-            storage.client.get_bucket(bucket_name)
+        bucket = storage.get_bucket(bucket_name)
+        if bucket.exists():
             click.echo(f"Bucket '{bucket_name}' already exists")
-        except Exception as e:
-            if settings.use_fake_gcs_server:
-                click.echo(f"Creating bucket '{bucket_name}'")
-                storage.client.create_bucket(bucket_name)
-            else:
-                raise RuntimeError(f"Failed to ensure bucket '{bucket_name}' exists: {e}")
+        elif settings.use_fake_gcs_server:
+            click.echo(f"Creating bucket '{bucket_name}'")
+            bucket.create()
+        else:
+            raise RuntimeError(
+                f"Failed to ensure bucket '{bucket_name}' exists. "
+                f"This bucket may need to be created manually in the cloud storage provider."
+            )
 
 
 if __name__ == "__main__":
diff --git a/nmdc_server/schemas.py b/nmdc_server/schemas.py
@@ -692,3 +692,34 @@ class VersionInfo(BaseModel):
     nmdc_schema: str = version("nmdc-schema")
     nmdc_submission_schema: str = version("nmdc-submission-schema")
     model_config = ConfigDict(frozen=False)
+
+
+class SignedUploadUrlRequest(BaseModel):
+    """Request to generate a signed upload URL for a file.
+
+    This model is used to generate a signed URL for uploading files to the object store.
+    """
+
+    file_name: str
+    file_size: int
+    content_type: str
+
+
+class SignedUrl(BaseModel):
+    """Response containing the signed URL and other metadata."""
+
+    url: str
+    object_name: str
+    expiration: datetime
+
+
+class UploadCompleteRequest(BaseModel):
+    """Request to mark an upload as complete.
+
+    This model is used to mark an upload as complete after the file has been uploaded
+    to the object store.
+    """
+
+    object_name: str
+    file_size: int
+    content_type: str
diff --git a/nmdc_server/storage.py b/nmdc_server/storage.py
@@ -1,29 +1,136 @@
+from datetime import datetime, timedelta, timezone
 from enum import StrEnum
+from functools import cached_property
 from typing import Any
 
-from google.auth.credentials import AnonymousCredentials
-from google.cloud import storage
+from google.cloud import storage as gcs, exceptions as gce
 
 from nmdc_server.config import settings
+from nmdc_server.schemas import SignedUrl
 
 
-class Bucket(StrEnum):
+class BucketName(StrEnum):
     """Enum for GCS bucket names"""
 
     SUBMISSION_IMAGES = "nmdc-submission-images"
 
 
-def _initialize_client() -> storage.Client:
-    client_args: dict[str, Any] = {
-        "project": settings.gcs_project_id,
-    }
+class Storage:
+    """A class to manage Google Cloud Storage interactions."""
 
-    if settings.use_fake_gcs_server:
-        # https://github.com/fsouza/fake-gcs-server/blob/cd43b03fcfb8149c6f57c1a92e19d1a07e291a3c/examples/python/python.py
-        client_args["credentials"] = AnonymousCredentials()
-        client_args["client_options"] = {"api_endpoint": "http://storage:4443"}
+    def __init__(self, project_id: str, use_fake_gcs_server: bool):
+        self.project_id = project_id
+        self.use_fake_gcs_server = use_fake_gcs_server
 
-    return storage.Client(**client_args)
+    @cached_property
+    def _client(self):
+        """Google Cloud Storage client."""
+        client_args: dict[str, Any] = {
+            "project": self.project_id,
+        }
 
+        if self.use_fake_gcs_server:
+            client_args["client_options"] = {"api_endpoint": "http://storage:4443"}
 
-client = _initialize_client()
+        return gcs.Client(**client_args)
+
+    def get_bucket(self, bucket_name: BucketName) -> gcs.Bucket:
+        """Get a GCS bucket by name.
+
+        :param bucket_name: The name of the bucket to retrieve.
+        """
+        return self._client.bucket(bucket_name)
+
+    def get_object(self, bucket_name: BucketName, object_name: str) -> gcs.Blob:
+        """Get an object from a GCS bucket.
+
+        :param bucket_name: The name of the bucket containing the object.
+        :param object_name: The name of the object to retrieve.
+        """
+        bucket = self.get_bucket(bucket_name)
+        return bucket.blob(object_name)
+
+    def delete_object(
+        self, bucket_name: BucketName, object_name: str, *, raise_if_not_found: bool = False
+    ) -> None:
+        """Delete an object from a GCS bucket.
+
+        :param bucket_name: The name of the bucket containing the object.
+        :param object_name: The name of the object to delete.
+        :param raise_if_not_found: If True, raise an exception if the object is not found. Default
+            is False.
+        """
+        bucket = self.get_bucket(bucket_name)
+        try:
+            bucket.delete_blob(object_name)
+        except gce.NotFound as e:
+            if raise_if_not_found:
+                raise e
+
+    def get_signed_upload_url(
+        self,
+        bucket_name: BucketName,
+        object_name: str,
+        *,
+        expiration: int = 15,
+        content_type: str | None = None,
+    ) -> SignedUrl:
+        """Get a signed URL for uploading to an object to a GCS bucket.
+
+        :param bucket_name: The name of the bucket that will contain the object.
+        :param object_name: The name of the object.
+        :param expiration: The expiration time for the signed URL in minutes. Default is 15 minutes.
+        :param content_type: The content type of the object being uploaded.
+        """
+        blob = self.get_object(bucket_name, object_name)
+        expiration_delta = timedelta(minutes=expiration)
+        expiration_time = datetime.now(timezone.utc) + expiration_delta
+
+        url = blob.generate_signed_url(
+            version="v4",
+            expiration=expiration_delta,
+            method="PUT",
+            content_type=content_type,
+        )
+
+        if self.use_fake_gcs_server:
+            # If using a fake GCS server, we need to adjust the URL to point localhost instead of
+            # the docker-compose service name.
+            url = url.replace("//storage:", "//localhost:")
+
+        return SignedUrl(url=url, expiration=expiration_time, object_name=blob.name)
+
+    def get_signed_download_url(
+        self, bucket_name: BucketName, object_name: str, *, expiration: int = 15
+    ) -> SignedUrl:
+        """Get a signed URL for downloading an object from a GCS bucket.
+
+        :param bucket_name: The name of the bucket containing the object.
+        :param object_name: The name of the object to download.
+        :param expiration: The expiration time for the signed URL in minutes. Default is 15 minutes.
+        """
+        expiration_delta = timedelta(minutes=expiration)
+        expiration_time = datetime.now(timezone.utc) + expiration_delta
+        blob = self.get_object(bucket_name, object_name)
+        url = blob.generate_signed_url(
+            version="v4",
+            expiration=expiration_delta,
+            method="GET",
+        )
+
+        if self.use_fake_gcs_server:
+            # If using a fake GCS server, we need to adjust the URL to point localhost instead of
+            # the docker-compose service name.
+            url = url.replace("//storage:", "//localhost:")
+
+        return SignedUrl(
+            url=url,
+            object_name=blob.name,
+            expiration=expiration_time,
+        )
+
+
+storage = Storage(
+    project_id=settings.gcs_project_id,
+    use_fake_gcs_server=settings.use_fake_gcs_server,
+)
