wehub-resource-sync
b7f52be4c9
CI / Run CI (push) Has been cancelled
CI / check-backend (push) Has been cancelled
CI / check-frontend (push) Has been cancelled
CI / tests (push) Has been cancelled
CI / e2e-tests (push) Has been cancelled
Copilot Setup Steps / copilot-setup-steps (push) Has been cancelled
105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
from typing import Any, Dict, Optional, Union
|
|
|
|
from google.auth import default
|
|
from google.cloud import storage # type: ignore
|
|
from google.oauth2 import service_account
|
|
|
|
from chainlit import make_async
|
|
from chainlit.data.storage_clients.base import BaseStorageClient, storage_expiry_time
|
|
from chainlit.logger import logger
|
|
|
|
|
|
class GCSStorageClient(BaseStorageClient):
|
|
def __init__(
|
|
self,
|
|
bucket_name: str,
|
|
project_id: Optional[str] = None,
|
|
client_email: Optional[str] = None,
|
|
private_key: Optional[str] = None,
|
|
):
|
|
if client_email and private_key and project_id:
|
|
# Go to IAM & Admin, click on Service Accounts, and generate a new JSON key
|
|
logger.info("Using Private Key from Environment Variable")
|
|
credentials = service_account.Credentials.from_service_account_info(
|
|
{
|
|
"type": "service_account",
|
|
"project_id": project_id,
|
|
"private_key": private_key,
|
|
"client_email": client_email,
|
|
"token_uri": "https://oauth2.googleapis.com/token",
|
|
}
|
|
)
|
|
else:
|
|
# Application Default Credentials (e.g. in Google Cloud Run)
|
|
logger.info("Using Application Default Credentials.")
|
|
credentials, default_project_id = default()
|
|
if not project_id:
|
|
project_id = default_project_id
|
|
|
|
self.client = storage.Client(project=project_id, credentials=credentials)
|
|
self.bucket = self.client.bucket(bucket_name)
|
|
logger.info("GCSStorageClient initialized")
|
|
|
|
def sync_get_read_url(self, object_key: str) -> str:
|
|
return self.bucket.blob(object_key).generate_signed_url(
|
|
version="v4", expiration=storage_expiry_time, method="GET"
|
|
)
|
|
|
|
async def get_read_url(self, object_key: str) -> str:
|
|
return await make_async(self.sync_get_read_url)(object_key)
|
|
|
|
def sync_upload_file(
|
|
self,
|
|
object_key: str,
|
|
data: Union[bytes, str],
|
|
mime: str = "application/octet-stream",
|
|
overwrite: bool = True,
|
|
) -> Dict[str, Any]:
|
|
try:
|
|
blob = self.bucket.blob(object_key)
|
|
|
|
if not overwrite and blob.exists():
|
|
raise Exception(
|
|
f"File {object_key} already exists and overwrite is False"
|
|
)
|
|
|
|
if isinstance(data, str):
|
|
data = data.encode("utf-8")
|
|
|
|
blob.upload_from_string(data, content_type=mime)
|
|
|
|
# Return signed URL
|
|
return {
|
|
"object_key": object_key,
|
|
"url": self.sync_get_read_url(object_key),
|
|
}
|
|
|
|
except Exception as e:
|
|
raise Exception(f"Failed to upload file to GCS: {e!s}")
|
|
|
|
async def upload_file(
|
|
self,
|
|
object_key: str,
|
|
data: Union[bytes, str],
|
|
mime: str = "application/octet-stream",
|
|
overwrite: bool = True,
|
|
content_disposition: str | None = None,
|
|
) -> Dict[str, Any]:
|
|
return await make_async(self.sync_upload_file)(
|
|
object_key, data, mime, overwrite
|
|
)
|
|
|
|
def sync_delete_file(self, object_key: str) -> bool:
|
|
try:
|
|
self.bucket.blob(object_key).delete()
|
|
return True
|
|
except Exception as e:
|
|
logger.warning(f"GCSStorageClient, delete_file error: {e}")
|
|
return False
|
|
|
|
async def delete_file(self, object_key: str) -> bool:
|
|
return await make_async(self.sync_delete_file)(object_key)
|
|
|
|
async def close(self) -> None:
|
|
self.client.close()
|