andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Add GridFS integration to large_file

Browse source
This commit is contained in:
Andras Schmelczer 2022-06-03 14:15:01 +02:00
parent a4e16d5739
commit a0d3a9cb2a
8 changed files with 187 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

2
great_ai/src/great_ai/__init__.py
View file

@ -1,3 +1,3 @@
from .great_ai import *
from .open_s3 import *
from .large_file import *
from .utilities import *

46
great_ai/src/great_ai/large_file/__main__.py
View file

@ -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)

1
great_ai/src/great_ai/large_file/large_file/__init__.py
View file

@ -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

26
great_ai/src/great_ai/large_file/large_file/large_file.py
View file

@ -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

121
great_ai/src/great_ai/large_file/large_file/large_file_mongo.py Normal file
View file

@ -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])

25
great_ai/src/great_ai/large_file/large_file/large_file_s3.py
View file

@ -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

5
great_ai/src/great_ai/large_file/parse_arguments.py
View file

@ -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,
)

3
mongo.ini Normal file
View file

@ -0,0 +1,3 @@
[DEFAULT]
connection_string=mongodb://localhost:27017/
database=large_file_tests