Files
annotations/build/cdn_manager.py
T
dzaitsev d92da6afa4 Errors sending to UI
notifying client of AI model conversion
2025-05-14 12:43:50 +03:00

88 lines
3.4 KiB
Python

import io
import sys
import boto3
import yaml
import os
class CDNCredentials:
def __init__(self, host, downloader_access_key, downloader_access_secret, uploader_access_key, uploader_access_secret):
self.host = host
self.downloader_access_key = downloader_access_key
self.downloader_access_secret = downloader_access_secret
self.uploader_access_key = uploader_access_key
self.uploader_access_secret = uploader_access_secret
class CDNManager:
def __init__(self, credentials: CDNCredentials):
self.creds = credentials
self.download_client = boto3.client('s3', endpoint_url=self.creds.host,
aws_access_key_id=self.creds.downloader_access_key,
aws_secret_access_key=self.creds.downloader_access_secret)
self.upload_client = boto3.client('s3', endpoint_url=self.creds.host,
aws_access_key_id=self.creds.uploader_access_key,
aws_secret_access_key=self.creds.uploader_access_secret)
def upload(self, bucket: str, filename: str, file_bytes: bytearray):
try:
self.upload_client.upload_fileobj(io.BytesIO(file_bytes), bucket, filename)
print(f'uploaded {filename} ({len(file_bytes)} bytes) to the {bucket}')
return True
except Exception as e:
print(e)
return False
def download(self, folder: str, filename: str):
try:
if filename is not None:
self.download_client.download_file(folder, filename, f'{folder}\\{filename}')
print(f'downloaded {filename} from the {folder} to current folder')
return True
else:
response = self.download_client.list_objects_v2(Bucket=folder)
if 'Contents' in response:
for obj in response['Contents']:
object_key = obj['Key']
local_filepath = os.path.join(folder, object_key)
local_dir = os.path.dirname(local_filepath)
if local_dir:
os.makedirs(local_dir, exist_ok=True)
if not object_key.endswith('/'):
try:
self.download_client.download_file(folder, object_key, local_filepath)
except Exception as e_file:
all_successful = False # Mark as failed if any file fails
return True
except Exception as e:
print(e)
return False
config_file = os.path.join('downloader_config.yaml')
with open(config_file, 'r') as f:
config = yaml.safe_load(f)
cdn_manager = CDNManager(CDNCredentials(
config['host'],
config['downloader_access_key'],
config['downloader_access_secret'],
config['uploader_access_key'],
config['uploader_access_secret']
))
input_action = sys.argv[1]
input_bucket = sys.argv[2]
input_filename = sys.argv[3] if len(sys.argv) > 3 else None
if len(sys.argv) > 4: # 0 is this script's path, hence 5 args is max
input_path = sys.argv[4]
if input_action == 'upload':
with open(input_path, 'rb') as f:
cdn_manager.upload(input_bucket, input_filename, bytearray(f.read()))
elif input_action == 'download':
cdn_manager.download(input_bucket, input_filename)