|
1 | | -from typing import Optional, Union |
| 1 | +from os import PathLike |
| 2 | +from typing import BinaryIO, Optional, Union |
2 | 3 |
|
3 | 4 | from upcloud_api.api import API |
4 | 5 | from upcloud_api.storage import Storage |
5 | 6 | from upcloud_api.storage_import import StorageImport |
6 | | -from upcloud_api.utils import get_raw_data_from_file |
7 | 7 |
|
8 | 8 |
|
9 | 9 | class StorageManager: |
@@ -41,10 +41,11 @@ def get_storage(self, storage: str) -> Storage: |
41 | 41 |
|
42 | 42 | def create_storage( |
43 | 43 | self, |
| 44 | + zone: str, |
44 | 45 | size: int = 10, |
45 | 46 | tier: str = 'maxiops', |
46 | 47 | title: str = 'Storage disk', |
47 | | - zone: str = 'fi-hel1', |
| 48 | + *, |
48 | 49 | backup_rule: Optional[dict] = None, |
49 | 50 | ) -> Storage: |
50 | 51 | """ |
@@ -184,33 +185,52 @@ def create_storage_import( |
184 | 185 | Creates an import task to import data into an existing storage. |
185 | 186 | Source types: http_import or direct_upload. |
186 | 187 | """ |
| 188 | + if source not in ("http_import", "direct_upload"): |
| 189 | + raise Exception(f"invalid storage import source: {source}") |
| 190 | + |
187 | 191 | url = f'/storage/{storage}/import' |
188 | 192 | body = {'storage_import': {'source': source}} |
189 | 193 | if source_location: |
190 | 194 | body['storage_import']['source_location'] = source_location |
191 | 195 | res = self.api.post_request(url, body) |
192 | 196 | return StorageImport(**res['storage_import']) |
193 | 197 |
|
194 | | - def upload_file_for_storage_import(self, storage_import, file): |
| 198 | + def upload_file_for_storage_import( |
| 199 | + self, |
| 200 | + storage_import: StorageImport, |
| 201 | + file: Union[str, PathLike, BinaryIO], |
| 202 | + timeout: int = 30, |
| 203 | + content_type: str = 'application/octet-stream', |
| 204 | + ): |
195 | 205 | """ |
196 | 206 | Uploads a file directly to UpCloud's uploader session. |
197 | 207 | """ |
198 | | - # TODO: this should not buffer the entire `file` into memory |
199 | | - |
200 | | - # This is importing and using `requests` directly since there doesn't |
201 | | - # seem to be a point in adding a `.api.raw_request()` call to the `API` class. |
202 | | - # That could be changed. |
203 | 208 |
|
204 | 209 | import requests |
205 | 210 |
|
206 | | - resp = requests.put( |
207 | | - url=storage_import.direct_upload_url, |
208 | | - data=get_raw_data_from_file(file), |
209 | | - headers={'Content-type': 'application/octet-stream'}, |
210 | | - timeout=600, |
211 | | - ) |
212 | | - resp.raise_for_status() |
213 | | - return resp.json() |
| 211 | + # This is importing and using `requests` directly since there doesn't |
| 212 | + # seem to be a point in adding a `.api.raw_request()` call to the `API` class. |
| 213 | + # That could be changed if there starts to be more of these cases. |
| 214 | + |
| 215 | + f = file |
| 216 | + needs_closing = False |
| 217 | + if not hasattr(file, 'read'): |
| 218 | + f = open(file, 'rb') |
| 219 | + needs_closing = True |
| 220 | + |
| 221 | + try: |
| 222 | + resp = requests.put( |
| 223 | + url=storage_import.direct_upload_url, |
| 224 | + data=f, |
| 225 | + headers={'Content-type': content_type}, |
| 226 | + timeout=timeout, |
| 227 | + ) |
| 228 | + |
| 229 | + resp.raise_for_status() |
| 230 | + return resp.json() |
| 231 | + finally: |
| 232 | + if needs_closing: |
| 233 | + f.close() |
214 | 234 |
|
215 | 235 | def get_storage_import_details(self, storage: str) -> StorageImport: |
216 | 236 | """ |
|
0 commit comments