gcp.py

import os
from pathlib import Path
from datetime import timedelta
from google.cloud import storage
from huggingface_hub import hf_hub_download
from google.cloud.storage import transfer_manager


def download_credentials():
    os.makedirs('assets', exist_ok=True)

    # Download credentials file
    hf_hub_download(
        repo_id=os.environ.get('DATA'), repo_type='dataset', filename="credentials.json",
        token=os.environ.get('HUB_TOKEN'), local_dir="assets"
    )


def upload_folder(bucket_name: str, source_directory: str) -> None:
    # Filter so the list only includes files, not directories themselves.
    string_paths = [
        str(path.relative_to(source_directory)) for path in Path(source_directory).rglob("*") if path.is_file()
    ]

    # Start the upload.
    bucket = STORAGE_CLIENT.bucket(bucket_name)
    results = transfer_manager.upload_many_from_filenames(
        bucket, string_paths, source_directory=source_directory, max_workers=2
    )

    for name, result in zip(string_paths, results):
        if isinstance(result, Exception):
            print(f"Failed to upload {name} due to exception: {result}")
        else:
            print(f"Uploaded {name} to {bucket.name}.")


def get_link_file(bucket_name: str, client_name: str, type_media: str, media_name: str):
    bucket = STORAGE_CLIENT.bucket(bucket_name)
    blobs = bucket.list_blobs(prefix=f'{client_name}/media/{type_media}/{media_name}')
    blob = next(blobs)
    signed_url = blob.generate_signed_url(expiration=timedelta(minutes=15))
    return signed_url


download_credentials()
STORAGE_CLIENT = storage.Client.from_service_account_json(os.getenv('GOOGLE_APPLICATION_CREDENTIALS'))