302 lines
12 KiB
Python
302 lines
12 KiB
Python
"""Archival utility."""
|
|
|
|
import queue
|
|
import collections
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
_IS_WINDOWS = sys.platform == "win32" or sys.platform == "cygwin"
|
|
|
|
if _IS_WINDOWS:
|
|
import ctypes
|
|
|
|
UploadArgs = collections.namedtuple("UploadArgs", [
|
|
"archival_file", "display_name", "local_file", "content_type", "s3_bucket", "s3_path",
|
|
"delete_file"
|
|
])
|
|
|
|
ArchiveArgs = collections.namedtuple("ArchiveArgs",
|
|
["archival_file", "display_name", "remote_file"])
|
|
|
|
|
|
def file_list_size(files):
|
|
"""Return size (in bytes) of all 'files' and their subdirectories."""
|
|
if isinstance(files, str):
|
|
files = [files]
|
|
file_bytes = 0
|
|
for ifile in files:
|
|
if not os.path.exists(ifile):
|
|
pass
|
|
elif os.path.isdir(ifile):
|
|
file_bytes += directory_size(ifile)
|
|
else:
|
|
file_bytes += os.path.getsize(ifile)
|
|
return file_bytes
|
|
|
|
|
|
def directory_size(directory):
|
|
"""Return size (in bytes) of files in 'directory' tree."""
|
|
dir_bytes = 0
|
|
for root_dir, _, files in os.walk(str(directory)):
|
|
for name in files:
|
|
full_name = os.path.join(root_dir, name)
|
|
try:
|
|
dir_bytes += os.path.getsize(full_name)
|
|
except OSError:
|
|
# A file might be deleted while we are looping through the os.walk() result.
|
|
pass
|
|
return dir_bytes
|
|
|
|
|
|
def free_space(path):
|
|
"""Return file system free space (in bytes) for 'path'."""
|
|
if _IS_WINDOWS:
|
|
dirname = os.path.dirname(path)
|
|
free_bytes = ctypes.c_ulonglong(0)
|
|
ctypes.windll.kernel32.GetDiskFreeSpaceExW(
|
|
ctypes.c_wchar_p(dirname), None, None, ctypes.pointer(free_bytes))
|
|
return free_bytes.value
|
|
stat = os.statvfs(path)
|
|
return stat.f_bavail * stat.f_bsize
|
|
|
|
|
|
def remove_file(file_name):
|
|
"""Attempt to remove file. Return status and message."""
|
|
try:
|
|
# File descriptors, on Windows, are inherited by all subprocesses and file removal may fail
|
|
# because the file is still open.
|
|
# See https://www.python.org/dev/peps/pep-0446/#issues-with-inheritable-file-descriptors
|
|
os.remove(file_name)
|
|
status = 0
|
|
message = "Successfully deleted file {}".format(file_name)
|
|
except Exception as err: # pylint: disable=broad-except
|
|
status = 1
|
|
message = "Error deleting file {}: {}".format(file_name, err)
|
|
return status, message
|
|
|
|
|
|
class Archival(object): # pylint: disable=too-many-instance-attributes
|
|
"""Class to support file archival to S3."""
|
|
|
|
def __init__( # pylint: disable=too-many-arguments
|
|
self, logger, archival_json_file="archive.json", limit_size_mb=0, limit_files=0,
|
|
s3_client=None):
|
|
"""Initialize Archival."""
|
|
|
|
self.archival_json_file = archival_json_file
|
|
self.limit_size_mb = limit_size_mb
|
|
self.limit_files = limit_files
|
|
self.size_mb = 0
|
|
self.num_files = 0
|
|
self.archive_time = 0
|
|
self.logger = logger
|
|
|
|
# Lock to control access from multiple threads.
|
|
self._lock = threading.Lock()
|
|
|
|
# Start the worker thread to update the 'archival_json_file'.
|
|
self._archive_file_queue = queue.Queue()
|
|
self._archive_file_worker = threading.Thread(target=self._update_archive_file_wkr,
|
|
args=(self._archive_file_queue,
|
|
logger), name="archive_file_worker")
|
|
self._archive_file_worker.setDaemon(True)
|
|
self._archive_file_worker.start()
|
|
if not s3_client:
|
|
self.s3_client = self._get_s3_client()
|
|
else:
|
|
self.s3_client = s3_client
|
|
|
|
# Start the worker thread which uploads the archive.
|
|
self._upload_queue = queue.Queue()
|
|
self._upload_worker = threading.Thread(
|
|
target=self._upload_to_s3_wkr, args=(self._upload_queue, self._archive_file_queue,
|
|
logger, self.s3_client), name="upload_worker")
|
|
self._upload_worker.setDaemon(True)
|
|
self._upload_worker.start()
|
|
|
|
@staticmethod
|
|
def _get_s3_client():
|
|
# Since boto3 is a 3rd party module, we import locally.
|
|
import boto3
|
|
return boto3.client("s3")
|
|
|
|
def archive_files_to_s3(self, display_name, input_files, s3_bucket, s3_path):
|
|
"""Archive 'input_files' to 's3_bucket' and 's3_path'.
|
|
|
|
Archive is not done if user specified limits are reached. The size limit is
|
|
enforced after it has been exceeded, since it can only be calculated after the
|
|
tar/gzip has been done.
|
|
|
|
Return status and message, where message contains information if status is non-0.
|
|
"""
|
|
|
|
start_time = time.time()
|
|
with self._lock:
|
|
if not input_files:
|
|
status = 1
|
|
message = "No input_files specified"
|
|
elif self.limit_size_mb and self.size_mb >= self.limit_size_mb:
|
|
status = 1
|
|
message = "Files not archived, {}MB size limit reached".format(self.limit_size_mb)
|
|
elif self.limit_files and self.num_files >= self.limit_files:
|
|
status = 1
|
|
message = "Files not archived, {} file limit reached".format(self.limit_files)
|
|
else:
|
|
status, message, file_size_mb = self._archive_files(display_name, input_files,
|
|
s3_bucket, s3_path)
|
|
|
|
if status == 0:
|
|
self.num_files += 1
|
|
self.size_mb += file_size_mb
|
|
self.archive_time += time.time() - start_time
|
|
|
|
return status, message
|
|
|
|
@staticmethod
|
|
def _update_archive_file_wkr(work_queue, logger):
|
|
"""Worker thread: Update the archival JSON file from 'work_queue'."""
|
|
archival_json = []
|
|
while True:
|
|
archive_args = work_queue.get()
|
|
# Exit worker thread when sentinel is received.
|
|
if archive_args is None:
|
|
work_queue.task_done()
|
|
break
|
|
archival_record = {
|
|
"name": archive_args.display_name, "link": archive_args.remote_file,
|
|
"visibility": "private"
|
|
}
|
|
logger.debug("Updating archive file %s with %s", archive_args.archival_file,
|
|
archival_record)
|
|
archival_json.append(archival_record)
|
|
with open(archive_args.archival_file, "w") as archival_fh:
|
|
json.dump(archival_json, archival_fh)
|
|
work_queue.task_done()
|
|
|
|
@staticmethod
|
|
def _upload_to_s3_wkr(work_queue, archive_file_work_queue, logger, s3_client):
|
|
"""Worker thread: Upload to S3 from 'work_queue', dispatch to 'archive_file_work_queue'."""
|
|
while True:
|
|
upload_args = work_queue.get()
|
|
# Exit worker thread when sentinel is received.
|
|
if upload_args is None:
|
|
work_queue.task_done()
|
|
archive_file_work_queue.put(None)
|
|
break
|
|
extra_args = {"ContentType": upload_args.content_type, "ACL": "public-read"}
|
|
logger.debug("Uploading to S3 %s to bucket %s path %s", upload_args.local_file,
|
|
upload_args.s3_bucket, upload_args.s3_path)
|
|
upload_completed = False
|
|
try:
|
|
s3_client.upload_file(upload_args.local_file, upload_args.s3_bucket,
|
|
upload_args.s3_path, ExtraArgs=extra_args)
|
|
upload_completed = True
|
|
logger.debug("Upload to S3 completed for %s to bucket %s path %s",
|
|
upload_args.local_file, upload_args.s3_bucket, upload_args.s3_path)
|
|
except Exception as err: # pylint: disable=broad-except
|
|
logger.exception("Upload to S3 error %s", err)
|
|
|
|
if upload_args.delete_file:
|
|
status, message = remove_file(upload_args.local_file)
|
|
if status:
|
|
logger.error("Upload to S3 delete file error %s", message)
|
|
|
|
remote_file = "https://s3.amazonaws.com/{}/{}".format(upload_args.s3_bucket,
|
|
upload_args.s3_path)
|
|
if upload_completed:
|
|
archive_file_work_queue.put(
|
|
ArchiveArgs(upload_args.archival_file, upload_args.display_name, remote_file))
|
|
|
|
work_queue.task_done()
|
|
|
|
def _archive_files(self, display_name, input_files, s3_bucket, s3_path):
|
|
"""
|
|
Gather 'input_files' into a single tar/gzip and archive to 's3_path'.
|
|
|
|
The caller waits until the list of files has been tar/gzipped to a temporary file.
|
|
The S3 upload and subsequent update to 'archival_json_file' will be done asynchronosly.
|
|
|
|
Returns status, message and size_mb of archive.
|
|
"""
|
|
|
|
# Parameter 'input_files' can either be a string or list of strings.
|
|
if isinstance(input_files, str):
|
|
input_files = [input_files]
|
|
|
|
message = "Tar/gzip {} files: {}".format(display_name, input_files)
|
|
status = 0
|
|
size_mb = 0
|
|
|
|
# Tar/gzip to a temporary file.
|
|
_, temp_file = tempfile.mkstemp(suffix=".tgz")
|
|
|
|
# Check if there is sufficient space for the temporary tgz file.
|
|
if file_list_size(input_files) > free_space(temp_file):
|
|
status, message = remove_file(temp_file)
|
|
if status:
|
|
self.logger.warning("Removing tarfile due to insufficient space - %s", message)
|
|
return 1, "Insufficient space for {}".format(message), 0
|
|
|
|
try:
|
|
with tarfile.open(temp_file, "w:gz") as tar_handle:
|
|
for input_file in input_files:
|
|
try:
|
|
tar_handle.add(input_file)
|
|
except (IOError, OSError, tarfile.TarError) as err:
|
|
message = "{}; Unable to add {} to archive file: {}".format(
|
|
message, input_file, err)
|
|
except (IOError, OSError, tarfile.TarError) as err:
|
|
status, message = remove_file(temp_file)
|
|
if status:
|
|
self.logger.warning("Removing tarfile due to creation failure - %s", message)
|
|
return 1, str(err), 0
|
|
|
|
# Round up the size of the archive.
|
|
size_mb = int(math.ceil(float(file_list_size(temp_file)) / (1024 * 1024)))
|
|
self._upload_queue.put(
|
|
UploadArgs(self.archival_json_file, display_name, temp_file, "application/x-gzip",
|
|
s3_bucket, s3_path, True))
|
|
|
|
return status, message, size_mb
|
|
|
|
def check_thread(self, thread, expected_alive):
|
|
"""Check if the thread is still active."""
|
|
if thread.isAlive() and not expected_alive:
|
|
self.logger.warning(
|
|
"The %s thread did not complete, some files might not have been uploaded"
|
|
" to S3 or archived to %s.", thread.name, self.archival_json_file)
|
|
elif not thread.isAlive() and expected_alive:
|
|
self.logger.warning(
|
|
"The %s thread is no longer running, some files might not have been uploaded"
|
|
" to S3 or archived to %s.", thread.name, self.archival_json_file)
|
|
|
|
def exit(self, timeout=30):
|
|
"""Wait for worker threads to finish."""
|
|
# Put sentinel on upload queue to trigger worker thread exit.
|
|
self._upload_queue.put(None)
|
|
self.check_thread(self._upload_worker, True)
|
|
self.check_thread(self._archive_file_worker, True)
|
|
self._upload_worker.join(timeout=timeout)
|
|
self.check_thread(self._upload_worker, False)
|
|
|
|
# Archive file worker thread exit should be triggered by upload thread worker.
|
|
self._archive_file_worker.join(timeout=timeout)
|
|
self.check_thread(self._archive_file_worker, False)
|
|
|
|
self.logger.info("Total tar/gzip archive time is %0.2f seconds, for %d file(s) %d MB",
|
|
self.archive_time, self.num_files, self.size_mb)
|
|
|
|
def files_archived_num(self):
|
|
"""Return the number of the archived files."""
|
|
return self.num_files
|
|
|
|
def files_archived_size_mb(self):
|
|
"""Return the size of the archived files."""
|
|
return self.size_mb
|