Remove symlink if exists
[osm/common.git] / osm_common / fsmongo.py
index c558d8e..740d540 100644 (file)
 # contact: eduardo.sousa@canonical.com
 ##
 
+import errno
+from http import HTTPStatus
 from io import BytesIO, StringIO
-from pymongo import MongoClient
-from gridfs import GridFSBucket, errors
 import logging
-from http import HTTPStatus
 import os
+
+from gridfs import GridFSBucket, errors
 from osm_common.fsbase import FsBase, FsException
+from pymongo import MongoClient
+
 
 __author__ = "Eduardo Sousa <eduardo.sousa@canonical.com>"
 
@@ -34,6 +37,7 @@ class GridByteStream(BytesIO):
         self.filename = filename
         self.fs = fs
         self.mode = mode
+        self.file_type = "file"  # Set "file" as default file_type
 
         self.__initialize__()
 
@@ -48,8 +52,9 @@ class GridByteStream(BytesIO):
             if exception_file:
                 raise FsException("Multiple files found", http_code=HTTPStatus.INTERNAL_SERVER_ERROR)
 
-            if requested_file.metadata["type"] == "file":
+            if requested_file.metadata["type"] in ("file", "sym"):
                 grid_file = requested_file
+                self.file_type = requested_file.metadata["type"]
             else:
                 raise FsException("Type isn't file", http_code=HTTPStatus.INTERNAL_SERVER_ERROR)
 
@@ -84,13 +89,13 @@ class GridByteStream(BytesIO):
                 self._id,
                 self.filename,
                 self,
-                metadata={"type": "file"}
+                metadata={"type": self.file_type}
             )
         else:
             self.fs.upload_from_stream(
                 self.filename,
                 self,
-                metadata={"type": "file"}
+                metadata={"type": self.file_type}
             )
         super(GridByteStream, self).close()
 
@@ -108,6 +113,7 @@ class GridStringStream(StringIO):
         self.filename = filename
         self.fs = fs
         self.mode = mode
+        self.file_type = "file"  # Set "file" as default file_type
 
         self.__initialize__()
 
@@ -122,8 +128,9 @@ class GridStringStream(StringIO):
             if exception_file:
                 raise FsException("Multiple files found", http_code=HTTPStatus.INTERNAL_SERVER_ERROR)
 
-            if requested_file.metadata["type"] == "file":
+            if requested_file.metadata["type"] in ("file", "dir"):
                 grid_file = requested_file
+                self.file_type = requested_file.metadata["type"]
             else:
                 raise FsException("File type isn't file", http_code=HTTPStatus.INTERNAL_SERVER_ERROR)
 
@@ -165,13 +172,13 @@ class GridStringStream(StringIO):
                 self._id,
                 self.filename,
                 stream,
-                metadata={"type": "file"}
+                metadata={"type": self.file_type}
             )
         else:
             self.fs.upload_from_stream(
                 self.filename,
                 stream,
-                metadata={"type": "file"}
+                metadata={"type": self.file_type}
             )
         stream.close()
         super(GridStringStream, self).close()
@@ -191,25 +198,41 @@ class FsMongo(FsBase):
         self.client = None
         self.fs = None
 
-    def __update_local_fs(self):
+    def __update_local_fs(self, from_path=None):
         dir_cursor = self.fs.find({"metadata.type": "dir"}, no_cursor_timeout=True)
 
         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)
 
-        file_cursor = self.fs.find({"metadata.type": "file"}, no_cursor_timeout=True)
+        file_cursor = self.fs.find({"metadata.type": {"$in": ["file", "sym"]}}, no_cursor_timeout=True)
 
         for writing_file in file_cursor:
+            if from_path and not writing_file.filename.startswith(from_path):
+                continue
             file_path = self.path + writing_file.filename
-            file_stream = open(file_path, 'wb+')
-            self.fs.download_to_stream(writing_file._id, file_stream)
-            file_stream.close()
-            if "permissions" in writing_file.metadata:
-                os.chmod(file_path, writing_file.metadata["permissions"])
 
-    def get_params(self):
-        self.__update_local_fs()
+            if writing_file.metadata["type"] == "sym":
+                with BytesIO() as b:
+                    self.fs.download_to_stream(writing_file._id, b)
+                    b.seek(0)
+                    link = b.read().decode("utf-8")
+
+                try:
+                    os.remove(file_path)
+                except OSError as e:
+                    if e.errno != errno.ENOENT:
+                        # This is probably permission denied or worse
+                        raise
+                os.symlink(link, file_path)
+            else:
+                with open(file_path, 'wb+') as file_stream:
+                    self.fs.download_to_stream(writing_file._id, file_stream)
+                if "permissions" in writing_file.metadata:
+                    os.chmod(file_path, writing_file.metadata["permissions"])
 
+    def get_params(self):
         return {"fs": "mongo", "path": self.path}
 
     def fs_connect(self, config):
@@ -305,6 +328,9 @@ class FsMongo(FsBase):
             if requested_file.metadata["type"] == mode:
                 return True
 
+            if requested_file.metadata["type"] == "sym" and mode == "file":
+                return True
+
         return False
 
     def file_size(self, storage):
@@ -337,11 +363,20 @@ class FsMongo(FsBase):
         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 member.isfile():
+                file_type = "file"
+            elif member.issym():
+                file_type = "sym"
+            else:
+                file_type = "dir"
+
             metadata = {
-                "type": "file" if member.isfile() else "dir",
+                "type": file_type,
                 "permissions": member.mode
             }
 
@@ -427,6 +462,14 @@ class FsMongo(FsBase):
                 else:
                     self.fs.delete(requested_file._id)
             if not found and not ignore_non_exist:
-                raise FsException("File {} does not exist".format(storage), http_code=HTTPStatus.NOT_FOUND)    
+                raise FsException("File {} does not exist".format(storage), http_code=HTTPStatus.NOT_FOUND)
         except IOError as e:
             raise FsException("File {} cannot be deleted: {}".format(f, e), http_code=HTTPStatus.INTERNAL_SERVER_ERROR)
+
+    def sync(self, from_path=None):
+        """
+        Sync from FSMongo to local storage
+        :param from_path: if supplied, only copy content from this path, not all
+        :return: None
+        """
+        self.__update_local_fs(from_path=from_path)