-
-
Notifications
You must be signed in to change notification settings - Fork 57
/
r2_utils.py
28 lines (22 loc) · 908 Bytes
/
r2_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import io
import os
import boto3
R2_ACCESS_KEY_ID = os.environ["R2_ACCESS_KEY_ID"]
R2_SECRET_ACCESS_KEY = os.environ["R2_SECRET_ACCESS_KEY"]
R2_S3_ENDPOINT_URL = os.environ["R2_S3_ENDPOINT_URL"]
R2_DOMAIN_NAME = os.environ["R2_DOMAIN_NAME"]
s3 = boto3.client(
service_name ="s3",
endpoint_url = R2_S3_ENDPOINT_URL,
aws_access_key_id = R2_ACCESS_KEY_ID,
aws_secret_access_key = R2_SECRET_ACCESS_KEY,
region_name="auto", # Must be one of: wnam, enam, weur, eeur, apac, auto
)
def upload_to_r2(bucket_name, arxiv_id, filename):
obj_name = os.path.join(arxiv_id, os.path.basename(filename))
with open(filename, 'rb') as f:
file_content = f.read()
s3.upload_fileobj(io.BytesIO(file_content), bucket_name, obj_name)
return f"https://{R2_DOMAIN_NAME}/{obj_name}"
def remove_r2_directory(bucket_name, arxiv_id):
s3.delete_object(Bucket=bucket_name, Key=arxiv_id)