forked from kiri-art/docker-diffusers-api
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathS3Storage.py
More file actions
128 lines (105 loc) · 3.74 KB
/
Copy pathS3Storage.py
File metadata and controls
128 lines (105 loc) · 3.74 KB
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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import boto3
import botocore
import re
import os
import time
from tqdm import tqdm
from botocore.client import Config
AWS_S3_ENDPOINT_URL = os.environ.get("AWS_S3_ENDPOINT_URL", None)
AWS_S3_DEFAULT_BUCKET = os.environ.get("AWS_S3_DEFAULT_BUCKET", None)
if AWS_S3_ENDPOINT_URL == "":
AWS_S3_ENDPOINT_URL = None
if AWS_S3_DEFAULT_BUCKET == "":
AWS_S3_DEFAULT_BUCKET = None
def get_now():
return round(time.time() * 1000)
class S3Storage:
def __init__(self, url, path=""):
self.url = url
if url.startswith("s3://"):
url = "https://" + url[5:]
elif url.startswith("http+s3://"):
url = "http" + url[7:]
elif url.startswith("https+s3://"):
url = "https" + url[8:]
s3_dest = re.match(
"^(?P<endpoint>https?://[^/]*)(/(?P<bucket>[^/]+))?(/(?P<path>.*))?$",
url,
).groupdict()
if not s3_dest["endpoint"] or s3_dest["endpoint"].endswith("//"):
s3_dest["endpoint"] = AWS_S3_ENDPOINT_URL
if not s3_dest["bucket"]:
s3_dest["bucket"] = AWS_S3_DEFAULT_BUCKET
if not s3_dest["path"] or s3_dest["path"] == "":
s3_dest["path"] = path
self.endpoint_url = s3_dest["endpoint"]
self.bucket_name = s3_dest["bucket"]
self.path = s3_dest["path"]
self._s3resource = None
self._s3client = None
self._bucket = None
print("self.endpoint_url", self.endpoint_url)
def s3resource(self):
if self._s3resource:
return self._s3resource
self._s3 = boto3.resource(
"s3",
endpoint_url=self.endpoint_url,
config=Config(signature_version="s3v4"),
)
return self._s3
def s3client(self):
if self._s3client:
return self._s3client
self._s3client = boto3.client(
"s3",
endpoint_url=self.endpoint_url,
config=Config(signature_version="s3v4"),
)
return self._s3client
def bucket(self):
if self._bucket:
return self._bucket
self._bucket = self.s3resource().Bucket(self.bucket_name)
return self._bucket
def upload_file(self, source, dest):
if not dest:
dest = self.path
upload_start = get_now()
file_size = os.stat(source).st_size
with tqdm(total=file_size, unit="B", unit_scale=True, desc="Uploading") as bar:
result = self.bucket().upload_file(
Filename=source,
Key=dest,
Callback=lambda bytes_transferred: bar.update(bytes_transferred),
)
print(result)
upload_total = get_now() - upload_start
return {"$time": upload_total}
def download_file(self, dest):
if not dest:
dest = self.path.split("/").pop()
print(f"Downloading {self.url} to {dest}...")
object = self.s3resource().Object(self.bucket_name, self.path)
object.load()
with tqdm(
total=object.content_length, unit="B", unit_scale=True, desc="Downloading"
) as bar:
object.download_file(
Filename=dest,
Callback=lambda bytes_transffered: bar.update(bytes_transffered),
)
def file_exists(self):
# res = self.s3client().list_objects_v2(
# Bucket=self.bucket_name, Prefix=self.path, MaxKeys=1
# )
# return "Contents" in res
object = self.s3resource().Object(self.bucket_name, self.path)
try:
object.load()
except botocore.exceptions.ClientError as error:
if error.response["Error"]["Code"] == "404":
return False
else:
raise
return True