- 11class ExecEnv(object):
-12 """Represents the basic resources regarding the engine execution environment.
-13
-14 Currently, it is solely used to encapsulate the logic to get a Spark session.
-15 """
-16
-17 SESSION: SparkSession
-18 _LOGGER = LoggingHandler(__name__).get_logger()
-19 DEFAULT_AWS_REGION = "eu-west-1"
-20
-21 @classmethod
-22 def get_or_create(
-23 cls,
-24 session: SparkSession = None,
-25 enable_hive_support: bool = True,
-26 app_name: str = None,
-27 config: dict = None,
-28 ) -> None:
-29 """Get or create an execution environment session (currently Spark).
-30
-31 It instantiates a singleton session that can be accessed anywhere from the
-32 lakehouse engine.
-33
-34 Args:
-35 session: spark session.
-36 enable_hive_support: whether to enable hive support or not.
-37 app_name: application name.
-38 config: extra spark configs to supply to the spark session.
-39 """
-40 default_config = {
-41 "spark.databricks.delta.optimizeWrite.enabled": True,
-42 "spark.sql.adaptive.enabled": True,
-43 "spark.databricks.delta.merge.enableLowShuffle": True,
-44 }
-45 cls._LOGGER.info(
-46 f"Using the following default configs you may want to override them for "
-47 f"your job: {default_config}"
-48 )
-49 final_config: dict = {**default_config, **(config if config else {})}
-50 cls._LOGGER.info(f"Final config is: {final_config}")
-51
-52 if session:
-53 cls.SESSION = session
-54 else:
-55 session_builder = SparkSession.builder.appName(app_name)
-56 if config:
-57 session_builder = session_builder.config(
-58 conf=SparkConf().setAll(final_config.items()) # type: ignore
-59 )
-60 if enable_hive_support:
-61 session_builder = session_builder.enableHiveSupport()
-62 cls.SESSION = session_builder.getOrCreate()
-63
-64 cls._set_environment_variables(final_config.get("os_env_vars"))
-65
-66 @classmethod
-67 def _set_environment_variables(cls, os_env_vars: dict = None) -> None:
-68 """Set environment variables at OS level.
-69
-70 By default, we are setting the AWS_DEFAULT_REGION as we have identified this is
-71 beneficial to avoid getBucketLocation permission problems.
+ 10class ExecEnv(object):
+11 """Represents the basic resources regarding the engine execution environment.
+12
+13 Currently, it is solely used to encapsulate the logic to get a Spark session.
+14 """
+15
+16 SESSION: SparkSession
+17 _LOGGER = LoggingHandler(__name__).get_logger()
+18 DEFAULT_AWS_REGION = "eu-west-1"
+19
+20 @classmethod
+21 def get_or_create(
+22 cls,
+23 session: SparkSession = None,
+24 enable_hive_support: bool = True,
+25 app_name: str = None,
+26 config: dict = None,
+27 ) -> None:
+28 """Get or create an execution environment session (currently Spark).
+29
+30 It instantiates a singleton session that can be accessed anywhere from the
+31 lakehouse engine.
+32
+33 Args:
+34 session: spark session.
+35 enable_hive_support: whether to enable hive support or not.
+36 app_name: application name.
+37 config: extra spark configs to supply to the spark session.
+38 """
+39 default_config = {
+40 "spark.databricks.delta.optimizeWrite.enabled": True,
+41 "spark.sql.adaptive.enabled": True,
+42 "spark.databricks.delta.merge.enableLowShuffle": True,
+43 }
+44 cls._LOGGER.info(
+45 f"Using the following default configs you may want to override them for "
+46 f"your job: {default_config}"
+47 )
+48 final_config: dict = {**default_config, **(config if config else {})}
+49 cls._LOGGER.info(f"Final config is: {final_config}")
+50
+51 if session:
+52 cls.SESSION = session
+53 else:
+54 # with active session we do not need app name
+55 if SparkSession.getActiveSession():
+56 app_name = SparkSession.getActiveSession().sparkContext.appName
+57 cls._LOGGER.info(f"Detected active session: {app_name}")
+58 elif not SparkSession.getActiveSession() and not app_name:
+59 cls._LOGGER.info("No active session or appname detected")
+60 app_name = "lakehouse_engine"
+61 # we will still add this part to set configs
+62 session_builder = SparkSession.builder.appName(app_name)
+63 if config:
+64 for k, v in final_config.items():
+65 session_builder.config(k, v)
+66
+67 if enable_hive_support:
+68 session_builder = session_builder.enableHiveSupport()
+69 cls.SESSION = session_builder.getOrCreate()
+70
+71 cls._set_environment_variables(final_config.get("os_env_vars"))
72
-73 Args:
-74 os_env_vars: this parameter can be used to pass the environment variables to
-75 be defined.
-76 """
-77 if os_env_vars is None:
-78 os_env_vars = {}
+73 @classmethod
+74 def _set_environment_variables(cls, os_env_vars: dict = None) -> None:
+75 """Set environment variables at OS level.
+76
+77 By default, we are setting the AWS_DEFAULT_REGION as we have identified this is
+78 beneficial to avoid getBucketLocation permission problems.
79
-80 for env_var in os_env_vars.items():
-81 os.environ[env_var[0]] = env_var[1]
-82
-83 if "AWS_DEFAULT_REGION" not in os_env_vars:
-84 os.environ["AWS_DEFAULT_REGION"] = cls.SESSION.sparkContext.getConf().get(
-85 "spark.databricks.clusterUsageTags.region", cls.DEFAULT_AWS_REGION
-86 )
+80 Args:
+81 os_env_vars: this parameter can be used to pass the environment variables to
+82 be defined.
+83 """
+84 if os_env_vars is None:
+85 os_env_vars = {}
+86
+87 for env_var in os_env_vars.items():
+88 os.environ[env_var[0]] = env_var[1]
+89
+90 if "AWS_DEFAULT_REGION" not in os_env_vars:
+91 os.environ["AWS_DEFAULT_REGION"] = cls.SESSION.sparkContext.getConf().get(
+92 "spark.databricks.clusterUsageTags.region", cls.DEFAULT_AWS_REGION
+93 )
21 @classmethod
-22 def get_or_create(
-23 cls,
-24 session: SparkSession = None,
-25 enable_hive_support: bool = True,
-26 app_name: str = None,
-27 config: dict = None,
-28 ) -> None:
-29 """Get or create an execution environment session (currently Spark).
-30
-31 It instantiates a singleton session that can be accessed anywhere from the
-32 lakehouse engine.
-33
-34 Args:
-35 session: spark session.
-36 enable_hive_support: whether to enable hive support or not.
-37 app_name: application name.
-38 config: extra spark configs to supply to the spark session.
-39 """
-40 default_config = {
-41 "spark.databricks.delta.optimizeWrite.enabled": True,
-42 "spark.sql.adaptive.enabled": True,
-43 "spark.databricks.delta.merge.enableLowShuffle": True,
-44 }
-45 cls._LOGGER.info(
-46 f"Using the following default configs you may want to override them for "
-47 f"your job: {default_config}"
-48 )
-49 final_config: dict = {**default_config, **(config if config else {})}
-50 cls._LOGGER.info(f"Final config is: {final_config}")
-51
-52 if session:
-53 cls.SESSION = session
-54 else:
-55 session_builder = SparkSession.builder.appName(app_name)
-56 if config:
-57 session_builder = session_builder.config(
-58 conf=SparkConf().setAll(final_config.items()) # type: ignore
-59 )
-60 if enable_hive_support:
-61 session_builder = session_builder.enableHiveSupport()
-62 cls.SESSION = session_builder.getOrCreate()
-63
-64 cls._set_environment_variables(final_config.get("os_env_vars"))
+ 20 @classmethod
+21 def get_or_create(
+22 cls,
+23 session: SparkSession = None,
+24 enable_hive_support: bool = True,
+25 app_name: str = None,
+26 config: dict = None,
+27 ) -> None:
+28 """Get or create an execution environment session (currently Spark).
+29
+30 It instantiates a singleton session that can be accessed anywhere from the
+31 lakehouse engine.
+32
+33 Args:
+34 session: spark session.
+35 enable_hive_support: whether to enable hive support or not.
+36 app_name: application name.
+37 config: extra spark configs to supply to the spark session.
+38 """
+39 default_config = {
+40 "spark.databricks.delta.optimizeWrite.enabled": True,
+41 "spark.sql.adaptive.enabled": True,
+42 "spark.databricks.delta.merge.enableLowShuffle": True,
+43 }
+44 cls._LOGGER.info(
+45 f"Using the following default configs you may want to override them for "
+46 f"your job: {default_config}"
+47 )
+48 final_config: dict = {**default_config, **(config if config else {})}
+49 cls._LOGGER.info(f"Final config is: {final_config}")
+50
+51 if session:
+52 cls.SESSION = session
+53 else:
+54 # with active session we do not need app name
+55 if SparkSession.getActiveSession():
+56 app_name = SparkSession.getActiveSession().sparkContext.appName
+57 cls._LOGGER.info(f"Detected active session: {app_name}")
+58 elif not SparkSession.getActiveSession() and not app_name:
+59 cls._LOGGER.info("No active session or appname detected")
+60 app_name = "lakehouse_engine"
+61 # we will still add this part to set configs
+62 session_builder = SparkSession.builder.appName(app_name)
+63 if config:
+64 for k, v in final_config.items():
+65 session_builder.config(k, v)
+66
+67 if enable_hive_support:
+68 session_builder = session_builder.enableHiveSupport()
+69 cls.SESSION = session_builder.getOrCreate()
+70
+71 cls._set_environment_variables(final_config.get("os_env_vars"))
1"""File manager module."""
2import time
- 3from typing import Any, Optional
+ 3from typing import Any, Optional, Tuple
4
5import boto3
6
@@ -111,516 +111,609 @@
14from lakehouse_engine.utils.logging_handler import LoggingHandler
15
16
- 17def _dry_run(bucket: str, object_paths: list) -> dict:
- 18 """Build the dry run request return format.
+ 17def _process_directory_path(path: str) -> str:
+ 18 """Add '/' to the end of the path of a directory.
19
20 Args:
- 21 bucket: name of bucket to perform operation.
- 22 object_paths: paths of object to list.
- 23
- 24 Returns:
- 25 A dict with a list of objects that would be copied/deleted.
- 26 """
- 27 response = {}
+ 21 path: directory to be processed
+ 22
+ 23 Returns:
+ 24 Directory path stripped and with '/' at the end.
+ 25 """
+ 26 path = path.strip()
+ 27 return path if path[-1] == "/" else path + "/"
28
- 29 for path in object_paths:
- 30 path = path.strip()
- 31 res = _list_objects_recursively(bucket=bucket, path=path)
+ 29
+ 30def _dry_run(bucket: str, object_paths: list) -> dict:
+ 31 """Build the dry run request return format.
32
- 33 if res:
- 34 response[path] = res
- 35 else:
- 36 response[path] = ["No such key"]
- 37
- 38 return response
- 39
- 40
- 41def _list_objects_recursively(bucket: str, path: str) -> list:
- 42 """Recursively list all objects given a prefix in s3.
- 43
- 44 Args:
- 45 bucket: name of bucket to perform the list.
- 46 path: path to be used as a prefix.
+ 33 Args:
+ 34 bucket: name of bucket to perform operation.
+ 35 object_paths: paths of object to list.
+ 36
+ 37 Returns:
+ 38 A dict with a list of objects that would be copied/deleted.
+ 39 """
+ 40 response = {}
+ 41
+ 42 for path in object_paths:
+ 43 if _check_directory(bucket, path):
+ 44 path = _process_directory_path(path)
+ 45
+ 46 res = _list_objects_recursively(bucket=bucket, path=path)
47
- 48 Returns:
- 49 A list of object names fetched recursively.
- 50 """
- 51 object_list = []
- 52 more_objects = True
- 53 pagination = ""
+ 48 if res:
+ 49 response[path] = res
+ 50 else:
+ 51 response[path] = ["No such key"]
+ 52
+ 53 return response
54
- 55 s3 = boto3.client("s3")
- 56
- 57 while more_objects:
- 58 if not pagination:
- 59 list_response = s3.list_objects_v2(Bucket=bucket, Prefix=path)
- 60 else:
- 61 list_response = s3.list_objects_v2(
- 62 Bucket=bucket,
- 63 Prefix=path,
- 64 ContinuationToken=pagination,
- 65 )
- 66
- 67 if FileManagerAPIKeys.CONTENTS.value in list_response:
- 68 for obj in list_response[FileManagerAPIKeys.CONTENTS.value]:
- 69 object_list.append(obj[FileManagerAPIKeys.KEY.value])
+ 55
+ 56def _list_objects(
+ 57 s3_client: Any, bucket: str, path: str, paginator: str = ""
+ 58) -> Tuple[list, str]:
+ 59 """List 1000 objects in a bucket given a prefix and paginator in s3.
+ 60
+ 61 Args:
+ 62 bucket: name of bucket to perform the list.
+ 63 path: path to be used as a prefix.
+ 64 paginator: paginator token to be used.
+ 65
+ 66 Returns:
+ 67 A list of object names.
+ 68 """
+ 69 object_list = []
70
- 71 if FileManagerAPIKeys.CONTINUATION.value in list_response:
- 72 pagination = list_response[FileManagerAPIKeys.CONTINUATION.value]
- 73 else:
- 74 more_objects = False
- 75
- 76 return object_list
- 77
- 78
- 79class FileManager(object):
- 80 """Set of actions to manipulate files in several ways."""
- 81
- 82 _logger = LoggingHandler(__name__).get_logger()
+ 71 if not paginator:
+ 72 list_response = s3_client.list_objects_v2(Bucket=bucket, Prefix=path)
+ 73 else:
+ 74 list_response = s3_client.list_objects_v2(
+ 75 Bucket=bucket,
+ 76 Prefix=path,
+ 77 ContinuationToken=paginator,
+ 78 )
+ 79
+ 80 if FileManagerAPIKeys.CONTENTS.value in list_response:
+ 81 for obj in list_response[FileManagerAPIKeys.CONTENTS.value]:
+ 82 object_list.append(obj[FileManagerAPIKeys.KEY.value])
83
- 84 def __init__(self, configs: dict):
- 85 """Construct FileManager algorithm instances.
- 86
- 87 Args:
- 88 configs: configurations for the FileManager algorithm.
- 89 """
- 90 self.configs = configs
- 91 self.function = self.configs["function"]
- 92
- 93 def get_function(self) -> None:
- 94 """Get a specific function to execute."""
- 95 available_functions = {
- 96 "delete_objects": self.delete_objects,
- 97 "copy_objects": self.copy_objects,
- 98 "request_restore": self.request_restore,
- 99 "check_restore_status": self.check_restore_status,
-100 "request_restore_to_destination_and_wait": (
-101 self.request_restore_to_destination_and_wait
-102 ),
-103 }
-104
-105 self._logger.info("Function being executed: {}".format(self.function))
-106 if self.function in available_functions.keys():
-107 func = available_functions[self.function]
-108 func()
-109 else:
-110 raise NotImplementedError(
-111 f"The requested function {self.function} is not implemented."
-112 )
-113
-114 def delete_objects(self) -> None:
-115 """Delete objects and 'directories' in s3.
-116
-117 If dry_run is set to True the function will print a dict with all the
-118 paths that would be deleted based on the given keys.
-119 """
-120 bucket = self.configs["bucket"]
-121 objects_paths = self.configs["object_paths"]
-122 dry_run = self.configs["dry_run"]
-123
-124 s3 = boto3.client("s3")
+ 84 if FileManagerAPIKeys.CONTINUATION.value in list_response:
+ 85 pagination = list_response[FileManagerAPIKeys.CONTINUATION.value]
+ 86 else:
+ 87 pagination = ""
+ 88
+ 89 return object_list, pagination
+ 90
+ 91
+ 92def _list_objects_recursively(bucket: str, path: str) -> list:
+ 93 """Recursively list all objects given a prefix in s3.
+ 94
+ 95 Args:
+ 96 bucket: name of bucket to perform the list.
+ 97 path: path to be used as a prefix.
+ 98
+ 99 Returns:
+100 A list of object names fetched recursively.
+101 """
+102 object_list = []
+103 more_objects = True
+104 paginator = ""
+105
+106 s3 = boto3.client("s3")
+107
+108 while more_objects:
+109 temp_list, paginator = _list_objects(s3, bucket, path, paginator)
+110
+111 object_list.extend(temp_list)
+112
+113 if not paginator:
+114 more_objects = False
+115
+116 return object_list
+117
+118
+119def _check_directory(bucket: str, path: str) -> bool:
+120 """Checks if the object is a 'directory' in s3.
+121
+122 Args:
+123 bucket: name of bucket to perform the check.
+124 path: path to be used as a prefix.
125
-126 if dry_run:
-127 response = _dry_run(bucket=bucket, object_paths=objects_paths)
-128
-129 self._logger.info("Paths that would be deleted:")
-130 else:
-131 objects_to_delete = []
-132 for path in objects_paths:
-133 for obj in _list_objects_recursively(bucket=bucket, path=path):
-134 objects_to_delete.append({FileManagerAPIKeys.KEY.value: obj})
-135
-136 response = s3.delete_objects(
-137 Bucket=bucket,
-138 Delete={FileManagerAPIKeys.OBJECTS.value: objects_to_delete},
-139 )
-140
-141 self._logger.info(response)
-142
-143 def copy_objects(self) -> None:
-144 """Copies objects and 'directories' in s3."""
-145 source_bucket = self.configs["bucket"]
-146 source_object = self.configs["source_object"]
-147 destination_bucket = self.configs["destination_bucket"]
-148 destination_object = self.configs["destination_object"]
-149 dry_run = self.configs["dry_run"]
-150
-151 FileManager._copy_objects(
-152 source_bucket=source_bucket,
-153 source_object=source_object,
-154 destination_bucket=destination_bucket,
-155 destination_object=destination_object,
-156 dry_run=dry_run,
-157 )
-158
-159 def request_restore(self) -> None:
-160 """Request the restore of archived data."""
-161 source_bucket = self.configs["bucket"]
-162 source_object = self.configs["source_object"]
-163 restore_expiration = self.configs["restore_expiration"]
-164 retrieval_tier = self.configs["retrieval_tier"]
-165 dry_run = self.configs["dry_run"]
-166
-167 ArchiveFileManager.request_restore(
-168 source_bucket,
-169 source_object,
-170 restore_expiration,
-171 retrieval_tier,
-172 dry_run,
-173 )
-174
-175 def check_restore_status(self) -> None:
-176 """Check the restore status of archived data."""
-177 source_bucket = self.configs["bucket"]
-178 source_object = self.configs["source_object"]
-179
-180 restore_status = ArchiveFileManager.check_restore_status(
-181 source_bucket, source_object
-182 )
+126 Returns:
+127 If path represents a 'directory'.
+128 """
+129 s3 = boto3.client("s3")
+130 objects, _ = _list_objects(s3, bucket, path)
+131 return len(objects) > 1
+132
+133
+134class FileManager(object):
+135 """Set of actions to manipulate files in several ways."""
+136
+137 _logger = LoggingHandler(__name__).get_logger()
+138
+139 def __init__(self, configs: dict):
+140 """Construct FileManager algorithm instances.
+141
+142 Args:
+143 configs: configurations for the FileManager algorithm.
+144 """
+145 self.configs = configs
+146 self.function = self.configs["function"]
+147
+148 def get_function(self) -> None:
+149 """Get a specific function to execute."""
+150 available_functions = {
+151 "delete_objects": self.delete_objects,
+152 "copy_objects": self.copy_objects,
+153 "request_restore": self.request_restore,
+154 "check_restore_status": self.check_restore_status,
+155 "request_restore_to_destination_and_wait": (
+156 self.request_restore_to_destination_and_wait
+157 ),
+158 }
+159
+160 self._logger.info("Function being executed: {}".format(self.function))
+161 if self.function in available_functions.keys():
+162 func = available_functions[self.function]
+163 func()
+164 else:
+165 raise NotImplementedError(
+166 f"The requested function {self.function} is not implemented."
+167 )
+168
+169 def _delete_objects(self, bucket: str, objects_paths: list) -> None:
+170 """Delete objects recursively in s3.
+171
+172 Params:
+173 bucket: name of bucket to perform the delete operation.
+174 objects_paths: objects to be deleted.
+175 """
+176 s3 = boto3.client("s3")
+177
+178 for path in objects_paths:
+179 if _check_directory(bucket, path):
+180 path = _process_directory_path(path)
+181 else:
+182 path = path.strip()
183
-184 self._logger.info(
-185 f"""
-186 Restore status:
-187 - Not Started: {restore_status.get('not_started_objects')}
-188 - Ongoing: {restore_status.get('ongoing_objects')}
-189 - Restored: {restore_status.get('restored_objects')}
-190 Total objects in this restore process: {restore_status.get('total_objects')}
-191 """
-192 )
-193
-194 def request_restore_to_destination_and_wait(self) -> None:
-195 """Request and wait for the restore to complete, polling the restore status.
-196
-197 After the restore is done, copy the restored files to destination
-198 """
-199 source_bucket = self.configs["bucket"]
-200 source_object = self.configs["source_object"]
-201 destination_bucket = self.configs["destination_bucket"]
-202 destination_object = self.configs["destination_object"]
-203 restore_expiration = self.configs["restore_expiration"]
-204 retrieval_tier = self.configs["retrieval_tier"]
-205 dry_run = self.configs["dry_run"]
-206
-207 ArchiveFileManager.request_restore_and_wait(
-208 source_bucket=source_bucket,
-209 source_object=source_object,
-210 restore_expiration=restore_expiration,
-211 retrieval_tier=retrieval_tier,
-212 dry_run=dry_run,
-213 )
+184 more_objects = True
+185 paginator = ""
+186 objects_to_delete = []
+187
+188 while more_objects:
+189 objects_found, paginator = _list_objects(
+190 s3_client=s3, bucket=bucket, path=path, paginator=paginator
+191 )
+192 for obj in objects_found:
+193 objects_to_delete.append({FileManagerAPIKeys.KEY.value: obj})
+194
+195 if not paginator:
+196 more_objects = False
+197
+198 response = s3.delete_objects(
+199 Bucket=bucket,
+200 Delete={FileManagerAPIKeys.OBJECTS.value: objects_to_delete},
+201 )
+202 self._logger.info(response)
+203 objects_to_delete = []
+204
+205 def delete_objects(self) -> None:
+206 """Delete objects and 'directories' in s3.
+207
+208 If dry_run is set to True the function will print a dict with all the
+209 paths that would be deleted based on the given keys.
+210 """
+211 bucket = self.configs["bucket"]
+212 objects_paths = self.configs["object_paths"]
+213 dry_run = self.configs["dry_run"]
214
-215 FileManager._logger.info(
-216 f"Restoration complete for {source_bucket} and {source_object}"
-217 )
-218 FileManager._logger.info(
-219 f"Starting to copy data from {source_bucket}/{source_object} to "
-220 f"{destination_bucket}/{destination_object}"
-221 )
-222 FileManager._copy_objects(
-223 source_bucket=source_bucket,
-224 source_object=source_object,
-225 destination_bucket=destination_bucket,
-226 destination_object=destination_object,
-227 dry_run=dry_run,
-228 )
-229 FileManager._logger.info(
-230 f"Finished copying data, data should be available on {destination_bucket}/"
-231 f"{destination_object}"
-232 )
-233
-234 @staticmethod
-235 def _copy_objects(
-236 source_bucket: str,
-237 source_object: str,
-238 destination_bucket: str,
-239 destination_object: str,
-240 dry_run: bool,
-241 ) -> None:
-242 """Copies objects and 'directories' in s3.
-243
-244 Args:
-245 source_bucket: name of bucket to perform the copy.
-246 source_object: object/folder to be copied.
-247 destination_bucket: name of the target bucket to copy.
-248 destination_object: target object/folder to copy.
-249 dry_run: if dry_run is set to True the function will print a dict with
-250 all the paths that would be deleted based on the given keys.
-251 """
-252 s3 = boto3.client("s3")
-253
-254 if dry_run:
-255 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
-256
-257 FileManager._logger.info("Paths that would be copied:")
-258 FileManager._logger.info(response)
-259 else:
-260 copy_object = _list_objects_recursively(
-261 bucket=source_bucket, path=source_object
-262 )
+215 if dry_run:
+216 response = _dry_run(bucket=bucket, object_paths=objects_paths)
+217
+218 self._logger.info("Paths that would be deleted:")
+219 self._logger.info(response)
+220 else:
+221 self._delete_objects(bucket, objects_paths)
+222
+223 def copy_objects(self) -> None:
+224 """Copies objects and 'directories' in s3."""
+225 source_bucket = self.configs["bucket"]
+226 source_object = self.configs["source_object"]
+227 destination_bucket = self.configs["destination_bucket"]
+228 destination_object = self.configs["destination_object"]
+229 dry_run = self.configs["dry_run"]
+230
+231 FileManager._copy_objects(
+232 source_bucket=source_bucket,
+233 source_object=source_object,
+234 destination_bucket=destination_bucket,
+235 destination_object=destination_object,
+236 dry_run=dry_run,
+237 )
+238
+239 def request_restore(self) -> None:
+240 """Request the restore of archived data."""
+241 source_bucket = self.configs["bucket"]
+242 source_object = self.configs["source_object"]
+243 restore_expiration = self.configs["restore_expiration"]
+244 retrieval_tier = self.configs["retrieval_tier"]
+245 dry_run = self.configs["dry_run"]
+246
+247 ArchiveFileManager.request_restore(
+248 source_bucket,
+249 source_object,
+250 restore_expiration,
+251 retrieval_tier,
+252 dry_run,
+253 )
+254
+255 def check_restore_status(self) -> None:
+256 """Check the restore status of archived data."""
+257 source_bucket = self.configs["bucket"]
+258 source_object = self.configs["source_object"]
+259
+260 restore_status = ArchiveFileManager.check_restore_status(
+261 source_bucket, source_object
+262 )
263
-264 if len(copy_object) == 1:
-265 FileManager._logger.info(f"Copying obj: {source_object}")
-266
-267 response = s3.copy_object(
-268 Bucket=destination_bucket,
-269 CopySource={
-270 FileManagerAPIKeys.BUCKET.value: source_bucket,
-271 FileManagerAPIKeys.KEY.value: source_object,
-272 },
-273 Key=f"""{destination_object}/{copy_object[0].split("/")[-1]}""",
-274 )
-275 FileManager._logger.info(response)
-276 else:
-277 for obj in copy_object:
-278 FileManager._logger.info(f"Copying obj: {obj}")
-279
-280 final_path = obj.replace(source_object, "")
-281
-282 response = s3.copy_object(
-283 Bucket=destination_bucket,
-284 CopySource={
-285 FileManagerAPIKeys.BUCKET.value: source_bucket,
-286 FileManagerAPIKeys.KEY.value: obj,
-287 },
-288 Key=f"{destination_object}{final_path}",
-289 )
-290 FileManager._logger.info(response)
-291
-292
-293class ArchiveFileManager(object):
-294 """Set of actions to restore archives."""
-295
-296 _logger = LoggingHandler(__name__).get_logger()
-297
-298 @staticmethod
-299 def _get_archived_object(bucket: str, object_key: str) -> Optional[Any]:
-300 """Get the archived object if it's an object.
-301
-302 Args:
-303 bucket: name of bucket to check get the object.
-304 object_key: object to get.
-305
-306 Returns:
-307 S3 Object if it's an archived object, otherwise None.
-308 """
-309 s3 = boto3.resource("s3")
-310 object_to_restore = s3.Object(bucket, object_key)
-311
-312 if (
-313 object_to_restore.storage_class is not None
-314 and object_to_restore.storage_class in ARCHIVE_STORAGE_CLASS
-315 ):
-316 return object_to_restore
-317 else:
-318 return None
-319
-320 @staticmethod
-321 def _check_object_restore_status(
-322 bucket: str, object_key: str
-323 ) -> Optional[RestoreStatus]:
-324 """Check the restore status of the archive.
-325
-326 Args:
-327 bucket: name of bucket to check the restore status.
-328 object_key: object to check the restore status.
-329
-330 Returns:
-331 The restore status represented by an enum, possible values are:
-332 NOT_STARTED, ONGOING or RESTORED
-333 """
-334 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
-335
-336 if archived_object is None:
-337 status = None
-338 elif archived_object.restore is None:
-339 status = RestoreStatus.NOT_STARTED
-340 elif 'ongoing-request="true"' in archived_object.restore:
-341 status = RestoreStatus.ONGOING
-342 else:
-343 status = RestoreStatus.RESTORED
+264 self._logger.info(
+265 f"""
+266 Restore status:
+267 - Not Started: {restore_status.get('not_started_objects')}
+268 - Ongoing: {restore_status.get('ongoing_objects')}
+269 - Restored: {restore_status.get('restored_objects')}
+270 Total objects in this restore process: {restore_status.get('total_objects')}
+271 """
+272 )
+273
+274 def request_restore_to_destination_and_wait(self) -> None:
+275 """Request and wait for the restore to complete, polling the restore status.
+276
+277 After the restore is done, copy the restored files to destination
+278 """
+279 source_bucket = self.configs["bucket"]
+280 source_object = self.configs["source_object"]
+281 destination_bucket = self.configs["destination_bucket"]
+282 destination_object = self.configs["destination_object"]
+283 restore_expiration = self.configs["restore_expiration"]
+284 retrieval_tier = self.configs["retrieval_tier"]
+285 dry_run = self.configs["dry_run"]
+286
+287 ArchiveFileManager.request_restore_and_wait(
+288 source_bucket=source_bucket,
+289 source_object=source_object,
+290 restore_expiration=restore_expiration,
+291 retrieval_tier=retrieval_tier,
+292 dry_run=dry_run,
+293 )
+294
+295 FileManager._logger.info(
+296 f"Restoration complete for {source_bucket} and {source_object}"
+297 )
+298 FileManager._logger.info(
+299 f"Starting to copy data from {source_bucket}/{source_object} to "
+300 f"{destination_bucket}/{destination_object}"
+301 )
+302 FileManager._copy_objects(
+303 source_bucket=source_bucket,
+304 source_object=source_object,
+305 destination_bucket=destination_bucket,
+306 destination_object=destination_object,
+307 dry_run=dry_run,
+308 )
+309 FileManager._logger.info(
+310 f"Finished copying data, data should be available on {destination_bucket}/"
+311 f"{destination_object}"
+312 )
+313
+314 @staticmethod
+315 def _copy_objects(
+316 source_bucket: str,
+317 source_object: str,
+318 destination_bucket: str,
+319 destination_object: str,
+320 dry_run: bool,
+321 ) -> None:
+322 """Copies objects and 'directories' in s3.
+323
+324 Args:
+325 source_bucket: name of bucket to perform the copy.
+326 source_object: object/folder to be copied.
+327 destination_bucket: name of the target bucket to copy.
+328 destination_object: target object/folder to copy.
+329 dry_run: if dry_run is set to True the function will print a dict with
+330 all the paths that would be deleted based on the given keys.
+331 """
+332 s3 = boto3.client("s3")
+333
+334 if dry_run:
+335 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
+336
+337 FileManager._logger.info("Paths that would be copied:")
+338 FileManager._logger.info(response)
+339 else:
+340 original_object_name = source_object.split("/")[-1]
+341
+342 if _check_directory(source_bucket, source_object):
+343 source_object = _process_directory_path(source_object)
344
-345 return status
-346
-347 @staticmethod
-348 def check_restore_status(source_bucket: str, source_object: str) -> dict:
-349 """Check the restore status of archived data.
-350
-351 Args:
-352 source_bucket: name of bucket to check the restore status.
-353 source_object: object to check the restore status.
-354
-355 Returns:
-356 A dict containing the amount of objects in each status.
-357 """
-358 not_started_objects = 0
-359 ongoing_objects = 0
-360 restored_objects = 0
-361 total_objects = 0
-362
-363 objects_to_restore = _list_objects_recursively(
-364 bucket=source_bucket, path=source_object
-365 )
-366
-367 for obj in objects_to_restore:
-368 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
-369
-370 restore_status = ArchiveFileManager._check_object_restore_status(
-371 source_bucket, obj
-372 )
-373 if not restore_status:
-374 ArchiveFileManager._logger.warning(
-375 f"Restore status not found for {source_bucket}/{obj}"
-376 )
-377 else:
-378 total_objects += 1
+345 copy_object = _list_objects_recursively(
+346 bucket=source_bucket, path=source_object
+347 )
+348
+349 for obj in copy_object:
+350 FileManager._logger.info(f"Copying obj: {obj}")
+351
+352 final_path = obj.replace(source_object, "")
+353
+354 response = s3.copy_object(
+355 Bucket=destination_bucket,
+356 CopySource={
+357 FileManagerAPIKeys.BUCKET.value: source_bucket,
+358 FileManagerAPIKeys.KEY.value: obj,
+359 },
+360 Key=f"{destination_object}/{original_object_name}/{final_path}",
+361 )
+362 FileManager._logger.info(response)
+363 else:
+364 FileManager._logger.info(f"Copying obj: {source_object}")
+365
+366 response = s3.copy_object(
+367 Bucket=destination_bucket,
+368 CopySource={
+369 FileManagerAPIKeys.BUCKET.value: source_bucket,
+370 FileManagerAPIKeys.KEY.value: source_object,
+371 },
+372 Key=f"""{destination_object}/{original_object_name}""",
+373 )
+374 FileManager._logger.info(response)
+375
+376
+377class ArchiveFileManager(object):
+378 """Set of actions to restore archives."""
379
-380 if RestoreStatus.NOT_STARTED == restore_status:
-381 not_started_objects += 1
-382 elif RestoreStatus.ONGOING == restore_status:
-383 ongoing_objects += 1
-384 else:
-385 restored_objects += 1
-386
-387 ArchiveFileManager._logger.info(
-388 f"{obj} restore status is {restore_status.value}"
-389 )
-390
-391 return {
-392 "total_objects": total_objects,
-393 "not_started_objects": not_started_objects,
-394 "ongoing_objects": ongoing_objects,
-395 "restored_objects": restored_objects,
-396 }
-397
-398 @staticmethod
-399 def _request_restore_object(
-400 bucket: str, object_key: str, expiration: int, retrieval_tier: str
-401 ) -> None:
-402 """Request a restore of the archive.
+380 _logger = LoggingHandler(__name__).get_logger()
+381
+382 @staticmethod
+383 def _get_archived_object(bucket: str, object_key: str) -> Optional[Any]:
+384 """Get the archived object if it's an object.
+385
+386 Args:
+387 bucket: name of bucket to check get the object.
+388 object_key: object to get.
+389
+390 Returns:
+391 S3 Object if it's an archived object, otherwise None.
+392 """
+393 s3 = boto3.resource("s3")
+394 object_to_restore = s3.Object(bucket, object_key)
+395
+396 if (
+397 object_to_restore.storage_class is not None
+398 and object_to_restore.storage_class in ARCHIVE_STORAGE_CLASS
+399 ):
+400 return object_to_restore
+401 else:
+402 return None
403
-404 Args:
-405 bucket: name of bucket to perform the restore.
-406 object_key: object to be restored.
-407 expiration: restore expiration.
-408 retrieval_tier: type of restore, possible values are:
-409 Bulk, Standard or Expedited.
-410 """
-411 if not RestoreType.exists(retrieval_tier):
-412 raise RestoreTypeNotFoundException(
-413 f"Restore type {retrieval_tier} not supported."
-414 )
-415
-416 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
-417
-418 if archived_object and archived_object.restore is None:
-419 ArchiveFileManager._logger.info(f"Restoring archive {bucket}/{object_key}.")
-420 archived_object.restore_object(
-421 RestoreRequest={
-422 "Days": expiration,
-423 "GlacierJobParameters": {"Tier": retrieval_tier},
-424 }
-425 )
+404 @staticmethod
+405 def _check_object_restore_status(
+406 bucket: str, object_key: str
+407 ) -> Optional[RestoreStatus]:
+408 """Check the restore status of the archive.
+409
+410 Args:
+411 bucket: name of bucket to check the restore status.
+412 object_key: object to check the restore status.
+413
+414 Returns:
+415 The restore status represented by an enum, possible values are:
+416 NOT_STARTED, ONGOING or RESTORED
+417 """
+418 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
+419
+420 if archived_object is None:
+421 status = None
+422 elif archived_object.restore is None:
+423 status = RestoreStatus.NOT_STARTED
+424 elif 'ongoing-request="true"' in archived_object.restore:
+425 status = RestoreStatus.ONGOING
426 else:
-427 ArchiveFileManager._logger.info(
-428 f"Restore request for {bucket}/{object_key} not performed."
-429 )
+427 status = RestoreStatus.RESTORED
+428
+429 return status
430
431 @staticmethod
-432 def request_restore(
-433 source_bucket: str,
-434 source_object: str,
-435 restore_expiration: int,
-436 retrieval_tier: str,
-437 dry_run: bool,
-438 ) -> None:
-439 """Request the restore of archived data.
-440
-441 Args:
-442 source_bucket: name of bucket to perform the restore.
-443 source_object: object to be restored.
-444 restore_expiration: restore expiration in days.
-445 retrieval_tier: type of restore, possible values are:
-446 Bulk, Standard or Expedited.
-447 dry_run: if dry_run is set to True the function will print a dict with
-448 all the paths that would be deleted based on the given keys.
-449 """
-450 if dry_run:
-451 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
-452
-453 ArchiveFileManager._logger.info("Paths that would be restored:")
-454 ArchiveFileManager._logger.info(response)
-455 else:
-456 objects_to_restore = _list_objects_recursively(
-457 bucket=source_bucket, path=source_object
-458 )
-459
-460 for obj in objects_to_restore:
-461 ArchiveFileManager._request_restore_object(
-462 source_bucket,
-463 obj,
-464 restore_expiration,
-465 retrieval_tier,
-466 )
-467
-468 @staticmethod
-469 def request_restore_and_wait(
-470 source_bucket: str,
-471 source_object: str,
-472 restore_expiration: int,
-473 retrieval_tier: str,
-474 dry_run: bool,
-475 ) -> None:
-476 """Request and wait for the restore to complete, polling the restore status.
+432 def check_restore_status(source_bucket: str, source_object: str) -> dict:
+433 """Check the restore status of archived data.
+434
+435 Args:
+436 source_bucket: name of bucket to check the restore status.
+437 source_object: object to check the restore status.
+438
+439 Returns:
+440 A dict containing the amount of objects in each status.
+441 """
+442 not_started_objects = 0
+443 ongoing_objects = 0
+444 restored_objects = 0
+445 total_objects = 0
+446
+447 if _check_directory(source_bucket, source_object):
+448 source_object = _process_directory_path(source_object)
+449
+450 objects_to_restore = _list_objects_recursively(
+451 bucket=source_bucket, path=source_object
+452 )
+453
+454 for obj in objects_to_restore:
+455 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
+456
+457 restore_status = ArchiveFileManager._check_object_restore_status(
+458 source_bucket, obj
+459 )
+460 if not restore_status:
+461 ArchiveFileManager._logger.warning(
+462 f"Restore status not found for {source_bucket}/{obj}"
+463 )
+464 else:
+465 total_objects += 1
+466
+467 if RestoreStatus.NOT_STARTED == restore_status:
+468 not_started_objects += 1
+469 elif RestoreStatus.ONGOING == restore_status:
+470 ongoing_objects += 1
+471 else:
+472 restored_objects += 1
+473
+474 ArchiveFileManager._logger.info(
+475 f"{obj} restore status is {restore_status.value}"
+476 )
477
-478 Args:
-479 source_bucket: name of bucket to perform the restore.
-480 source_object: object to be restored.
-481 restore_expiration: restore expiration in days.
-482 retrieval_tier: type of restore, possible values are:
-483 Bulk, Standard or Expedited.
-484 dry_run: if dry_run is set to True the function will print a dict with
-485 all the paths that would be deleted based on the given keys.
-486 """
-487 if retrieval_tier != RestoreType.EXPEDITED.value:
-488 ArchiveFileManager._logger.error(
-489 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-490 "kind of restore should be used just with `Expedited` retrieval tier "
-491 "to save cluster costs."
-492 )
-493 raise ValueError(
-494 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-495 "kind of restore should be used just with `Expedited` retrieval tier "
-496 "to save cluster costs."
-497 )
-498
-499 ArchiveFileManager.request_restore(
-500 source_bucket=source_bucket,
-501 source_object=source_object,
-502 restore_expiration=restore_expiration,
-503 retrieval_tier=retrieval_tier,
-504 dry_run=dry_run,
-505 )
-506 restore_status = ArchiveFileManager.check_restore_status(
-507 source_bucket, source_object
-508 )
-509 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
-510
-511 if not dry_run:
-512 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
-513 wait_time = 300
-514 while restore_status.get("total_objects") > restore_status.get(
-515 "restored_objects"
-516 ):
-517 ArchiveFileManager._logger.info(
-518 "Not all objects have been restored yet, checking the status again "
-519 f"in {wait_time} seconds."
-520 )
-521 time.sleep(wait_time)
-522 wait_time = 30
-523 restore_status = ArchiveFileManager.check_restore_status(
-524 source_bucket, source_object
-525 )
-526 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+478 return {
+479 "total_objects": total_objects,
+480 "not_started_objects": not_started_objects,
+481 "ongoing_objects": ongoing_objects,
+482 "restored_objects": restored_objects,
+483 }
+484
+485 @staticmethod
+486 def _request_restore_object(
+487 bucket: str, object_key: str, expiration: int, retrieval_tier: str
+488 ) -> None:
+489 """Request a restore of the archive.
+490
+491 Args:
+492 bucket: name of bucket to perform the restore.
+493 object_key: object to be restored.
+494 expiration: restore expiration.
+495 retrieval_tier: type of restore, possible values are:
+496 Bulk, Standard or Expedited.
+497 """
+498 if not RestoreType.exists(retrieval_tier):
+499 raise RestoreTypeNotFoundException(
+500 f"Restore type {retrieval_tier} not supported."
+501 )
+502
+503 if _check_directory(bucket, object_key):
+504 object_key = _process_directory_path(object_key)
+505
+506 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
+507
+508 if archived_object and archived_object.restore is None:
+509 ArchiveFileManager._logger.info(f"Restoring archive {bucket}/{object_key}.")
+510 archived_object.restore_object(
+511 RestoreRequest={
+512 "Days": expiration,
+513 "GlacierJobParameters": {"Tier": retrieval_tier},
+514 }
+515 )
+516 else:
+517 ArchiveFileManager._logger.info(
+518 f"Restore request for {bucket}/{object_key} not performed."
+519 )
+520
+521 @staticmethod
+522 def request_restore(
+523 source_bucket: str,
+524 source_object: str,
+525 restore_expiration: int,
+526 retrieval_tier: str,
+527 dry_run: bool,
+528 ) -> None:
+529 """Request the restore of archived data.
+530
+531 Args:
+532 source_bucket: name of bucket to perform the restore.
+533 source_object: object to be restored.
+534 restore_expiration: restore expiration in days.
+535 retrieval_tier: type of restore, possible values are:
+536 Bulk, Standard or Expedited.
+537 dry_run: if dry_run is set to True the function will print a dict with
+538 all the paths that would be deleted based on the given keys.
+539 """
+540 if _check_directory(source_bucket, source_object):
+541 source_object = _process_directory_path(source_object)
+542
+543 if dry_run:
+544 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
+545
+546 ArchiveFileManager._logger.info("Paths that would be restored:")
+547 ArchiveFileManager._logger.info(response)
+548 else:
+549 objects_to_restore = _list_objects_recursively(
+550 bucket=source_bucket, path=source_object
+551 )
+552
+553 for obj in objects_to_restore:
+554 ArchiveFileManager._request_restore_object(
+555 source_bucket,
+556 obj,
+557 restore_expiration,
+558 retrieval_tier,
+559 )
+560
+561 @staticmethod
+562 def request_restore_and_wait(
+563 source_bucket: str,
+564 source_object: str,
+565 restore_expiration: int,
+566 retrieval_tier: str,
+567 dry_run: bool,
+568 ) -> None:
+569 """Request and wait for the restore to complete, polling the restore status.
+570
+571 Args:
+572 source_bucket: name of bucket to perform the restore.
+573 source_object: object to be restored.
+574 restore_expiration: restore expiration in days.
+575 retrieval_tier: type of restore, possible values are:
+576 Bulk, Standard or Expedited.
+577 dry_run: if dry_run is set to True the function will print a dict with
+578 all the paths that would be deleted based on the given keys.
+579 """
+580 if retrieval_tier != RestoreType.EXPEDITED.value:
+581 ArchiveFileManager._logger.error(
+582 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+583 "kind of restore should be used just with `Expedited` retrieval tier "
+584 "to save cluster costs."
+585 )
+586 raise ValueError(
+587 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+588 "kind of restore should be used just with `Expedited` retrieval tier "
+589 "to save cluster costs."
+590 )
+591
+592 ArchiveFileManager.request_restore(
+593 source_bucket=source_bucket,
+594 source_object=source_object,
+595 restore_expiration=restore_expiration,
+596 retrieval_tier=retrieval_tier,
+597 dry_run=dry_run,
+598 )
+599 restore_status = ArchiveFileManager.check_restore_status(
+600 source_bucket, source_object
+601 )
+602 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+603
+604 if not dry_run:
+605 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
+606 wait_time = 300
+607 while restore_status.get("total_objects") > restore_status.get(
+608 "restored_objects"
+609 ):
+610 ArchiveFileManager._logger.info(
+611 "Not all objects have been restored yet, checking the status again "
+612 f"in {wait_time} seconds."
+613 )
+614 time.sleep(wait_time)
+615 wait_time = 30
+616 restore_status = ArchiveFileManager.check_restore_status(
+617 source_bucket, source_object
+618 )
+619 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
80class FileManager(object):
- 81 """Set of actions to manipulate files in several ways."""
- 82
- 83 _logger = LoggingHandler(__name__).get_logger()
- 84
- 85 def __init__(self, configs: dict):
- 86 """Construct FileManager algorithm instances.
- 87
- 88 Args:
- 89 configs: configurations for the FileManager algorithm.
- 90 """
- 91 self.configs = configs
- 92 self.function = self.configs["function"]
- 93
- 94 def get_function(self) -> None:
- 95 """Get a specific function to execute."""
- 96 available_functions = {
- 97 "delete_objects": self.delete_objects,
- 98 "copy_objects": self.copy_objects,
- 99 "request_restore": self.request_restore,
-100 "check_restore_status": self.check_restore_status,
-101 "request_restore_to_destination_and_wait": (
-102 self.request_restore_to_destination_and_wait
-103 ),
-104 }
-105
-106 self._logger.info("Function being executed: {}".format(self.function))
-107 if self.function in available_functions.keys():
-108 func = available_functions[self.function]
-109 func()
-110 else:
-111 raise NotImplementedError(
-112 f"The requested function {self.function} is not implemented."
-113 )
-114
-115 def delete_objects(self) -> None:
-116 """Delete objects and 'directories' in s3.
-117
-118 If dry_run is set to True the function will print a dict with all the
-119 paths that would be deleted based on the given keys.
-120 """
-121 bucket = self.configs["bucket"]
-122 objects_paths = self.configs["object_paths"]
-123 dry_run = self.configs["dry_run"]
-124
-125 s3 = boto3.client("s3")
-126
-127 if dry_run:
-128 response = _dry_run(bucket=bucket, object_paths=objects_paths)
-129
-130 self._logger.info("Paths that would be deleted:")
-131 else:
-132 objects_to_delete = []
-133 for path in objects_paths:
-134 for obj in _list_objects_recursively(bucket=bucket, path=path):
-135 objects_to_delete.append({FileManagerAPIKeys.KEY.value: obj})
-136
-137 response = s3.delete_objects(
-138 Bucket=bucket,
-139 Delete={FileManagerAPIKeys.OBJECTS.value: objects_to_delete},
-140 )
-141
-142 self._logger.info(response)
-143
-144 def copy_objects(self) -> None:
-145 """Copies objects and 'directories' in s3."""
-146 source_bucket = self.configs["bucket"]
-147 source_object = self.configs["source_object"]
-148 destination_bucket = self.configs["destination_bucket"]
-149 destination_object = self.configs["destination_object"]
-150 dry_run = self.configs["dry_run"]
-151
-152 FileManager._copy_objects(
-153 source_bucket=source_bucket,
-154 source_object=source_object,
-155 destination_bucket=destination_bucket,
-156 destination_object=destination_object,
-157 dry_run=dry_run,
-158 )
-159
-160 def request_restore(self) -> None:
-161 """Request the restore of archived data."""
-162 source_bucket = self.configs["bucket"]
-163 source_object = self.configs["source_object"]
-164 restore_expiration = self.configs["restore_expiration"]
-165 retrieval_tier = self.configs["retrieval_tier"]
-166 dry_run = self.configs["dry_run"]
-167
-168 ArchiveFileManager.request_restore(
-169 source_bucket,
-170 source_object,
-171 restore_expiration,
-172 retrieval_tier,
-173 dry_run,
-174 )
-175
-176 def check_restore_status(self) -> None:
-177 """Check the restore status of archived data."""
-178 source_bucket = self.configs["bucket"]
-179 source_object = self.configs["source_object"]
-180
-181 restore_status = ArchiveFileManager.check_restore_status(
-182 source_bucket, source_object
-183 )
+ 135class FileManager(object):
+136 """Set of actions to manipulate files in several ways."""
+137
+138 _logger = LoggingHandler(__name__).get_logger()
+139
+140 def __init__(self, configs: dict):
+141 """Construct FileManager algorithm instances.
+142
+143 Args:
+144 configs: configurations for the FileManager algorithm.
+145 """
+146 self.configs = configs
+147 self.function = self.configs["function"]
+148
+149 def get_function(self) -> None:
+150 """Get a specific function to execute."""
+151 available_functions = {
+152 "delete_objects": self.delete_objects,
+153 "copy_objects": self.copy_objects,
+154 "request_restore": self.request_restore,
+155 "check_restore_status": self.check_restore_status,
+156 "request_restore_to_destination_and_wait": (
+157 self.request_restore_to_destination_and_wait
+158 ),
+159 }
+160
+161 self._logger.info("Function being executed: {}".format(self.function))
+162 if self.function in available_functions.keys():
+163 func = available_functions[self.function]
+164 func()
+165 else:
+166 raise NotImplementedError(
+167 f"The requested function {self.function} is not implemented."
+168 )
+169
+170 def _delete_objects(self, bucket: str, objects_paths: list) -> None:
+171 """Delete objects recursively in s3.
+172
+173 Params:
+174 bucket: name of bucket to perform the delete operation.
+175 objects_paths: objects to be deleted.
+176 """
+177 s3 = boto3.client("s3")
+178
+179 for path in objects_paths:
+180 if _check_directory(bucket, path):
+181 path = _process_directory_path(path)
+182 else:
+183 path = path.strip()
184
-185 self._logger.info(
-186 f"""
-187 Restore status:
-188 - Not Started: {restore_status.get('not_started_objects')}
-189 - Ongoing: {restore_status.get('ongoing_objects')}
-190 - Restored: {restore_status.get('restored_objects')}
-191 Total objects in this restore process: {restore_status.get('total_objects')}
-192 """
-193 )
-194
-195 def request_restore_to_destination_and_wait(self) -> None:
-196 """Request and wait for the restore to complete, polling the restore status.
-197
-198 After the restore is done, copy the restored files to destination
-199 """
-200 source_bucket = self.configs["bucket"]
-201 source_object = self.configs["source_object"]
-202 destination_bucket = self.configs["destination_bucket"]
-203 destination_object = self.configs["destination_object"]
-204 restore_expiration = self.configs["restore_expiration"]
-205 retrieval_tier = self.configs["retrieval_tier"]
-206 dry_run = self.configs["dry_run"]
-207
-208 ArchiveFileManager.request_restore_and_wait(
-209 source_bucket=source_bucket,
-210 source_object=source_object,
-211 restore_expiration=restore_expiration,
-212 retrieval_tier=retrieval_tier,
-213 dry_run=dry_run,
-214 )
+185 more_objects = True
+186 paginator = ""
+187 objects_to_delete = []
+188
+189 while more_objects:
+190 objects_found, paginator = _list_objects(
+191 s3_client=s3, bucket=bucket, path=path, paginator=paginator
+192 )
+193 for obj in objects_found:
+194 objects_to_delete.append({FileManagerAPIKeys.KEY.value: obj})
+195
+196 if not paginator:
+197 more_objects = False
+198
+199 response = s3.delete_objects(
+200 Bucket=bucket,
+201 Delete={FileManagerAPIKeys.OBJECTS.value: objects_to_delete},
+202 )
+203 self._logger.info(response)
+204 objects_to_delete = []
+205
+206 def delete_objects(self) -> None:
+207 """Delete objects and 'directories' in s3.
+208
+209 If dry_run is set to True the function will print a dict with all the
+210 paths that would be deleted based on the given keys.
+211 """
+212 bucket = self.configs["bucket"]
+213 objects_paths = self.configs["object_paths"]
+214 dry_run = self.configs["dry_run"]
215
-216 FileManager._logger.info(
-217 f"Restoration complete for {source_bucket} and {source_object}"
-218 )
-219 FileManager._logger.info(
-220 f"Starting to copy data from {source_bucket}/{source_object} to "
-221 f"{destination_bucket}/{destination_object}"
-222 )
-223 FileManager._copy_objects(
-224 source_bucket=source_bucket,
-225 source_object=source_object,
-226 destination_bucket=destination_bucket,
-227 destination_object=destination_object,
-228 dry_run=dry_run,
-229 )
-230 FileManager._logger.info(
-231 f"Finished copying data, data should be available on {destination_bucket}/"
-232 f"{destination_object}"
-233 )
-234
-235 @staticmethod
-236 def _copy_objects(
-237 source_bucket: str,
-238 source_object: str,
-239 destination_bucket: str,
-240 destination_object: str,
-241 dry_run: bool,
-242 ) -> None:
-243 """Copies objects and 'directories' in s3.
-244
-245 Args:
-246 source_bucket: name of bucket to perform the copy.
-247 source_object: object/folder to be copied.
-248 destination_bucket: name of the target bucket to copy.
-249 destination_object: target object/folder to copy.
-250 dry_run: if dry_run is set to True the function will print a dict with
-251 all the paths that would be deleted based on the given keys.
-252 """
-253 s3 = boto3.client("s3")
-254
-255 if dry_run:
-256 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
-257
-258 FileManager._logger.info("Paths that would be copied:")
-259 FileManager._logger.info(response)
-260 else:
-261 copy_object = _list_objects_recursively(
-262 bucket=source_bucket, path=source_object
-263 )
+216 if dry_run:
+217 response = _dry_run(bucket=bucket, object_paths=objects_paths)
+218
+219 self._logger.info("Paths that would be deleted:")
+220 self._logger.info(response)
+221 else:
+222 self._delete_objects(bucket, objects_paths)
+223
+224 def copy_objects(self) -> None:
+225 """Copies objects and 'directories' in s3."""
+226 source_bucket = self.configs["bucket"]
+227 source_object = self.configs["source_object"]
+228 destination_bucket = self.configs["destination_bucket"]
+229 destination_object = self.configs["destination_object"]
+230 dry_run = self.configs["dry_run"]
+231
+232 FileManager._copy_objects(
+233 source_bucket=source_bucket,
+234 source_object=source_object,
+235 destination_bucket=destination_bucket,
+236 destination_object=destination_object,
+237 dry_run=dry_run,
+238 )
+239
+240 def request_restore(self) -> None:
+241 """Request the restore of archived data."""
+242 source_bucket = self.configs["bucket"]
+243 source_object = self.configs["source_object"]
+244 restore_expiration = self.configs["restore_expiration"]
+245 retrieval_tier = self.configs["retrieval_tier"]
+246 dry_run = self.configs["dry_run"]
+247
+248 ArchiveFileManager.request_restore(
+249 source_bucket,
+250 source_object,
+251 restore_expiration,
+252 retrieval_tier,
+253 dry_run,
+254 )
+255
+256 def check_restore_status(self) -> None:
+257 """Check the restore status of archived data."""
+258 source_bucket = self.configs["bucket"]
+259 source_object = self.configs["source_object"]
+260
+261 restore_status = ArchiveFileManager.check_restore_status(
+262 source_bucket, source_object
+263 )
264
-265 if len(copy_object) == 1:
-266 FileManager._logger.info(f"Copying obj: {source_object}")
-267
-268 response = s3.copy_object(
-269 Bucket=destination_bucket,
-270 CopySource={
-271 FileManagerAPIKeys.BUCKET.value: source_bucket,
-272 FileManagerAPIKeys.KEY.value: source_object,
-273 },
-274 Key=f"""{destination_object}/{copy_object[0].split("/")[-1]}""",
-275 )
-276 FileManager._logger.info(response)
-277 else:
-278 for obj in copy_object:
-279 FileManager._logger.info(f"Copying obj: {obj}")
-280
-281 final_path = obj.replace(source_object, "")
-282
-283 response = s3.copy_object(
-284 Bucket=destination_bucket,
-285 CopySource={
-286 FileManagerAPIKeys.BUCKET.value: source_bucket,
-287 FileManagerAPIKeys.KEY.value: obj,
-288 },
-289 Key=f"{destination_object}{final_path}",
-290 )
-291 FileManager._logger.info(response)
+265 self._logger.info(
+266 f"""
+267 Restore status:
+268 - Not Started: {restore_status.get('not_started_objects')}
+269 - Ongoing: {restore_status.get('ongoing_objects')}
+270 - Restored: {restore_status.get('restored_objects')}
+271 Total objects in this restore process: {restore_status.get('total_objects')}
+272 """
+273 )
+274
+275 def request_restore_to_destination_and_wait(self) -> None:
+276 """Request and wait for the restore to complete, polling the restore status.
+277
+278 After the restore is done, copy the restored files to destination
+279 """
+280 source_bucket = self.configs["bucket"]
+281 source_object = self.configs["source_object"]
+282 destination_bucket = self.configs["destination_bucket"]
+283 destination_object = self.configs["destination_object"]
+284 restore_expiration = self.configs["restore_expiration"]
+285 retrieval_tier = self.configs["retrieval_tier"]
+286 dry_run = self.configs["dry_run"]
+287
+288 ArchiveFileManager.request_restore_and_wait(
+289 source_bucket=source_bucket,
+290 source_object=source_object,
+291 restore_expiration=restore_expiration,
+292 retrieval_tier=retrieval_tier,
+293 dry_run=dry_run,
+294 )
+295
+296 FileManager._logger.info(
+297 f"Restoration complete for {source_bucket} and {source_object}"
+298 )
+299 FileManager._logger.info(
+300 f"Starting to copy data from {source_bucket}/{source_object} to "
+301 f"{destination_bucket}/{destination_object}"
+302 )
+303 FileManager._copy_objects(
+304 source_bucket=source_bucket,
+305 source_object=source_object,
+306 destination_bucket=destination_bucket,
+307 destination_object=destination_object,
+308 dry_run=dry_run,
+309 )
+310 FileManager._logger.info(
+311 f"Finished copying data, data should be available on {destination_bucket}/"
+312 f"{destination_object}"
+313 )
+314
+315 @staticmethod
+316 def _copy_objects(
+317 source_bucket: str,
+318 source_object: str,
+319 destination_bucket: str,
+320 destination_object: str,
+321 dry_run: bool,
+322 ) -> None:
+323 """Copies objects and 'directories' in s3.
+324
+325 Args:
+326 source_bucket: name of bucket to perform the copy.
+327 source_object: object/folder to be copied.
+328 destination_bucket: name of the target bucket to copy.
+329 destination_object: target object/folder to copy.
+330 dry_run: if dry_run is set to True the function will print a dict with
+331 all the paths that would be deleted based on the given keys.
+332 """
+333 s3 = boto3.client("s3")
+334
+335 if dry_run:
+336 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
+337
+338 FileManager._logger.info("Paths that would be copied:")
+339 FileManager._logger.info(response)
+340 else:
+341 original_object_name = source_object.split("/")[-1]
+342
+343 if _check_directory(source_bucket, source_object):
+344 source_object = _process_directory_path(source_object)
+345
+346 copy_object = _list_objects_recursively(
+347 bucket=source_bucket, path=source_object
+348 )
+349
+350 for obj in copy_object:
+351 FileManager._logger.info(f"Copying obj: {obj}")
+352
+353 final_path = obj.replace(source_object, "")
+354
+355 response = s3.copy_object(
+356 Bucket=destination_bucket,
+357 CopySource={
+358 FileManagerAPIKeys.BUCKET.value: source_bucket,
+359 FileManagerAPIKeys.KEY.value: obj,
+360 },
+361 Key=f"{destination_object}/{original_object_name}/{final_path}",
+362 )
+363 FileManager._logger.info(response)
+364 else:
+365 FileManager._logger.info(f"Copying obj: {source_object}")
+366
+367 response = s3.copy_object(
+368 Bucket=destination_bucket,
+369 CopySource={
+370 FileManagerAPIKeys.BUCKET.value: source_bucket,
+371 FileManagerAPIKeys.KEY.value: source_object,
+372 },
+373 Key=f"""{destination_object}/{original_object_name}""",
+374 )
+375 FileManager._logger.info(response)
85 def __init__(self, configs: dict):
-86 """Construct FileManager algorithm instances.
-87
-88 Args:
-89 configs: configurations for the FileManager algorithm.
-90 """
-91 self.configs = configs
-92 self.function = self.configs["function"]
+ 140 def __init__(self, configs: dict):
+141 """Construct FileManager algorithm instances.
+142
+143 Args:
+144 configs: configurations for the FileManager algorithm.
+145 """
+146 self.configs = configs
+147 self.function = self.configs["function"]
94 def get_function(self) -> None:
- 95 """Get a specific function to execute."""
- 96 available_functions = {
- 97 "delete_objects": self.delete_objects,
- 98 "copy_objects": self.copy_objects,
- 99 "request_restore": self.request_restore,
-100 "check_restore_status": self.check_restore_status,
-101 "request_restore_to_destination_and_wait": (
-102 self.request_restore_to_destination_and_wait
-103 ),
-104 }
-105
-106 self._logger.info("Function being executed: {}".format(self.function))
-107 if self.function in available_functions.keys():
-108 func = available_functions[self.function]
-109 func()
-110 else:
-111 raise NotImplementedError(
-112 f"The requested function {self.function} is not implemented."
-113 )
+ 149 def get_function(self) -> None:
+150 """Get a specific function to execute."""
+151 available_functions = {
+152 "delete_objects": self.delete_objects,
+153 "copy_objects": self.copy_objects,
+154 "request_restore": self.request_restore,
+155 "check_restore_status": self.check_restore_status,
+156 "request_restore_to_destination_and_wait": (
+157 self.request_restore_to_destination_and_wait
+158 ),
+159 }
+160
+161 self._logger.info("Function being executed: {}".format(self.function))
+162 if self.function in available_functions.keys():
+163 func = available_functions[self.function]
+164 func()
+165 else:
+166 raise NotImplementedError(
+167 f"The requested function {self.function} is not implemented."
+168 )
115 def delete_objects(self) -> None:
-116 """Delete objects and 'directories' in s3.
-117
-118 If dry_run is set to True the function will print a dict with all the
-119 paths that would be deleted based on the given keys.
-120 """
-121 bucket = self.configs["bucket"]
-122 objects_paths = self.configs["object_paths"]
-123 dry_run = self.configs["dry_run"]
-124
-125 s3 = boto3.client("s3")
-126
-127 if dry_run:
-128 response = _dry_run(bucket=bucket, object_paths=objects_paths)
-129
-130 self._logger.info("Paths that would be deleted:")
-131 else:
-132 objects_to_delete = []
-133 for path in objects_paths:
-134 for obj in _list_objects_recursively(bucket=bucket, path=path):
-135 objects_to_delete.append({FileManagerAPIKeys.KEY.value: obj})
-136
-137 response = s3.delete_objects(
-138 Bucket=bucket,
-139 Delete={FileManagerAPIKeys.OBJECTS.value: objects_to_delete},
-140 )
-141
-142 self._logger.info(response)
+ 206 def delete_objects(self) -> None:
+207 """Delete objects and 'directories' in s3.
+208
+209 If dry_run is set to True the function will print a dict with all the
+210 paths that would be deleted based on the given keys.
+211 """
+212 bucket = self.configs["bucket"]
+213 objects_paths = self.configs["object_paths"]
+214 dry_run = self.configs["dry_run"]
+215
+216 if dry_run:
+217 response = _dry_run(bucket=bucket, object_paths=objects_paths)
+218
+219 self._logger.info("Paths that would be deleted:")
+220 self._logger.info(response)
+221 else:
+222 self._delete_objects(bucket, objects_paths)
144 def copy_objects(self) -> None:
-145 """Copies objects and 'directories' in s3."""
-146 source_bucket = self.configs["bucket"]
-147 source_object = self.configs["source_object"]
-148 destination_bucket = self.configs["destination_bucket"]
-149 destination_object = self.configs["destination_object"]
-150 dry_run = self.configs["dry_run"]
-151
-152 FileManager._copy_objects(
-153 source_bucket=source_bucket,
-154 source_object=source_object,
-155 destination_bucket=destination_bucket,
-156 destination_object=destination_object,
-157 dry_run=dry_run,
-158 )
+ 224 def copy_objects(self) -> None:
+225 """Copies objects and 'directories' in s3."""
+226 source_bucket = self.configs["bucket"]
+227 source_object = self.configs["source_object"]
+228 destination_bucket = self.configs["destination_bucket"]
+229 destination_object = self.configs["destination_object"]
+230 dry_run = self.configs["dry_run"]
+231
+232 FileManager._copy_objects(
+233 source_bucket=source_bucket,
+234 source_object=source_object,
+235 destination_bucket=destination_bucket,
+236 destination_object=destination_object,
+237 dry_run=dry_run,
+238 )
160 def request_restore(self) -> None:
-161 """Request the restore of archived data."""
-162 source_bucket = self.configs["bucket"]
-163 source_object = self.configs["source_object"]
-164 restore_expiration = self.configs["restore_expiration"]
-165 retrieval_tier = self.configs["retrieval_tier"]
-166 dry_run = self.configs["dry_run"]
-167
-168 ArchiveFileManager.request_restore(
-169 source_bucket,
-170 source_object,
-171 restore_expiration,
-172 retrieval_tier,
-173 dry_run,
-174 )
+ 240 def request_restore(self) -> None:
+241 """Request the restore of archived data."""
+242 source_bucket = self.configs["bucket"]
+243 source_object = self.configs["source_object"]
+244 restore_expiration = self.configs["restore_expiration"]
+245 retrieval_tier = self.configs["retrieval_tier"]
+246 dry_run = self.configs["dry_run"]
+247
+248 ArchiveFileManager.request_restore(
+249 source_bucket,
+250 source_object,
+251 restore_expiration,
+252 retrieval_tier,
+253 dry_run,
+254 )
176 def check_restore_status(self) -> None:
-177 """Check the restore status of archived data."""
-178 source_bucket = self.configs["bucket"]
-179 source_object = self.configs["source_object"]
-180
-181 restore_status = ArchiveFileManager.check_restore_status(
-182 source_bucket, source_object
-183 )
-184
-185 self._logger.info(
-186 f"""
-187 Restore status:
-188 - Not Started: {restore_status.get('not_started_objects')}
-189 - Ongoing: {restore_status.get('ongoing_objects')}
-190 - Restored: {restore_status.get('restored_objects')}
-191 Total objects in this restore process: {restore_status.get('total_objects')}
-192 """
-193 )
+ 256 def check_restore_status(self) -> None:
+257 """Check the restore status of archived data."""
+258 source_bucket = self.configs["bucket"]
+259 source_object = self.configs["source_object"]
+260
+261 restore_status = ArchiveFileManager.check_restore_status(
+262 source_bucket, source_object
+263 )
+264
+265 self._logger.info(
+266 f"""
+267 Restore status:
+268 - Not Started: {restore_status.get('not_started_objects')}
+269 - Ongoing: {restore_status.get('ongoing_objects')}
+270 - Restored: {restore_status.get('restored_objects')}
+271 Total objects in this restore process: {restore_status.get('total_objects')}
+272 """
+273 )
195 def request_restore_to_destination_and_wait(self) -> None:
-196 """Request and wait for the restore to complete, polling the restore status.
-197
-198 After the restore is done, copy the restored files to destination
-199 """
-200 source_bucket = self.configs["bucket"]
-201 source_object = self.configs["source_object"]
-202 destination_bucket = self.configs["destination_bucket"]
-203 destination_object = self.configs["destination_object"]
-204 restore_expiration = self.configs["restore_expiration"]
-205 retrieval_tier = self.configs["retrieval_tier"]
-206 dry_run = self.configs["dry_run"]
-207
-208 ArchiveFileManager.request_restore_and_wait(
-209 source_bucket=source_bucket,
-210 source_object=source_object,
-211 restore_expiration=restore_expiration,
-212 retrieval_tier=retrieval_tier,
-213 dry_run=dry_run,
-214 )
-215
-216 FileManager._logger.info(
-217 f"Restoration complete for {source_bucket} and {source_object}"
-218 )
-219 FileManager._logger.info(
-220 f"Starting to copy data from {source_bucket}/{source_object} to "
-221 f"{destination_bucket}/{destination_object}"
-222 )
-223 FileManager._copy_objects(
-224 source_bucket=source_bucket,
-225 source_object=source_object,
-226 destination_bucket=destination_bucket,
-227 destination_object=destination_object,
-228 dry_run=dry_run,
-229 )
-230 FileManager._logger.info(
-231 f"Finished copying data, data should be available on {destination_bucket}/"
-232 f"{destination_object}"
-233 )
+ 275 def request_restore_to_destination_and_wait(self) -> None:
+276 """Request and wait for the restore to complete, polling the restore status.
+277
+278 After the restore is done, copy the restored files to destination
+279 """
+280 source_bucket = self.configs["bucket"]
+281 source_object = self.configs["source_object"]
+282 destination_bucket = self.configs["destination_bucket"]
+283 destination_object = self.configs["destination_object"]
+284 restore_expiration = self.configs["restore_expiration"]
+285 retrieval_tier = self.configs["retrieval_tier"]
+286 dry_run = self.configs["dry_run"]
+287
+288 ArchiveFileManager.request_restore_and_wait(
+289 source_bucket=source_bucket,
+290 source_object=source_object,
+291 restore_expiration=restore_expiration,
+292 retrieval_tier=retrieval_tier,
+293 dry_run=dry_run,
+294 )
+295
+296 FileManager._logger.info(
+297 f"Restoration complete for {source_bucket} and {source_object}"
+298 )
+299 FileManager._logger.info(
+300 f"Starting to copy data from {source_bucket}/{source_object} to "
+301 f"{destination_bucket}/{destination_object}"
+302 )
+303 FileManager._copy_objects(
+304 source_bucket=source_bucket,
+305 source_object=source_object,
+306 destination_bucket=destination_bucket,
+307 destination_object=destination_object,
+308 dry_run=dry_run,
+309 )
+310 FileManager._logger.info(
+311 f"Finished copying data, data should be available on {destination_bucket}/"
+312 f"{destination_object}"
+313 )
294class ArchiveFileManager(object):
-295 """Set of actions to restore archives."""
-296
-297 _logger = LoggingHandler(__name__).get_logger()
-298
-299 @staticmethod
-300 def _get_archived_object(bucket: str, object_key: str) -> Optional[Any]:
-301 """Get the archived object if it's an object.
-302
-303 Args:
-304 bucket: name of bucket to check get the object.
-305 object_key: object to get.
-306
-307 Returns:
-308 S3 Object if it's an archived object, otherwise None.
-309 """
-310 s3 = boto3.resource("s3")
-311 object_to_restore = s3.Object(bucket, object_key)
-312
-313 if (
-314 object_to_restore.storage_class is not None
-315 and object_to_restore.storage_class in ARCHIVE_STORAGE_CLASS
-316 ):
-317 return object_to_restore
-318 else:
-319 return None
-320
-321 @staticmethod
-322 def _check_object_restore_status(
-323 bucket: str, object_key: str
-324 ) -> Optional[RestoreStatus]:
-325 """Check the restore status of the archive.
-326
-327 Args:
-328 bucket: name of bucket to check the restore status.
-329 object_key: object to check the restore status.
-330
-331 Returns:
-332 The restore status represented by an enum, possible values are:
-333 NOT_STARTED, ONGOING or RESTORED
-334 """
-335 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
-336
-337 if archived_object is None:
-338 status = None
-339 elif archived_object.restore is None:
-340 status = RestoreStatus.NOT_STARTED
-341 elif 'ongoing-request="true"' in archived_object.restore:
-342 status = RestoreStatus.ONGOING
-343 else:
-344 status = RestoreStatus.RESTORED
-345
-346 return status
-347
-348 @staticmethod
-349 def check_restore_status(source_bucket: str, source_object: str) -> dict:
-350 """Check the restore status of archived data.
-351
-352 Args:
-353 source_bucket: name of bucket to check the restore status.
-354 source_object: object to check the restore status.
-355
-356 Returns:
-357 A dict containing the amount of objects in each status.
-358 """
-359 not_started_objects = 0
-360 ongoing_objects = 0
-361 restored_objects = 0
-362 total_objects = 0
-363
-364 objects_to_restore = _list_objects_recursively(
-365 bucket=source_bucket, path=source_object
-366 )
-367
-368 for obj in objects_to_restore:
-369 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
-370
-371 restore_status = ArchiveFileManager._check_object_restore_status(
-372 source_bucket, obj
-373 )
-374 if not restore_status:
-375 ArchiveFileManager._logger.warning(
-376 f"Restore status not found for {source_bucket}/{obj}"
-377 )
-378 else:
-379 total_objects += 1
+ 378class ArchiveFileManager(object):
+379 """Set of actions to restore archives."""
380
-381 if RestoreStatus.NOT_STARTED == restore_status:
-382 not_started_objects += 1
-383 elif RestoreStatus.ONGOING == restore_status:
-384 ongoing_objects += 1
-385 else:
-386 restored_objects += 1
-387
-388 ArchiveFileManager._logger.info(
-389 f"{obj} restore status is {restore_status.value}"
-390 )
-391
-392 return {
-393 "total_objects": total_objects,
-394 "not_started_objects": not_started_objects,
-395 "ongoing_objects": ongoing_objects,
-396 "restored_objects": restored_objects,
-397 }
-398
-399 @staticmethod
-400 def _request_restore_object(
-401 bucket: str, object_key: str, expiration: int, retrieval_tier: str
-402 ) -> None:
-403 """Request a restore of the archive.
+381 _logger = LoggingHandler(__name__).get_logger()
+382
+383 @staticmethod
+384 def _get_archived_object(bucket: str, object_key: str) -> Optional[Any]:
+385 """Get the archived object if it's an object.
+386
+387 Args:
+388 bucket: name of bucket to check get the object.
+389 object_key: object to get.
+390
+391 Returns:
+392 S3 Object if it's an archived object, otherwise None.
+393 """
+394 s3 = boto3.resource("s3")
+395 object_to_restore = s3.Object(bucket, object_key)
+396
+397 if (
+398 object_to_restore.storage_class is not None
+399 and object_to_restore.storage_class in ARCHIVE_STORAGE_CLASS
+400 ):
+401 return object_to_restore
+402 else:
+403 return None
404
-405 Args:
-406 bucket: name of bucket to perform the restore.
-407 object_key: object to be restored.
-408 expiration: restore expiration.
-409 retrieval_tier: type of restore, possible values are:
-410 Bulk, Standard or Expedited.
-411 """
-412 if not RestoreType.exists(retrieval_tier):
-413 raise RestoreTypeNotFoundException(
-414 f"Restore type {retrieval_tier} not supported."
-415 )
-416
-417 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
-418
-419 if archived_object and archived_object.restore is None:
-420 ArchiveFileManager._logger.info(f"Restoring archive {bucket}/{object_key}.")
-421 archived_object.restore_object(
-422 RestoreRequest={
-423 "Days": expiration,
-424 "GlacierJobParameters": {"Tier": retrieval_tier},
-425 }
-426 )
+405 @staticmethod
+406 def _check_object_restore_status(
+407 bucket: str, object_key: str
+408 ) -> Optional[RestoreStatus]:
+409 """Check the restore status of the archive.
+410
+411 Args:
+412 bucket: name of bucket to check the restore status.
+413 object_key: object to check the restore status.
+414
+415 Returns:
+416 The restore status represented by an enum, possible values are:
+417 NOT_STARTED, ONGOING or RESTORED
+418 """
+419 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
+420
+421 if archived_object is None:
+422 status = None
+423 elif archived_object.restore is None:
+424 status = RestoreStatus.NOT_STARTED
+425 elif 'ongoing-request="true"' in archived_object.restore:
+426 status = RestoreStatus.ONGOING
427 else:
-428 ArchiveFileManager._logger.info(
-429 f"Restore request for {bucket}/{object_key} not performed."
-430 )
+428 status = RestoreStatus.RESTORED
+429
+430 return status
431
432 @staticmethod
-433 def request_restore(
-434 source_bucket: str,
-435 source_object: str,
-436 restore_expiration: int,
-437 retrieval_tier: str,
-438 dry_run: bool,
-439 ) -> None:
-440 """Request the restore of archived data.
-441
-442 Args:
-443 source_bucket: name of bucket to perform the restore.
-444 source_object: object to be restored.
-445 restore_expiration: restore expiration in days.
-446 retrieval_tier: type of restore, possible values are:
-447 Bulk, Standard or Expedited.
-448 dry_run: if dry_run is set to True the function will print a dict with
-449 all the paths that would be deleted based on the given keys.
-450 """
-451 if dry_run:
-452 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
-453
-454 ArchiveFileManager._logger.info("Paths that would be restored:")
-455 ArchiveFileManager._logger.info(response)
-456 else:
-457 objects_to_restore = _list_objects_recursively(
-458 bucket=source_bucket, path=source_object
-459 )
-460
-461 for obj in objects_to_restore:
-462 ArchiveFileManager._request_restore_object(
-463 source_bucket,
-464 obj,
-465 restore_expiration,
-466 retrieval_tier,
-467 )
-468
-469 @staticmethod
-470 def request_restore_and_wait(
-471 source_bucket: str,
-472 source_object: str,
-473 restore_expiration: int,
-474 retrieval_tier: str,
-475 dry_run: bool,
-476 ) -> None:
-477 """Request and wait for the restore to complete, polling the restore status.
+433 def check_restore_status(source_bucket: str, source_object: str) -> dict:
+434 """Check the restore status of archived data.
+435
+436 Args:
+437 source_bucket: name of bucket to check the restore status.
+438 source_object: object to check the restore status.
+439
+440 Returns:
+441 A dict containing the amount of objects in each status.
+442 """
+443 not_started_objects = 0
+444 ongoing_objects = 0
+445 restored_objects = 0
+446 total_objects = 0
+447
+448 if _check_directory(source_bucket, source_object):
+449 source_object = _process_directory_path(source_object)
+450
+451 objects_to_restore = _list_objects_recursively(
+452 bucket=source_bucket, path=source_object
+453 )
+454
+455 for obj in objects_to_restore:
+456 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
+457
+458 restore_status = ArchiveFileManager._check_object_restore_status(
+459 source_bucket, obj
+460 )
+461 if not restore_status:
+462 ArchiveFileManager._logger.warning(
+463 f"Restore status not found for {source_bucket}/{obj}"
+464 )
+465 else:
+466 total_objects += 1
+467
+468 if RestoreStatus.NOT_STARTED == restore_status:
+469 not_started_objects += 1
+470 elif RestoreStatus.ONGOING == restore_status:
+471 ongoing_objects += 1
+472 else:
+473 restored_objects += 1
+474
+475 ArchiveFileManager._logger.info(
+476 f"{obj} restore status is {restore_status.value}"
+477 )
478
-479 Args:
-480 source_bucket: name of bucket to perform the restore.
-481 source_object: object to be restored.
-482 restore_expiration: restore expiration in days.
-483 retrieval_tier: type of restore, possible values are:
-484 Bulk, Standard or Expedited.
-485 dry_run: if dry_run is set to True the function will print a dict with
-486 all the paths that would be deleted based on the given keys.
-487 """
-488 if retrieval_tier != RestoreType.EXPEDITED.value:
-489 ArchiveFileManager._logger.error(
-490 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-491 "kind of restore should be used just with `Expedited` retrieval tier "
-492 "to save cluster costs."
-493 )
-494 raise ValueError(
-495 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-496 "kind of restore should be used just with `Expedited` retrieval tier "
-497 "to save cluster costs."
-498 )
-499
-500 ArchiveFileManager.request_restore(
-501 source_bucket=source_bucket,
-502 source_object=source_object,
-503 restore_expiration=restore_expiration,
-504 retrieval_tier=retrieval_tier,
-505 dry_run=dry_run,
-506 )
-507 restore_status = ArchiveFileManager.check_restore_status(
-508 source_bucket, source_object
-509 )
-510 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
-511
-512 if not dry_run:
-513 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
-514 wait_time = 300
-515 while restore_status.get("total_objects") > restore_status.get(
-516 "restored_objects"
-517 ):
-518 ArchiveFileManager._logger.info(
-519 "Not all objects have been restored yet, checking the status again "
-520 f"in {wait_time} seconds."
-521 )
-522 time.sleep(wait_time)
-523 wait_time = 30
-524 restore_status = ArchiveFileManager.check_restore_status(
-525 source_bucket, source_object
-526 )
-527 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+479 return {
+480 "total_objects": total_objects,
+481 "not_started_objects": not_started_objects,
+482 "ongoing_objects": ongoing_objects,
+483 "restored_objects": restored_objects,
+484 }
+485
+486 @staticmethod
+487 def _request_restore_object(
+488 bucket: str, object_key: str, expiration: int, retrieval_tier: str
+489 ) -> None:
+490 """Request a restore of the archive.
+491
+492 Args:
+493 bucket: name of bucket to perform the restore.
+494 object_key: object to be restored.
+495 expiration: restore expiration.
+496 retrieval_tier: type of restore, possible values are:
+497 Bulk, Standard or Expedited.
+498 """
+499 if not RestoreType.exists(retrieval_tier):
+500 raise RestoreTypeNotFoundException(
+501 f"Restore type {retrieval_tier} not supported."
+502 )
+503
+504 if _check_directory(bucket, object_key):
+505 object_key = _process_directory_path(object_key)
+506
+507 archived_object = ArchiveFileManager._get_archived_object(bucket, object_key)
+508
+509 if archived_object and archived_object.restore is None:
+510 ArchiveFileManager._logger.info(f"Restoring archive {bucket}/{object_key}.")
+511 archived_object.restore_object(
+512 RestoreRequest={
+513 "Days": expiration,
+514 "GlacierJobParameters": {"Tier": retrieval_tier},
+515 }
+516 )
+517 else:
+518 ArchiveFileManager._logger.info(
+519 f"Restore request for {bucket}/{object_key} not performed."
+520 )
+521
+522 @staticmethod
+523 def request_restore(
+524 source_bucket: str,
+525 source_object: str,
+526 restore_expiration: int,
+527 retrieval_tier: str,
+528 dry_run: bool,
+529 ) -> None:
+530 """Request the restore of archived data.
+531
+532 Args:
+533 source_bucket: name of bucket to perform the restore.
+534 source_object: object to be restored.
+535 restore_expiration: restore expiration in days.
+536 retrieval_tier: type of restore, possible values are:
+537 Bulk, Standard or Expedited.
+538 dry_run: if dry_run is set to True the function will print a dict with
+539 all the paths that would be deleted based on the given keys.
+540 """
+541 if _check_directory(source_bucket, source_object):
+542 source_object = _process_directory_path(source_object)
+543
+544 if dry_run:
+545 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
+546
+547 ArchiveFileManager._logger.info("Paths that would be restored:")
+548 ArchiveFileManager._logger.info(response)
+549 else:
+550 objects_to_restore = _list_objects_recursively(
+551 bucket=source_bucket, path=source_object
+552 )
+553
+554 for obj in objects_to_restore:
+555 ArchiveFileManager._request_restore_object(
+556 source_bucket,
+557 obj,
+558 restore_expiration,
+559 retrieval_tier,
+560 )
+561
+562 @staticmethod
+563 def request_restore_and_wait(
+564 source_bucket: str,
+565 source_object: str,
+566 restore_expiration: int,
+567 retrieval_tier: str,
+568 dry_run: bool,
+569 ) -> None:
+570 """Request and wait for the restore to complete, polling the restore status.
+571
+572 Args:
+573 source_bucket: name of bucket to perform the restore.
+574 source_object: object to be restored.
+575 restore_expiration: restore expiration in days.
+576 retrieval_tier: type of restore, possible values are:
+577 Bulk, Standard or Expedited.
+578 dry_run: if dry_run is set to True the function will print a dict with
+579 all the paths that would be deleted based on the given keys.
+580 """
+581 if retrieval_tier != RestoreType.EXPEDITED.value:
+582 ArchiveFileManager._logger.error(
+583 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+584 "kind of restore should be used just with `Expedited` retrieval tier "
+585 "to save cluster costs."
+586 )
+587 raise ValueError(
+588 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+589 "kind of restore should be used just with `Expedited` retrieval tier "
+590 "to save cluster costs."
+591 )
+592
+593 ArchiveFileManager.request_restore(
+594 source_bucket=source_bucket,
+595 source_object=source_object,
+596 restore_expiration=restore_expiration,
+597 retrieval_tier=retrieval_tier,
+598 dry_run=dry_run,
+599 )
+600 restore_status = ArchiveFileManager.check_restore_status(
+601 source_bucket, source_object
+602 )
+603 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+604
+605 if not dry_run:
+606 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
+607 wait_time = 300
+608 while restore_status.get("total_objects") > restore_status.get(
+609 "restored_objects"
+610 ):
+611 ArchiveFileManager._logger.info(
+612 "Not all objects have been restored yet, checking the status again "
+613 f"in {wait_time} seconds."
+614 )
+615 time.sleep(wait_time)
+616 wait_time = 30
+617 restore_status = ArchiveFileManager.check_restore_status(
+618 source_bucket, source_object
+619 )
+620 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
348 @staticmethod
-349 def check_restore_status(source_bucket: str, source_object: str) -> dict:
-350 """Check the restore status of archived data.
-351
-352 Args:
-353 source_bucket: name of bucket to check the restore status.
-354 source_object: object to check the restore status.
-355
-356 Returns:
-357 A dict containing the amount of objects in each status.
-358 """
-359 not_started_objects = 0
-360 ongoing_objects = 0
-361 restored_objects = 0
-362 total_objects = 0
-363
-364 objects_to_restore = _list_objects_recursively(
-365 bucket=source_bucket, path=source_object
-366 )
-367
-368 for obj in objects_to_restore:
-369 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
-370
-371 restore_status = ArchiveFileManager._check_object_restore_status(
-372 source_bucket, obj
-373 )
-374 if not restore_status:
-375 ArchiveFileManager._logger.warning(
-376 f"Restore status not found for {source_bucket}/{obj}"
-377 )
-378 else:
-379 total_objects += 1
-380
-381 if RestoreStatus.NOT_STARTED == restore_status:
-382 not_started_objects += 1
-383 elif RestoreStatus.ONGOING == restore_status:
-384 ongoing_objects += 1
-385 else:
-386 restored_objects += 1
-387
-388 ArchiveFileManager._logger.info(
-389 f"{obj} restore status is {restore_status.value}"
-390 )
-391
-392 return {
-393 "total_objects": total_objects,
-394 "not_started_objects": not_started_objects,
-395 "ongoing_objects": ongoing_objects,
-396 "restored_objects": restored_objects,
-397 }
+ 432 @staticmethod
+433 def check_restore_status(source_bucket: str, source_object: str) -> dict:
+434 """Check the restore status of archived data.
+435
+436 Args:
+437 source_bucket: name of bucket to check the restore status.
+438 source_object: object to check the restore status.
+439
+440 Returns:
+441 A dict containing the amount of objects in each status.
+442 """
+443 not_started_objects = 0
+444 ongoing_objects = 0
+445 restored_objects = 0
+446 total_objects = 0
+447
+448 if _check_directory(source_bucket, source_object):
+449 source_object = _process_directory_path(source_object)
+450
+451 objects_to_restore = _list_objects_recursively(
+452 bucket=source_bucket, path=source_object
+453 )
+454
+455 for obj in objects_to_restore:
+456 ArchiveFileManager._logger.info(f"Checking restore status for: {obj}")
+457
+458 restore_status = ArchiveFileManager._check_object_restore_status(
+459 source_bucket, obj
+460 )
+461 if not restore_status:
+462 ArchiveFileManager._logger.warning(
+463 f"Restore status not found for {source_bucket}/{obj}"
+464 )
+465 else:
+466 total_objects += 1
+467
+468 if RestoreStatus.NOT_STARTED == restore_status:
+469 not_started_objects += 1
+470 elif RestoreStatus.ONGOING == restore_status:
+471 ongoing_objects += 1
+472 else:
+473 restored_objects += 1
+474
+475 ArchiveFileManager._logger.info(
+476 f"{obj} restore status is {restore_status.value}"
+477 )
+478
+479 return {
+480 "total_objects": total_objects,
+481 "not_started_objects": not_started_objects,
+482 "ongoing_objects": ongoing_objects,
+483 "restored_objects": restored_objects,
+484 }
432 @staticmethod
-433 def request_restore(
-434 source_bucket: str,
-435 source_object: str,
-436 restore_expiration: int,
-437 retrieval_tier: str,
-438 dry_run: bool,
-439 ) -> None:
-440 """Request the restore of archived data.
-441
-442 Args:
-443 source_bucket: name of bucket to perform the restore.
-444 source_object: object to be restored.
-445 restore_expiration: restore expiration in days.
-446 retrieval_tier: type of restore, possible values are:
-447 Bulk, Standard or Expedited.
-448 dry_run: if dry_run is set to True the function will print a dict with
-449 all the paths that would be deleted based on the given keys.
-450 """
-451 if dry_run:
-452 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
-453
-454 ArchiveFileManager._logger.info("Paths that would be restored:")
-455 ArchiveFileManager._logger.info(response)
-456 else:
-457 objects_to_restore = _list_objects_recursively(
-458 bucket=source_bucket, path=source_object
-459 )
-460
-461 for obj in objects_to_restore:
-462 ArchiveFileManager._request_restore_object(
-463 source_bucket,
-464 obj,
-465 restore_expiration,
-466 retrieval_tier,
-467 )
+ 522 @staticmethod
+523 def request_restore(
+524 source_bucket: str,
+525 source_object: str,
+526 restore_expiration: int,
+527 retrieval_tier: str,
+528 dry_run: bool,
+529 ) -> None:
+530 """Request the restore of archived data.
+531
+532 Args:
+533 source_bucket: name of bucket to perform the restore.
+534 source_object: object to be restored.
+535 restore_expiration: restore expiration in days.
+536 retrieval_tier: type of restore, possible values are:
+537 Bulk, Standard or Expedited.
+538 dry_run: if dry_run is set to True the function will print a dict with
+539 all the paths that would be deleted based on the given keys.
+540 """
+541 if _check_directory(source_bucket, source_object):
+542 source_object = _process_directory_path(source_object)
+543
+544 if dry_run:
+545 response = _dry_run(bucket=source_bucket, object_paths=[source_object])
+546
+547 ArchiveFileManager._logger.info("Paths that would be restored:")
+548 ArchiveFileManager._logger.info(response)
+549 else:
+550 objects_to_restore = _list_objects_recursively(
+551 bucket=source_bucket, path=source_object
+552 )
+553
+554 for obj in objects_to_restore:
+555 ArchiveFileManager._request_restore_object(
+556 source_bucket,
+557 obj,
+558 restore_expiration,
+559 retrieval_tier,
+560 )
469 @staticmethod
-470 def request_restore_and_wait(
-471 source_bucket: str,
-472 source_object: str,
-473 restore_expiration: int,
-474 retrieval_tier: str,
-475 dry_run: bool,
-476 ) -> None:
-477 """Request and wait for the restore to complete, polling the restore status.
-478
-479 Args:
-480 source_bucket: name of bucket to perform the restore.
-481 source_object: object to be restored.
-482 restore_expiration: restore expiration in days.
-483 retrieval_tier: type of restore, possible values are:
-484 Bulk, Standard or Expedited.
-485 dry_run: if dry_run is set to True the function will print a dict with
-486 all the paths that would be deleted based on the given keys.
-487 """
-488 if retrieval_tier != RestoreType.EXPEDITED.value:
-489 ArchiveFileManager._logger.error(
-490 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-491 "kind of restore should be used just with `Expedited` retrieval tier "
-492 "to save cluster costs."
-493 )
-494 raise ValueError(
-495 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
-496 "kind of restore should be used just with `Expedited` retrieval tier "
-497 "to save cluster costs."
-498 )
-499
-500 ArchiveFileManager.request_restore(
-501 source_bucket=source_bucket,
-502 source_object=source_object,
-503 restore_expiration=restore_expiration,
-504 retrieval_tier=retrieval_tier,
-505 dry_run=dry_run,
-506 )
-507 restore_status = ArchiveFileManager.check_restore_status(
-508 source_bucket, source_object
-509 )
-510 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
-511
-512 if not dry_run:
-513 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
-514 wait_time = 300
-515 while restore_status.get("total_objects") > restore_status.get(
-516 "restored_objects"
-517 ):
-518 ArchiveFileManager._logger.info(
-519 "Not all objects have been restored yet, checking the status again "
-520 f"in {wait_time} seconds."
-521 )
-522 time.sleep(wait_time)
-523 wait_time = 30
-524 restore_status = ArchiveFileManager.check_restore_status(
-525 source_bucket, source_object
-526 )
-527 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+ 562 @staticmethod
+563 def request_restore_and_wait(
+564 source_bucket: str,
+565 source_object: str,
+566 restore_expiration: int,
+567 retrieval_tier: str,
+568 dry_run: bool,
+569 ) -> None:
+570 """Request and wait for the restore to complete, polling the restore status.
+571
+572 Args:
+573 source_bucket: name of bucket to perform the restore.
+574 source_object: object to be restored.
+575 restore_expiration: restore expiration in days.
+576 retrieval_tier: type of restore, possible values are:
+577 Bulk, Standard or Expedited.
+578 dry_run: if dry_run is set to True the function will print a dict with
+579 all the paths that would be deleted based on the given keys.
+580 """
+581 if retrieval_tier != RestoreType.EXPEDITED.value:
+582 ArchiveFileManager._logger.error(
+583 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+584 "kind of restore should be used just with `Expedited` retrieval tier "
+585 "to save cluster costs."
+586 )
+587 raise ValueError(
+588 f"Retrieval Tier {retrieval_tier} not allowed on this operation! This "
+589 "kind of restore should be used just with `Expedited` retrieval tier "
+590 "to save cluster costs."
+591 )
+592
+593 ArchiveFileManager.request_restore(
+594 source_bucket=source_bucket,
+595 source_object=source_object,
+596 restore_expiration=restore_expiration,
+597 retrieval_tier=retrieval_tier,
+598 dry_run=dry_run,
+599 )
+600 restore_status = ArchiveFileManager.check_restore_status(
+601 source_bucket, source_object
+602 )
+603 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
+604
+605 if not dry_run:
+606 ArchiveFileManager._logger.info("Checking the restore status in 5 minutes.")
+607 wait_time = 300
+608 while restore_status.get("total_objects") > restore_status.get(
+609 "restored_objects"
+610 ):
+611 ArchiveFileManager._logger.info(
+612 "Not all objects have been restored yet, checking the status again "
+613 f"in {wait_time} seconds."
+614 )
+615 time.sleep(wait_time)
+616 wait_time = 30
+617 restore_status = ArchiveFileManager.check_restore_status(
+618 source_bucket, source_object
+619 )
+620 ArchiveFileManager._logger.info(f"Restore status: {restore_status}")
44 else:
45 raise WrongArgumentsException("Hashing approach is not supported.")
46
-47 if suffix and suffix != "":
-48 masked_df = masked_df.drop(col)
-49
-50 return masked_df
-51
-52 return inner
-53
-54 @classmethod
-55 def column_dropper(cls, cols: List[str]) -> Callable:
-56 """Drop specific columns.
+47 return masked_df
+48
+49 return inner
+50
+51 @classmethod
+52 def column_dropper(cls, cols: List[str]) -> Callable:
+53 """Drop specific columns.
+54
+55 Args:
+56 cols: list of column names to drop.
57
-58 Args:
-59 cols: list of column names to drop.
-60
-61 Returns:
-62 A function to be called in .transform() spark function.
-63 """
-64
-65 def inner(df: DataFrame) -> DataFrame:
-66 drop_df = df
-67 for col in cols:
-68 drop_df = drop_df.drop(col)
-69
-70 return drop_df
-71
-72 return inner
+58 Returns:
+59 A function to be called in .transform() spark function.
+60 """
+61
+62 def inner(df: DataFrame) -> DataFrame:
+63 drop_df = df
+64 for col in cols:
+65 drop_df = drop_df.drop(col)
+66
+67 return drop_df
+68
+69 return inner
45 else:
46 raise WrongArgumentsException("Hashing approach is not supported.")
47
-48 if suffix and suffix != "":
-49 masked_df = masked_df.drop(col)
-50
-51 return masked_df
-52
-53 return inner
-54
-55 @classmethod
-56 def column_dropper(cls, cols: List[str]) -> Callable:
-57 """Drop specific columns.
+48 return masked_df
+49
+50 return inner
+51
+52 @classmethod
+53 def column_dropper(cls, cols: List[str]) -> Callable:
+54 """Drop specific columns.
+55
+56 Args:
+57 cols: list of column names to drop.
58
-59 Args:
-60 cols: list of column names to drop.
-61
-62 Returns:
-63 A function to be called in .transform() spark function.
-64 """
-65
-66 def inner(df: DataFrame) -> DataFrame:
-67 drop_df = df
-68 for col in cols:
-69 drop_df = drop_df.drop(col)
-70
-71 return drop_df
-72
-73 return inner
+59 Returns:
+60 A function to be called in .transform() spark function.
+61 """
+62
+63 def inner(df: DataFrame) -> DataFrame:
+64 drop_df = df
+65 for col in cols:
+66 drop_df = drop_df.drop(col)
+67
+68 return drop_df
+69
+70 return inner
@@ -264,12 +258,9 @@
45 else:
46 raise WrongArgumentsException("Hashing approach is not supported.")
47
-48 if suffix and suffix != "":
-49 masked_df = masked_df.drop(col)
-50
-51 return masked_df
-52
-53 return inner
+48 return masked_df
+49
+50 return inner
@@ -321,25 +312,25 @@
- 55 @classmethod
-56 def column_dropper(cls, cols: List[str]) -> Callable:
-57 """Drop specific columns.
+ 52 @classmethod
+53 def column_dropper(cls, cols: List[str]) -> Callable:
+54 """Drop specific columns.
+55
+56 Args:
+57 cols: list of column names to drop.
58
-59 Args:
-60 cols: list of column names to drop.
-61
-62 Returns:
-63 A function to be called in .transform() spark function.
-64 """
-65
-66 def inner(df: DataFrame) -> DataFrame:
-67 drop_df = df
-68 for col in cols:
-69 drop_df = drop_df.drop(col)
-70
-71 return drop_df
-72
-73 return inner
+59 Returns:
+60 A function to be called in .transform() spark function.
+61 """
+62
+63 def inner(df: DataFrame) -> DataFrame:
+64 drop_df = df
+65 for col in cols:
+66 drop_df = drop_df.drop(col)
+67
+68 return drop_df
+69
+70 return inner
Inherited Members