gcs_utils.py

import os
import logging
from typing import Optional
from pathlib import Path

from google.cloud import storage
from google.api_core.exceptions import GoogleAPIError
from google.api_core.retry import Retry
from google.oauth2 import service_account


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class GCSClient:

    def __init__(self, gcp_credentials_path: Optional[str], project: Optional[str] = None, retry: Optional[Retry] = None):
        if gcp_credentials_path:
            credentials = service_account.Credentials.from_service_account_file(
                gcp_credentials_path
            )
            self.client = storage.Client(credentials=credentials)
        else:
            self.client = storage.Client(project=project)
            
        self.retry = retry or Retry()

    def close(self):
        """Close the storage client."""
        self.client.close()

    def upload_file(
        self,
        bucket_name: str,
        source_file: str,
        destination_blob: str,
        overwrite: bool = True,
    ) -> None:
        try:
            if not os.path.exists(source_file):
                raise FileNotFoundError(f"Source file not found: {source_file}")

            bucket = self.client.bucket(bucket_name)
            blob = bucket.blob(destination_blob)

            if not overwrite and blob.exists():
                raise FileExistsError(f"Blob already exists: {destination_blob}")

            logger.info(f"Uploading {source_file} → gs://{bucket_name}/{destination_blob}")
            blob.upload_from_filename(source_file, retry=self.retry)
            logger.info("Upload successful.")

        except GoogleAPIError as e:
            logger.error(f"GCS Upload Error: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected Error: {e}")
            raise

    def download_file(
        self,
        bucket_name: str,
        source_blob: str,
        destination_file: str,
    ) -> None:
        try:
            bucket = self.client.bucket(bucket_name)
            blob = bucket.blob(source_blob)

            if not blob.exists():
                raise FileNotFoundError(f"Blob not found: {source_blob}")

            os.makedirs(os.path.dirname(destination_file), exist_ok=True)

            logger.info(f"Downloading gs://{bucket_name}/{source_blob} → {destination_file}")
            blob.download_to_filename(destination_file, retry=self.retry)
            logger.info("Download successful.")

        except GoogleAPIError as e:
            logger.error(f"GCS Download Error: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected Error: {e}")
            raise

    def download_folder(
        self,
        bucket_name: str,
        prefix: str,
        local_dir: str,
    ) -> None:
        try:
            bucket = self.client.bucket(bucket_name)
            blobs = bucket.list_blobs(prefix=prefix)

            downloaded_any = False

            for blob in blobs:
                if blob.name.endswith("/"):
                    continue

                relative_path = os.path.relpath(blob.name, prefix)
                local_path = Path(local_dir) / relative_path
                local_path.parent.mkdir(parents=True, exist_ok=True)

                logger.info(f"Downloading {blob.name} → {local_path}")
                blob.download_to_filename(str(local_path), retry=self.retry)
                downloaded_any = True

            if not downloaded_any:
                logger.warning(f"No files found with prefix: {prefix}")

            logger.info("Folder download complete.")

        except GoogleAPIError as e:
            logger.error(f"GCS Folder Download Error: {e}")
            raise
        except Exception as e:
            logger.error(f"Unexpected Error: {e}")
            raise


    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()