stable-diffusion-aws-extension/middleware_api/lambda/common/util.py

45 lines
1.6 KiB
Python

from typing import Dict
import boto3
def publish_msg(topic_arn, msg, subject):
client = boto3.client('sns')
client.publish(
TopicArn=topic_arn,
Message=str(msg),
Subject=subject
)
def get_s3_presign_urls(bucket_name, base_key, filenames) -> Dict[str, str]:
return _get_s3_presign_urls(bucket_name, base_key, filenames, expires=3600 * 24 * 7, method='put_object')
def get_s3_get_presign_urls(bucket_name, base_key, filenames) -> Dict[str, str]:
return _get_s3_presign_urls(bucket_name, base_key, filenames, expires=3600 * 24, method='get_object')
def _get_s3_presign_urls(bucket_name, base_key, filenames, expires=3600, method='put_object') -> Dict[str, str]:
s3 = boto3.client('s3')
presign_url_map = {}
for filename in filenames:
key = f'{base_key}/{filename}'
url = s3.generate_presigned_url(method,
Params={'Bucket': bucket_name,
'Key': key,
},
ExpiresIn=expires)
presign_url_map[filename] = url
return presign_url_map
def generate_presign_url(bucket_name, key, expires=3600, method='put_object') -> Dict[str, str]:
s3 = boto3.client('s3')
return s3.generate_presigned_url(method,
Params={'Bucket': bucket_name,
'Key': key,
},
ExpiresIn=expires)