From a0d3a9cb2ae0b715e140ed743a953c75e19d9512 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= Date: Fri, 3 Jun 2022 14:15:01 +0200 Subject: [PATCH] Add GridFS integration to large_file --- great_ai/src/great_ai/__init__.py | 2 +- great_ai/src/great_ai/large_file/__main__.py | 46 ++++--- .../large_file/large_file/__init__.py | 1 + .../large_file/large_file/large_file.py | 26 +++- .../large_file/large_file/large_file_mongo.py | 121 ++++++++++++++++++ .../large_file/large_file/large_file_s3.py | 25 +--- .../great_ai/large_file/parse_arguments.py | 5 +- mongo.ini | 3 + 8 files changed, 187 insertions(+), 42 deletions(-) create mode 100644 great_ai/src/great_ai/large_file/large_file/large_file_mongo.py create mode 100644 mongo.ini diff --git a/great_ai/src/great_ai/__init__.py b/great_ai/src/great_ai/__init__.py index 2644c5b..a31d902 100644 --- a/great_ai/src/great_ai/__init__.py +++ b/great_ai/src/great_ai/__init__.py @@ -1,3 +1,3 @@ from .great_ai import * -from .open_s3 import * +from .large_file import * from .utilities import * diff --git a/great_ai/src/great_ai/large_file/__main__.py b/great_ai/src/great_ai/large_file/__main__.py index 9881a22..6079d9d 100644 --- a/great_ai/src/great_ai/large_file/__main__.py +++ b/great_ai/src/great_ai/large_file/__main__.py @@ -1,10 +1,12 @@ #!/usr/bin/env python3 +from argparse import Namespace from pathlib import Path +from typing import Mapping, Type from great_ai.utilities.logger import get_logger -from .large_file import LargeFileLocal, LargeFileS3 +from .large_file import LargeFile, LargeFileLocal, LargeFileMongo, LargeFileS3 from .parse_arguments import parse_arguments logger = get_logger("large_file") @@ -13,21 +15,7 @@ logger = get_logger("large_file") def main() -> None: parser, args = parse_arguments() - factory = {"s3": LargeFileS3, "local": LargeFileLocal} - - if args.driver not in factory: - raise ValueError( - f"Driver {args.driver} does not exits, available options: {' ,'.join(factory.keys())}" - ) - - large_file = factory[args.driver] - - if args.driver != "local": - if args.secrets is None: - raise ValueError( - "Providing a credentials file is required when the driver mode is not `local`." - ) - large_file.configure_credentials_from_file(args.secrets) # type: ignore + large_file = get_class(args) if not args.cache and not args.push and not args.delete: logger.warning("No action required.") @@ -50,6 +38,30 @@ def main() -> None: large_file(f).delete() +def get_class(args: Namespace) -> Type[LargeFile]: + factory: Mapping[str, Type[LargeFile]] = { + "s3": LargeFileS3, + "local": LargeFileLocal, + "mongodb": LargeFileMongo, + } + + if args.backend not in factory: + raise ValueError( + f"Backend {args.backend} does not exits, available options: {' ,'.join(factory.keys())}" + ) + + large_file = factory[args.backend] + + if args.backend != "local": + if args.secrets is None: + raise ValueError( + "Providing a credentials file is required when the backend mode is not `local`." + ) + large_file.configure_credentials_from_file(args.secrets) + + return large_file + + if __name__ == "__main__": try: main() @@ -57,4 +69,4 @@ if __name__ == "__main__": logger.warning("Exiting") exit() except Exception as e: - logger.error(e) + logger.exception(e) diff --git a/great_ai/src/great_ai/large_file/large_file/__init__.py b/great_ai/src/great_ai/large_file/large_file/__init__.py index cd28af0..60b5aaf 100644 --- a/great_ai/src/great_ai/large_file/large_file/__init__.py +++ b/great_ai/src/great_ai/large_file/large_file/__init__.py @@ -1,3 +1,4 @@ from .large_file import LargeFile from .large_file_local import LargeFileLocal +from .large_file_mongo import LargeFileMongo from .large_file_s3 import LargeFileS3 diff --git a/great_ai/src/great_ai/large_file/large_file/large_file.py b/great_ai/src/great_ai/large_file/large_file/large_file.py index ca3a434..57883dc 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file.py @@ -1,7 +1,8 @@ +import configparser import os import shutil import tempfile -from abc import ABC, abstractmethod +from abc import ABC, abstractclassmethod, abstractmethod from pathlib import Path from types import TracebackType from typing import IO, Any, List, Optional, Type, Union, cast @@ -79,6 +80,28 @@ class LargeFile(ABC): self._find_instances() self._check_mode_and_set_version() + @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]) + + @abstractclassmethod + def configure_credentials( + cls, + ) -> None: + pass + def __enter__(self) -> IO: self._file: IO[Any] = ( tempfile.NamedTemporaryFile( @@ -168,6 +191,7 @@ class LargeFile(ABC): try: # Make local copy in the cache + shutil.rmtree(self.cache_path / self._local_name, ignore_errors=True) copy(str(path), str(self.cache_path / self._local_name)) except shutil.SameFileError: pass # No worries diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py b/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py new file mode 100644 index 0000000..836b828 --- /dev/null +++ b/great_ai/src/great_ai/large_file/large_file/large_file_mongo.py @@ -0,0 +1,121 @@ +import re +from functools import cached_property +from pathlib import Path +from typing import Any, List, Mapping + +from gridfs import DEFAULT_CHUNK_SIZE, Database, GridFSBucket +from pymongo import MongoClient + +from great_ai.utilities.logger import get_logger + +from ..helper import DownloadProgressBar, UploadProgressBar +from ..models import DataInstance +from .large_file import LargeFile + +logger = get_logger("large_file") + + +MONGO_NAME_VERSION_SEPARATOR = "_" + + +class LargeFileMongo(LargeFile): + connection_string = None + database = None + + @classmethod + def configure_credentials( # type: ignore + cls, + *, + connection_string: str, + database: str, + **_: Mapping[str, Any], + ) -> None: + cls.connection_string = connection_string + cls.database = database + + @cached_property + def _client(self) -> GridFSBucket: + if self.connection_string is None or self.database is None: + raise ValueError( + "Please configure the MongoDB access options by calling LargeFileMongo.configure_credentials or set offline_mode=True in the constructor." + ) + + db: Database = MongoClient(self.connection_string)[self.database] + return GridFSBucket(db) + + def _find_remote_instances(self) -> List[DataInstance]: + logger.debug(f"Fetching Mongo (GridFS) versions of {self._name}") + + return [ + DataInstance( + name=MONGO_NAME_VERSION_SEPARATOR.join( + f.name.split(MONGO_NAME_VERSION_SEPARATOR)[:-1] + ), + version=int(f.name.split(MONGO_NAME_VERSION_SEPARATOR)[-1]), + remote_path=(f._id, f.length), + origin="mongodb", + ) + for f in self._client.find( + { + "filename": re.compile( + re.escape(self._name + MONGO_NAME_VERSION_SEPARATOR) + ".*" + ) + } + ) + ] + + def _download( + self, remote_path: Any, local_path: Path, hide_progress: bool + ) -> None: + logger.info(f"Downloading {remote_path[0]} from Mongo (GridFS)") + + progress = ( + DownloadProgressBar( + name=str(remote_path[0]), size=remote_path[1], logger=logger + ) + if not hide_progress + else None + ) + with self._client.open_download_stream(remote_path[0]) as stream: + with open(local_path, "wb") as f: + while True: + content = stream.read(DEFAULT_CHUNK_SIZE) + f.write(content) + + if progress: + progress(len(content)) + if len(content) < DEFAULT_CHUNK_SIZE: + break + + def _upload(self, local_path: Path, hide_progress: bool) -> None: + logger.info(f"Uploading {local_path} to Mongo (GridFS)") + + progress = ( + UploadProgressBar(path=local_path, logger=logger) + if not hide_progress + else None + ) + with self._client.open_upload_stream( + f"{self._name}{MONGO_NAME_VERSION_SEPARATOR}{self.version}" + ) as stream: + with open(local_path, "rb") as f: + while True: + content = f.read(DEFAULT_CHUNK_SIZE) + stream.write(content) + + if progress: + progress(len(content)) + if len(content) < DEFAULT_CHUNK_SIZE: + break + + def _delete_old_remote_versions(self) -> None: + if self._keep_last_n is not None: + for i in ( + self._instances[: -self._keep_last_n] + if self._keep_last_n > 0 + else self._instances + ): + logger.info( + f"Removing old version from MongoDB (GridFS) (keep_last_n={self._keep_last_n}): {i.name}{MONGO_NAME_VERSION_SEPARATOR}{i.version}" + ) + self._client.delete(i.remote_path[0]) diff --git a/great_ai/src/great_ai/large_file/large_file/large_file_s3.py b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py index 5c7a8e0..ad15df7 100644 --- a/great_ai/src/great_ai/large_file/large_file/large_file_s3.py +++ b/great_ai/src/great_ai/large_file/large_file/large_file_s3.py @@ -1,7 +1,6 @@ -import configparser from functools import cached_property from pathlib import Path -from typing import Any, List, Mapping, Optional, Union +from typing import Any, List, Mapping, Optional import boto3 @@ -52,23 +51,7 @@ class LargeFileS3(LargeFile): endpoint_url = None @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]) - - @classmethod - def configure_credentials( + def configure_credentials( # type: ignore cls, *, aws_region_name: str, @@ -93,7 +76,7 @@ class LargeFileS3(LargeFile): 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." + "Please configure the S3 access options by calling LargeFileS3.configure_credentials or set offline_mode=True in the constructor." ) return boto3.client( @@ -136,7 +119,7 @@ class LargeFileS3(LargeFile): self._client.download_file( Bucket=self.bucket_name, - Key=hide_progress, + Key=remote_path, Filename=str(local_path), Callback=None if hide_progress diff --git a/great_ai/src/great_ai/large_file/parse_arguments.py b/great_ai/src/great_ai/large_file/parse_arguments.py index 18c395a..b8cb966 100644 --- a/great_ai/src/great_ai/large_file/parse_arguments.py +++ b/great_ai/src/great_ai/large_file/parse_arguments.py @@ -8,7 +8,8 @@ def parse_arguments() -> Tuple[ArgumentParser, Namespace]: ) parser.add_argument( - "--driver", + "-b", + "--backend", type=str, help="choose which backend to use, available options: `local`, `s3`, `mongodb`", required=True, @@ -27,7 +28,7 @@ def parse_arguments() -> Tuple[ArgumentParser, Namespace]: "--cache", nargs="+", type=str, - help="download file into local cache, example: file_name:version", + help="download file into local cache, example: file_name.txt:version", required=False, ) diff --git a/mongo.ini b/mongo.ini new file mode 100644 index 0000000..3f5dfb1 --- /dev/null +++ b/mongo.ini @@ -0,0 +1,3 @@ +[DEFAULT] +connection_string=mongodb://localhost:27017/ +database=large_file_tests