diff --git a/fbpcp/entity/file_information.py b/fbpcp/entity/file_information.py index 7d9eda08..a687d86a 100644 --- a/fbpcp/entity/file_information.py +++ b/fbpcp/entity/file_information.py @@ -6,6 +6,7 @@ # pyre-strict from dataclasses import dataclass +from typing import Any, Dict, Optional @dataclass @@ -13,3 +14,4 @@ class FileInfo: file_name: str last_modified: str file_size: int + metadata: Optional[Dict[str, Any]] = None diff --git a/fbpcp/gateway/s3.py b/fbpcp/gateway/s3.py index 38d9c6bc..a597e2d1 100644 --- a/fbpcp/gateway/s3.py +++ b/fbpcp/gateway/s3.py @@ -47,13 +47,21 @@ def delete_bucket(self, bucket: str) -> None: self.client.delete_bucket(Bucket=bucket) @error_handler - def upload_file(self, file_name: str, bucket: str, key: str) -> None: + def upload_file( + self, + file_name: str, + bucket: str, + key: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: file_size = os.path.getsize(file_name) + metadata = metadata if metadata else {} self.client.upload_file( file_name, bucket, key, Callback=self.ProgressPercentage(file_name, file_size), + ExtraArgs={"Metadata": metadata}, ) @error_handler diff --git a/fbpcp/service/storage.py b/fbpcp/service/storage.py index 7d91d50d..a2bb8105 100644 --- a/fbpcp/service/storage.py +++ b/fbpcp/service/storage.py @@ -9,7 +9,7 @@ import abc import re from enum import Enum -from typing import List +from typing import Any, Dict, List, Optional from fbpcp.entity.file_information import FileInfo from fbpcp.entity.policy_statement import PolicyStatement, PublicAccessBlockConfig @@ -75,3 +75,9 @@ def get_bucket_public_access_block(self, bucket: str) -> PublicAccessBlockConfig @abc.abstractmethod def list_files(self, dirPath: str) -> List[str]: pass + + @abc.abstractmethod + def upload_file( + self, source: str, destination: str, metadata: Optional[Dict[str, Any]] = None + ) -> None: + pass diff --git a/fbpcp/service/storage_gcs.py b/fbpcp/service/storage_gcs.py index 7d332e6b..c4973f92 100644 --- a/fbpcp/service/storage_gcs.py +++ b/fbpcp/service/storage_gcs.py @@ -127,6 +127,11 @@ def copy(self, source: str, destination: str, recursive: bool = False) -> None: dest_key=destination_gcs_path.key, ) + def upload_file( + self, source: str, destination: str, metadata: Optional[Dict[str, Any]] + ) -> None: + raise NotImplementedError + def upload_dir(self, source: str, gcs_path_bucket: str, gcs_path_key: str) -> None: """Upload a directory from the filesystem to GCS diff --git a/fbpcp/service/storage_s3.py b/fbpcp/service/storage_s3.py index ca2d47cb..ea1c6189 100644 --- a/fbpcp/service/storage_s3.py +++ b/fbpcp/service/storage_s3.py @@ -64,7 +64,7 @@ def copy(self, source: str, destination: str, recursive: bool = False) -> None: raise ValueError(f"Source {source} is a folder. Use --recursive") self.upload_dir(source, s3_path.bucket, s3_path.key) else: - self.s3_gateway.upload_file(source, s3_path.bucket, s3_path.key) + self.upload_file(source, destination) else: source_s3_path = S3Path(source) if StorageService.path_type(destination) == PathType.S3: @@ -111,6 +111,12 @@ def copy(self, source: str, destination: str, recursive: bool = False) -> None: source_s3_path.bucket, source_s3_path.key, destination ) + def upload_file( + self, source: str, destination: str, metadata: Optional[Dict[str, Any]] = None + ) -> None: + s3_path = S3Path(destination) + self.s3_gateway.upload_file(source, s3_path.bucket, s3_path.key, metadata) + def upload_dir(self, source: str, s3_path_bucket: str, s3_path_key: str) -> None: for root, dirs, files in os.walk(source): for file in files: @@ -205,6 +211,7 @@ def get_file_info(self, filename: str) -> FileInfo: file_name=filename, last_modified=file_info_dict.get("LastModified").ctime(), file_size=file_info_dict.get("ContentLength"), + metadata=file_info_dict.get("Metadata"), ) def get_file_size(self, filename: str) -> int: diff --git a/onedocker/repository/onedocker_package.py b/onedocker/repository/onedocker_package.py index 22a82ca7..f12fe85d 100644 --- a/onedocker/repository/onedocker_package.py +++ b/onedocker/repository/onedocker_package.py @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. # pyre-strict -from typing import List +from typing import Any, Dict, List, Optional from fbpcp.service.storage import StorageService from onedocker.entity.package_info import PackageInfo @@ -19,9 +19,15 @@ def __init__(self, storage_svc: StorageService, repository_path: str) -> None: def _build_package_path(self, package_name: str, version: str) -> str: return f"{self.repository_path}{package_name}/{version}/{package_name.split('/')[-1]}" - def upload(self, package_name: str, version: str, source: str) -> None: + def upload( + self, + package_name: str, + version: str, + source: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: package_path = self._build_package_path(package_name, version) - self.storage_svc.copy(source, package_path) + self.storage_svc.upload_file(source, package_path, metadata) def download(self, package_name: str, version: str, destination: str) -> None: package_path = self._build_package_path(package_name, version) diff --git a/onedocker/repository/onedocker_repository_service.py b/onedocker/repository/onedocker_repository_service.py index 3d5de746..eff3ca2b 100644 --- a/onedocker/repository/onedocker_repository_service.py +++ b/onedocker/repository/onedocker_repository_service.py @@ -1,8 +1,10 @@ +#!/usr/bin/env python3 # Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from datetime import datetime from typing import Optional from fbpcp.service.storage import StorageService @@ -10,7 +12,7 @@ from onedocker.repository.onedocker_package import OneDockerPackageRepository -class OnedockerRepositoryService: +class OneDockerRepositoryService: def __init__( self, storage_svc: StorageService, @@ -32,7 +34,10 @@ def upload( source: str, metadata: Optional[dict] = None, ) -> None: - raise NotImplementedError + today = datetime.today().strftime("%Y-%m-%d") + metadata = metadata if metadata else {} + metadata["upload_date"] = today + self.package_repo.upload(package_name, version, source, metadata) def download(self, package_name: str, version: str, destination: str) -> None: raise NotImplementedError diff --git a/onedocker/tests/repository/test_onedocker_package.py b/onedocker/tests/repository/test_onedocker_package.py index 285c0299..9e01d01f 100644 --- a/onedocker/tests/repository/test_onedocker_package.py +++ b/onedocker/tests/repository/test_onedocker_package.py @@ -36,8 +36,8 @@ def test_onedockerrepo_upload(self): ) # Assert - self.onedocker_repository.storage_svc.copy.assert_called_with( - source, self.expected_s3_dest + self.onedocker_repository.storage_svc.upload_file.assert_called_with( + source, self.expected_s3_dest, None ) def test_onedockerrepo_download(self): diff --git a/onedocker/tests/repository/test_onedocker_repository_service.py b/onedocker/tests/repository/test_onedocker_repository_service.py new file mode 100644 index 00000000..62a9c884 --- /dev/null +++ b/onedocker/tests/repository/test_onedocker_repository_service.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from datetime import datetime +from unittest.mock import MagicMock, patch + +from onedocker.repository.onedocker_repository_service import OneDockerRepositoryService + + +class TestOneDockerRepositoryService(unittest.TestCase): + TEST_PACKAGE_PATH = "private_lift/lift" + TEST_PACKAGE_NAME = TEST_PACKAGE_PATH.split("/")[-1] + TEST_PACKAGE_VERSION = "latest" + + @patch( + "onedocker.repository.onedocker_repository_service.OneDockerChecksumRepository" + ) + @patch( + "onedocker.repository.onedocker_repository_service.OneDockerPackageRepository" + ) + @patch("fbpcp.service.storage_s3.S3StorageService") + def setUp( + self, mockStorageService, mockPackageRepoCall, mockChecksumRepoCall + ) -> None: + package_repo_path = "/package_repo_path/" + checksum_repo_path = "/checksum_repo_path/" + self.package_repo = MagicMock() + mockPackageRepoCall.return_value = self.package_repo + self.repo_service = OneDockerRepositoryService( + mockStorageService, package_repo_path, checksum_repo_path + ) + + def test_onedocker_repo_service_upload(self) -> None: + # Arrange + source_path = "test_source_path" + today = datetime.today().strftime("%Y-%m-%d") + metadata = {"upload_date": today} + + # Act + self.repo_service.upload( + self.TEST_PACKAGE_PATH, self.TEST_PACKAGE_VERSION, source_path + ) + + # Assert + self.package_repo.upload.assert_called_with( + self.TEST_PACKAGE_PATH, self.TEST_PACKAGE_VERSION, source_path, metadata + ) diff --git a/tests/service/test_storage_s3.py b/tests/service/test_storage_s3.py index d9b49c95..fa081769 100644 --- a/tests/service/test_storage_s3.py +++ b/tests/service/test_storage_s3.py @@ -48,7 +48,7 @@ def test_copy_local_to_s3(self, MockS3Gateway): service.s3_gateway.upload_file = MagicMock(return_value=None) service.copy(self.LOCAL_FILE, self.S3_FILE) service.s3_gateway.upload_file.assert_called_with( - str(self.LOCAL_FILE), "bucket", "test_file" + str(self.LOCAL_FILE), "bucket", "test_file", None ) def test_copy_local_dir_to_s3_recursive_false(self):