bugfix: fix for folders missing in mongoFS

[osm/common.git] / osm_common / fsmongo.py

diff --git a/osm_common/fsmongo.py b/osm_common/fsmongo.py
index ff37c42..b04057e 100644 (file)
--- a/osm_common/fsmongo.py
+++ b/osm_common/fsmongo.py
@@ -22,6 +22,8 @@ from io import BytesIO, StringIO
 import logging
 import os
 import datetime
+import tarfile
+import zipfile
 
 from gridfs import GridFSBucket, errors
 from osm_common.fsbase import FsBase, FsException
@@ -203,10 +205,13 @@ class FsMongo(FsBase):
     def __update_local_fs(self, from_path=None):
         dir_cursor = self.fs.find({"metadata.type": "dir"}, no_cursor_timeout=True)
 
+        valid_paths = []
+
         for directory in dir_cursor:
             if from_path and not directory.filename.startswith(from_path):
                 continue
             os.makedirs(self.path + directory.filename, exist_ok=True)
+            valid_paths.append(self.path + directory.filename)
 
         file_cursor = self.fs.find(
             {"metadata.type": {"$in": ["file", "sym"]}}, no_cursor_timeout=True
@@ -231,6 +236,9 @@ class FsMongo(FsBase):
                         raise
                 os.symlink(link, file_path)
             else:
+                folder = os.path.dirname(file_path)
+                if folder not in valid_paths:
+                    os.makedirs(folder, exist_ok=True)
                 with open(file_path, "wb+") as file_stream:
                     self.fs.download_to_stream(writing_file._id, file_stream)
                 if "permissions" in writing_file.metadata:
@@ -336,6 +344,8 @@ class FsMongo(FsBase):
                     "Multiple files found", http_code=HTTPStatus.INTERNAL_SERVER_ERROR
                 )
 
+            print(requested_file.metadata)
+
             # if no special mode is required just check it does exists
             if not mode:
                 return True
@@ -368,35 +378,60 @@ class FsMongo(FsBase):
 
             return requested_file.length
 
-    def file_extract(self, tar_object, path):
+    def file_extract(self, compressed_object, path):
         """
         extract a tar file
-        :param tar_object: object of type tar
+        :param compressed_object: object of type tar or zip
         :param path: can be a str or a str list, or a tar object where to extract the tar_object
         :return: None
         """
         f = path if isinstance(path, str) else "/".join(path)
 
-        for member in tar_object.getmembers():
-            if member.isfile():
-                stream = tar_object.extractfile(member)
-            elif member.issym():
-                stream = BytesIO(member.linkname.encode("utf-8"))
-            else:
-                stream = BytesIO()
+        if type(compressed_object) is tarfile.TarFile:
+            for member in compressed_object.getmembers():
+                if member.isfile():
+                    stream = compressed_object.extractfile(member)
+                elif member.issym():
+                    stream = BytesIO(member.linkname.encode("utf-8"))
+                else:
+                    stream = BytesIO()
 
-            if member.isfile():
-                file_type = "file"
-            elif member.issym():
-                file_type = "sym"
-            else:
-                file_type = "dir"
+                if member.isfile():
+                    file_type = "file"
+                elif member.issym():
+                    file_type = "sym"
+                else:
+                    file_type = "dir"
 
-            metadata = {"type": file_type, "permissions": member.mode}
+                metadata = {"type": file_type, "permissions": member.mode}
 
-            self.fs.upload_from_stream(f + "/" + member.name, stream, metadata=metadata)
+                self.fs.upload_from_stream(
+                    f + "/" + member.name, stream, metadata=metadata
+                )
 
-            stream.close()
+                stream.close()
+        elif type(compressed_object) is zipfile.ZipFile:
+            for member in compressed_object.infolist():
+                if member.is_dir():
+                    stream = BytesIO()
+                else:
+                    stream = compressed_object.read(member)
+
+                if member.is_dir():
+                    file_type = "dir"
+                else:
+                    file_type = "file"
+
+                metadata = {"type": file_type}
+
+                print("Now uploading...")
+                print(f + "/" + member.filename)
+                self.fs.upload_from_stream(
+                    f + "/" + member.filename, stream, metadata=metadata
+                )
+
+                if member.is_dir():
+                    stream.close()
 
     def file_open(self, storage, mode):
         """
@@ -447,6 +482,9 @@ class FsMongo(FsBase):
                         http_code=HTTPStatus.NOT_FOUND,
                     )
 
+                if f.endswith("/"):
+                    f = f[:-1]
+
                 files_cursor = self.fs.find(
                     {"filename": {"$regex": "^{}/([^/])*".format(f)}}
                 )