366 lines
12 KiB
Python
366 lines
12 KiB
Python
import configparser
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
from types import TracebackType
|
|
from typing import IO, Any, Dict, List, Optional, Type, Union
|
|
|
|
import boto3
|
|
from helper import DownloadProgressBar, UploadProgressBar, human_readable_to_byte
|
|
|
|
logger = logging.getLogger("open_s3")
|
|
|
|
|
|
class LargeFile:
|
|
"""
|
|
Store large files in S3. Use local cache for speed up.
|
|
|
|
Examples:
|
|
|
|
```
|
|
with LargeFile("test.txt", "w", keep_last_n=3) as f:
|
|
for i in range(1000000):
|
|
f.write('test\n')
|
|
|
|
with LargeFile("test.txt", "r") as f:
|
|
print(f.readlines()[0])
|
|
|
|
path_to_cached_text_file = LargeFile("test.txt", version=0).get()
|
|
```
|
|
|
|
By default, files are stored in the ".cache" folder and the
|
|
least recently use is deleted after the overall size reaches 30 GBs.
|
|
|
|
Change it with the following properties.
|
|
|
|
```
|
|
LargeFile.cache_path = Path(".cache")
|
|
LargeFile.max_cache_size = "30GB"
|
|
```
|
|
"""
|
|
|
|
region_name = None
|
|
access_key_id = None
|
|
secret_access_key = None
|
|
bucket_name = None
|
|
endpoint_url = None
|
|
|
|
cache_path = Path(".cache")
|
|
max_cache_size: Optional[str] = "30GB"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
mode: str = "r",
|
|
*,
|
|
buffering: int = -1,
|
|
encoding: Optional[str] = None,
|
|
errors: Optional[str] = None,
|
|
newline: Optional[str] = None,
|
|
version: Optional[int] = None,
|
|
keep_last_n: Optional[int] = None,
|
|
offline_mode: bool = False,
|
|
):
|
|
self._name: str = name
|
|
self._version = version
|
|
self._mode: str = mode
|
|
self._keep_last_n = keep_last_n
|
|
self._offline_mode = offline_mode
|
|
|
|
self._buffering = buffering
|
|
self._encoding = encoding
|
|
self._errors = errors
|
|
self._newline = newline
|
|
|
|
LargeFile.cache_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
self._find_versions()
|
|
self._check_mode_and_set_version()
|
|
|
|
@classmethod
|
|
def configure_credentials(
|
|
cls,
|
|
*,
|
|
aws_region_name: str,
|
|
aws_access_key_id: str,
|
|
aws_secret_access_key: str,
|
|
large_files_bucket_name: str,
|
|
endpoint_url: Optional[str] = None,
|
|
**_: Dict[str, Any],
|
|
) -> None:
|
|
cls.region_name = aws_region_name
|
|
cls.access_key_id = aws_access_key_id
|
|
cls.secret_access_key = aws_secret_access_key
|
|
cls.bucket_name = large_files_bucket_name
|
|
cls.endpoint_url = endpoint_url
|
|
|
|
@classmethod
|
|
def configure_credentials_from_file(
|
|
cls,
|
|
secrets_path: Union[Path, str],
|
|
) -> None:
|
|
|
|
if isinstance(secrets_path, str):
|
|
secrets_path = Path(secrets_path)
|
|
|
|
if not secrets_path.exists():
|
|
raise FileNotFoundError(secrets_path.resolve())
|
|
|
|
credentials = configparser.ConfigParser()
|
|
credentials.read(secrets_path)
|
|
credentials.default_section
|
|
cls.configure_credentials(**credentials[credentials.default_section])
|
|
|
|
def __enter__(self) -> IO:
|
|
self._file: IO[Any] = (
|
|
tempfile.NamedTemporaryFile(
|
|
mode=self._mode,
|
|
buffering=self._buffering,
|
|
encoding=self._encoding,
|
|
newline=self._newline,
|
|
errors=self._errors,
|
|
delete=False,
|
|
prefix="large-file-",
|
|
)
|
|
if "w" in self._mode
|
|
else open(
|
|
self.get(),
|
|
mode=self._mode,
|
|
buffering=self._buffering,
|
|
encoding=self._encoding,
|
|
newline=self._newline,
|
|
errors=self._errors,
|
|
)
|
|
)
|
|
|
|
return self._file
|
|
|
|
def __exit__(
|
|
self,
|
|
type: Optional[Type[BaseException]],
|
|
exc: Optional[BaseException],
|
|
traceback: Optional[TracebackType],
|
|
) -> bool:
|
|
self._file.close()
|
|
|
|
if type is None:
|
|
if "w" in self._mode:
|
|
self.push(Path(self._file.name))
|
|
os.unlink(self._file.name)
|
|
else:
|
|
logger.exception("Could not finish operation.")
|
|
|
|
return True
|
|
|
|
@property
|
|
def version_ids(self) -> List[int]:
|
|
return [self._get_version_from_key(key) for key in self._versions]
|
|
|
|
def get(self, hide_progress: bool = False) -> Path:
|
|
key = next(
|
|
key
|
|
for key in self._versions
|
|
if self._get_version_from_key(key) == self._version
|
|
)
|
|
|
|
destination = self.cache_path / self._local_name
|
|
if not destination.exists():
|
|
logger.info(
|
|
f"File {self._local_name} does not exist locally, starting download from S3"
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_file_archive = Path(tmp) / f"{self._local_name}.tar.gz"
|
|
|
|
size = self._client.head_object(Bucket=self.bucket_name, Key=key)[
|
|
"ContentLength"
|
|
]
|
|
self._client.download_file(
|
|
Bucket=self.bucket_name,
|
|
Key=key,
|
|
Filename=str(tmp_file_archive),
|
|
Callback=None
|
|
if hide_progress
|
|
else DownloadProgressBar(size=size, name=key, logger=logger),
|
|
)
|
|
logger.info(f"Decompressing {self._local_name}")
|
|
shutil.unpack_archive(str(tmp_file_archive), tmp, "gztar")
|
|
tmp_file = Path(tmp) / self._local_name
|
|
shutil.move(str(tmp_file), str(destination))
|
|
else:
|
|
logger.info(f"File {self._local_name} found in cache")
|
|
|
|
return destination
|
|
|
|
def push(self, path: Union[Path, str], hide_progress: bool = False) -> None:
|
|
if isinstance(path, str):
|
|
path = Path(path)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
if path.is_file():
|
|
logger.info(f"Copying file for {self._local_name}")
|
|
copy: Any = shutil.copy
|
|
else:
|
|
logger.info(f"Copying directory for {self._local_name}")
|
|
copy = shutil.copytree
|
|
|
|
try:
|
|
# Make local copy in the cache
|
|
copy(str(path), str(self.cache_path / self._local_name))
|
|
except shutil.SameFileError:
|
|
pass # No worries
|
|
|
|
copy(str(path), str(Path(tmp) / self._local_name))
|
|
|
|
logger.info(f"Compressing {self._local_name}")
|
|
shutil.make_archive(
|
|
str(Path(tmp) / self._local_name),
|
|
"gztar",
|
|
tmp,
|
|
)
|
|
|
|
logger.info(f"Uploading {self._local_name} to S3 from {path}")
|
|
|
|
file_to_be_uploaded = Path(tmp) / f"{self._local_name}.tar.gz"
|
|
self._client.upload_file(
|
|
Filename=str(file_to_be_uploaded),
|
|
Bucket=self.bucket_name,
|
|
Key=self._s3_name,
|
|
Callback=None
|
|
if hide_progress
|
|
else UploadProgressBar(file_to_be_uploaded, logger=logger),
|
|
)
|
|
|
|
self.clean_up()
|
|
|
|
def delete(self) -> None:
|
|
self._keep_last_n = 0
|
|
self._delete_old_versions_from_s3()
|
|
|
|
def clean_up(self) -> None:
|
|
self._delete_old_versions_from_s3()
|
|
self._delete_old_versions_from_disk()
|
|
|
|
def _create_client(self) -> None:
|
|
if (
|
|
self.region_name is None
|
|
or self.access_key_id is None
|
|
or self.secret_access_key is None
|
|
or self.bucket_name is None
|
|
):
|
|
raise ValueError(
|
|
"Please configure the S3 access options by calling LargeFile.configure_credentials or set offline_mode=True in the constructor."
|
|
)
|
|
|
|
self._client = boto3.client(
|
|
"s3",
|
|
aws_access_key_id=self.access_key_id,
|
|
aws_secret_access_key=self.secret_access_key,
|
|
region_name=self.region_name,
|
|
endpoint_url=self.endpoint_url,
|
|
)
|
|
|
|
def _find_versions(self) -> None:
|
|
if self._offline_mode:
|
|
self._fetch_versions_from_cache()
|
|
else:
|
|
self._create_client()
|
|
self._fetch_versions_from_s3()
|
|
|
|
if self._versions:
|
|
logger.info(f"Found versions: {self.version_ids}")
|
|
else:
|
|
logger.info("No versions found")
|
|
|
|
def _fetch_versions_from_cache(self) -> None:
|
|
logger.info(f"Fetching offline versions of {self._name}")
|
|
|
|
self._versions = [
|
|
path for path in self.cache_path.glob(f"{self._local_name}-*")
|
|
]
|
|
|
|
def _fetch_versions_from_s3(self) -> None:
|
|
logger.info(f"Fetching online versions of {self._name}")
|
|
found_objects = self._client.list_objects_v2(
|
|
Bucket=self.bucket_name, Prefix=self._name
|
|
)
|
|
self._versions = (
|
|
sorted(o["Key"] for o in found_objects["Contents"])
|
|
if "Contents" in found_objects
|
|
else []
|
|
)
|
|
|
|
def _check_mode_and_set_version(self) -> None:
|
|
if "+" in self._mode:
|
|
raise ValueError("Read-write mode is not allowed.")
|
|
|
|
if "w" in self._mode:
|
|
if self._version is not None:
|
|
raise ValueError("Providing a version is not allowed in write mode.")
|
|
|
|
self._version = self.version_ids[-1] + 1 if self.version_ids else 0
|
|
|
|
elif "r" in self._mode:
|
|
if not self.version_ids:
|
|
raise FileNotFoundError(
|
|
f"File {self._name} not found. No versions are available."
|
|
)
|
|
|
|
if self._version is None:
|
|
self._version = self.version_ids[-1]
|
|
logger.info(f"Lastest version of {self._local_name} is {self._version}")
|
|
|
|
elif self._version not in self.version_ids:
|
|
raise FileNotFoundError(
|
|
f"File {self._name} not found with version {self._version}. Available versions: {self.version_ids}"
|
|
)
|
|
else:
|
|
raise ValueError("Unsupported file mode.")
|
|
|
|
@property
|
|
def _local_name(self) -> str:
|
|
return f"{self._name}-{self._version}"
|
|
|
|
@property
|
|
def _s3_name(self) -> str:
|
|
return f"{self._name}/{self._version}"
|
|
|
|
@staticmethod
|
|
def _get_version_from_key(key: Union[str, Path]) -> int:
|
|
if isinstance(key, Path):
|
|
return int(key.name.split("-")[-1])
|
|
return int(key.split("/")[-1])
|
|
|
|
def _delete_old_versions_from_s3(self) -> None:
|
|
if self._keep_last_n is not None:
|
|
for key in (
|
|
self._versions[: -self._keep_last_n]
|
|
if self._keep_last_n > 0
|
|
else self._versions
|
|
):
|
|
logger.info(
|
|
f"Removing old version (keep_last_n={self._keep_last_n}): {key}"
|
|
)
|
|
self._client.delete_object(Bucket=self.bucket_name, Key=key)
|
|
|
|
def _delete_old_versions_from_disk(self) -> None:
|
|
self.cache_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
if self.max_cache_size is None:
|
|
return
|
|
|
|
allowed_size = human_readable_to_byte(self.max_cache_size)
|
|
assert allowed_size >= 0
|
|
|
|
least_recently_read = sorted(
|
|
[f for f in self.cache_path.glob("*")], key=lambda f: f.stat().st_atime
|
|
)
|
|
|
|
while sum(os.path.getsize(f) for f in least_recently_read) > allowed_size:
|
|
file = least_recently_read.pop(0)
|
|
logger.info(
|
|
f"Deleting file from cache to meet quota (max_cache_size={self.max_cache_size}): {file}"
|
|
)
|
|
os.unlink(file)
|