Add GridFS integration to large_file

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 *

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}
+    large_file = get_class(args)
-
-    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
 
     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)

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

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

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

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(
+    def configure_credentials(  # type: ignore
-        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(
         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

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

mongo.ini

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