mcp_server/gcs_utils.py

import os
import sys
from google.cloud import storage

def _get_storage_client(credentials_path=None):
    """Initializes and returns a GCS storage client, allowing for explicit credential passing."""
    if credentials_path:
        print(f"Authenticating with service account key.", file=sys.stderr)
        return storage.Client.from_service_account_json(credentials_path)
    else:
        print("Attempting to authenticate using Application Default Credentials.", file=sys.stderr)
        return storage.Client()

def download_blob(bucket_name, source_blob_name, destination_file_name, credentials_path=None):
    """Downloads a blob from the bucket."""
    try:
        storage_client = _get_storage_client(credentials_path)
        bucket = storage_client.bucket(bucket_name)
        blob = bucket.blob(source_blob_name)
        
        os.makedirs(os.path.dirname(destination_file_name), exist_ok=True)
        blob.download_to_filename(destination_file_name)
        
        print(f"Downloaded storage object {source_blob_name} from bucket {bucket_name} to local file {destination_file_name}.", file=sys.stderr)
    except Exception as e:
        print(f"Failed to download {source_blob_name} from {bucket_name}: {e}", file=sys.stderr)
        raise

def download_directory(bucket_name, prefix, local_dir, credentials_path=None):
    """Downloads a directory from the bucket."""
    try:
        storage_client = _get_storage_client(credentials_path)
        bucket = storage_client.bucket(bucket_name)
        blobs = list(bucket.list_blobs(prefix=prefix))
        
        if not blobs:
            print(f"Warning: No files found in bucket '{bucket_name}' with prefix '{prefix}'.", file=sys.stderr)
            return

        if not os.path.exists(local_dir):
            os.makedirs(local_dir)

        for blob in blobs:
            if blob.name.endswith("/"):
                continue
            
            relative_path = os.path.relpath(blob.name, prefix)
            local_path = os.path.join(local_dir, relative_path)
            
            local_dir_path = os.path.dirname(local_path)
            if not os.path.exists(local_dir_path):
                os.makedirs(local_dir_path)
                
            blob.download_to_filename(local_path)
            print(f"Downloaded {blob.name} to {local_path}", file=sys.stderr)
            
    except Exception as e:
        print(f"Failed to download directory {prefix} from {bucket_name}: {e}", file=sys.stderr)
        raise