|
4 | 4 | from okdata.sdk import SDK |
5 | 5 | from okdata.sdk.exceptions import ApiAuthenticateError |
6 | 6 |
|
7 | | - |
8 | 7 | log = logging.getLogger() |
9 | 8 |
|
10 | 9 |
|
11 | 10 | class Upload(SDK): |
12 | | - def __init__(self, config=None, auth=None): |
| 11 | + def __init__(self, config=None, auth=None, env=None): |
13 | 12 | self.__name__ = "upload" |
14 | | - super().__init__(config, auth) |
| 13 | + super().__init__(config, auth, env) |
15 | 14 |
|
16 | | - def upload(self, fileName, datasetid, versionid, editionid, retries=0): |
| 15 | + def upload(self, filename, dataset_id, version, edition, retries=0): |
17 | 16 | url = self.config.get("s3BucketUrl") |
18 | | - log.info(f"Uploading {fileName} to {datasetid} on: {url}") |
| 17 | + log.info(f"Uploading {filename} to {dataset_id} on: {url}") |
19 | 18 | if url is None: |
20 | | - raise KeyError("No s3 Bucket URL set") |
| 19 | + raise KeyError("No S3 bucket URL set") |
21 | 20 |
|
22 | | - s3SignedData = self.create_s3_signed_data( |
23 | | - fileName, datasetid, versionid, editionid, retries=retries |
| 21 | + s3_signed_data = self.create_s3_signed_data( |
| 22 | + filename, dataset_id, version, edition, retries=retries |
24 | 23 | ) |
25 | | - s3Data = {} |
26 | | - if "message" in s3SignedData: |
27 | | - # TODO: very specific error raised by the Lambda function, remove later |
28 | | - raise ApiAuthenticateError(s3SignedData["message"]) |
| 24 | + s3_data = {} |
| 25 | + if "message" in s3_signed_data: |
| 26 | + # TODO: Very specific error raised by the Lambda function, remove later |
| 27 | + raise ApiAuthenticateError(s3_signed_data["message"]) |
29 | 28 |
|
30 | | - for var in s3SignedData["fields"]: |
31 | | - s3Data[var] = s3SignedData["fields"][var] |
| 29 | + for var in s3_signed_data["fields"]: |
| 30 | + s3_data[var] = s3_signed_data["fields"][var] |
32 | 31 |
|
33 | | - with open(fileName, "rb") as file: |
| 32 | + with open(filename, "rb") as file: |
34 | 33 | files = {"file": file} |
35 | 34 | upload_session = self.prepared_request_with_retries(retries=retries) |
36 | | - result = upload_session.post(url, data=s3Data, files=files) |
37 | | - trace_id = s3SignedData.get("trace_id") |
38 | | - data = {"result": result.status_code == 204, "trace_id": trace_id} |
39 | | - return data |
40 | | - |
41 | | - def create_s3_signed_data( |
42 | | - self, fileName, datasetid, versionid, editionid, retries=0 |
43 | | - ): |
44 | | - edition = f"{datasetid}/{versionid}/{editionid}" |
45 | | - data = {"filename": os.path.basename(fileName), "editionId": edition} |
| 35 | + result = upload_session.post(url, data=s3_data, files=files) |
| 36 | + trace_id = s3_signed_data.get("trace_id") |
| 37 | + return {"result": result.status_code == 204, "trace_id": trace_id} |
| 38 | + |
| 39 | + def create_s3_signed_data(self, filename, dataset_id, version, edition, retries=0): |
| 40 | + edition_id = f"{dataset_id}/{version}/{edition}" |
| 41 | + data = {"filename": os.path.basename(filename), "editionId": edition_id} |
46 | 42 | url = self.config.get("uploadUrl") |
47 | | - log.info(f"Creating s3 signed data with payload: {data} on: {url}") |
| 43 | + log.info(f"Creating S3 signed data with payload: {data} on: {url}") |
48 | 44 | if url is None: |
49 | | - raise KeyError("No Signed S3 URL set") |
| 45 | + raise KeyError("No signed S3 URL set") |
50 | 46 | return self.post(url, data, retries=retries).json() |
0 commit comments