From 3a2819f89d18baeed4f6f3038f0cf0c6f559dfb9 Mon Sep 17 00:00:00 2001
From: yan <yan_sushchynski@epam.com>
Date: Thu, 10 Dec 2020 17:01:08 +0300
Subject: [PATCH] Add content type to upload file. Fix typo
---
src/dags/libs/upload_file.py | 33 +++++++++++++++----
...est_file_uplaod.py => test_file_upload.py} | 3 +-
2 files changed, 29 insertions(+), 7 deletions(-)
rename tests/plugin-unit-tests/{test_file_uplaod.py => test_file_upload.py} (97%)
diff --git a/src/dags/libs/upload_file.py b/src/dags/libs/upload_file.py
index 73e260d..4014e97 100644
--- a/src/dags/libs/upload_file.py
+++ b/src/dags/libs/upload_file.py
@@ -136,10 +136,14 @@ class GCSFileUploader(FileUploader):
raise GCSObjectURIError
@tenacity.retry(**RETRY_SETTINGS)
- def get_file_from_bucket(self, bucket_name: str, source_blob_name: str) -> io.BytesIO:
+ def get_file_from_bucket(
+ self,
+ bucket_name: str,
+ source_blob_name: str
+ ) -> Tuple[io.BytesIO, str]:
storage_client = storage.Client()
bucket = storage_client.bucket(bucket_name)
- blob = bucket.blob(source_blob_name)
+ blob = bucket.get_blob(source_blob_name)
does_exist = blob.exists()
if not does_exist:
@@ -149,9 +153,26 @@ class GCSFileUploader(FileUploader):
file = io.BytesIO()
blob.download_to_file(file)
logger.debug("File got from landing zone")
- return file
+ return file, blob.content_type
- def get_file_from_preload_path(self, preload_file_path: str) -> io.BytesIO:
+ def get_file_from_preload_path(self, preload_file_path: str) -> Tuple[io.BytesIO, str]:
bucket_name, blob_name = self._parse_object_uri(preload_file_path)
- buffer = self.get_file_from_bucket(bucket_name, blob_name)
- return buffer
+ buffer, content_type = self.get_file_from_bucket(bucket_name, blob_name)
+ return buffer, content_type
+
+ def upload_file(self, preload_file_path: str) -> str:
+ """
+ Copy file from Landing zone(preload_file_path) onto OSDU platform using File service.
+ Get Content-Type of this file, refresh Content-Type with this value in headers
+ while this file is being uploaded onto OSDU platform.
+ Return file_location.
+ """
+ buffer, content_type = self.get_file_from_preload_path(preload_file_path)
+ file_id, signed_url = self._get_signed_url_request(self.request_headers)
+
+ headers = self.request_headers
+ headers["Content-Type"] = content_type
+ self._upload_file_request(headers, signed_url, buffer)
+
+ file_location = self._get_file_location_request(self.request_headers, file_id)
+ return file_location
diff --git a/tests/plugin-unit-tests/test_file_uplaod.py b/tests/plugin-unit-tests/test_file_upload.py
similarity index 97%
rename from tests/plugin-unit-tests/test_file_uplaod.py
rename to tests/plugin-unit-tests/test_file_upload.py
index 9042b30..246d768 100644
--- a/tests/plugin-unit-tests/test_file_uplaod.py
+++ b/tests/plugin-unit-tests/test_file_upload.py
@@ -48,7 +48,8 @@ class TestSourceFileChecker:
file_uploader: GCSFileUploader
):
file = io.RawIOBase()
- monkeypatch.setattr(file_uploader, "get_file_from_bucket", lambda *args, **kwargs: file)
+ monkeypatch.setattr(file_uploader, "get_file_from_bucket",
+ lambda *args, **kwargs: (file, "test"))
file_uploader.upload_file("gs://test/test")
@pytest.mark.parametrize(
--
GitLab