Merge pull request ckan#3756 from smotornyuk/close-temporary-files-after-request

wardi · web-flow · commit 7333332e4515 · 2018-02-01T17:36:19.000-05:00
[ckan#3752] CloseWSGIInput middleware + update in ResourceUploader
diff --git a/ckan/config/middleware/common_middleware.py b/ckan/config/middleware/common_middleware.py
@@ -5,8 +5,10 @@
 import urllib2
 import hashlib
 import json
+import cgi
 
 import sqlalchemy as sa
+from webob.request import FakeCGIBody
 
 
 class RootPathMiddleware(object):
@@ -28,6 +30,32 @@ def __call__(self, environ, start_response):
         return self.app(environ, start_response)
 
 
+class CloseWSGIInputMiddleware(object):
+    '''
+    webob.request.Request has habit to create FakeCGIBody. This leads(
+    during file upload) to creating temporary files that are not closed.
+    For long lived processes this means that for each upload you will
+    spend the same amount of temporary space as size of uploaded
+    file additionally, until server restart(this will automatically
+    close temporary files).
+
+    This middleware is supposed to close such files after each request.
+    '''
+    def __init__(self, app, config):
+        self.app = app
+
+    def __call__(self, environ, start_response):
+        wsgi_input = environ['wsgi.input']
+        if isinstance(wsgi_input, FakeCGIBody):
+            for _, item in wsgi_input.vars.items():
+                if not isinstance(item, cgi.FieldStorage):
+                    continue
+                fp = getattr(item, 'fp', None)
+                if fp is not None:
+                    fp.close()
+        return self.app(environ, start_response)
+
+
 class TrackingMiddleware(object):
 
     def __init__(self, app, config):
diff --git a/ckan/config/middleware/pylons_app.py b/ckan/config/middleware/pylons_app.py
@@ -57,8 +57,8 @@ def make_pylons_stack(conf, full_stack=True, static_files=True,
     for plugin in PluginImplementations(IMiddleware):
         app = plugin.make_middleware(app, config)
 
+    app = common_middleware.CloseWSGIInputMiddleware(app, config)
     app = common_middleware.RootPathMiddleware(app, config)
-
     # Routing/Session/Cache Middleware
     app = RoutesMiddleware(app, config['routes.map'])
     # we want to be able to retrieve the routes middleware to be able to update
diff --git a/ckan/lib/uploader.py b/ckan/lib/uploader.py
@@ -15,13 +15,30 @@
 from ckan.common import config
 
 ALLOWED_UPLOAD_TYPES = (cgi.FieldStorage, FlaskFileStorage)
+MB = 1 << 20
+
 log = logging.getLogger(__name__)
 
 _storage_path = None
 _max_resource_size = None
 _max_image_size = None
 
 
+def _copy_file(input_file, output_file, max_size):
+    input_file.seek(0)
+    current_size = 0
+    while True:
+        current_size = current_size + 1
+        # MB chunks
+        data = input_file.read(MB)
+
+        if not data:
+            break
+        output_file.write(data)
+        if current_size > max_size:
+            raise logic.ValidationError({'upload': ['File upload too large']})
+
+
 def _get_underlying_file(wrapper):
     if isinstance(wrapper, FlaskFileStorage):
         return wrapper.stream
@@ -162,22 +179,14 @@ def upload(self, max_size=2):
         max_size is size in MB maximum of the file'''
 
         if self.filename:
-            output_file = open(self.tmp_filepath, 'wb')
-            self.upload_file.seek(0)
-            current_size = 0
-            while True:
-                current_size = current_size + 1
-                # MB chunks
-                data = self.upload_file.read(2 ** 20)
-                if not data:
-                    break
-                output_file.write(data)
-                if current_size > max_size:
+            with open(self.tmp_filepath, 'wb+') as output_file:
+                try:
+                    _copy_file(self.upload_file, output_file, max_size)
+                except logic.ValidationError:
                     os.remove(self.tmp_filepath)
-                    raise logic.ValidationError(
-                        {self.file_field: ['File upload too large']}
-                    )
-            output_file.close()
+                    raise
+                finally:
+                    self.upload_file.close()
             os.rename(self.tmp_filepath, self.filepath)
             self.clear = True
 
@@ -285,24 +294,14 @@ def upload(self, id, max_size=10):
                 if e.errno != 17:
                     raise
             tmp_filepath = filepath + '~'
-            output_file = open(tmp_filepath, 'wb+')
-            self.upload_file.seek(0)
-            current_size = 0
-            while True:
-                current_size = current_size + 1
-                # MB chunks
-                data = self.upload_file.read(2 ** 20)
-
-                if not data:
-                    break
-                output_file.write(data)
-                if current_size > max_size:
+            with open(tmp_filepath, 'wb+') as output_file:
+                try:
+                    _copy_file(self.upload_file, output_file, max_size)
+                except logic.ValidationError:
                     os.remove(tmp_filepath)
-                    raise logic.ValidationError(
-                        {'upload': ['File upload too large']}
-                    )
-
-            output_file.close()
+                    raise
+                finally:
+                    self.upload_file.close()
             os.rename(tmp_filepath, filepath)
             return
 
