Merge commit '9408938' from 565_workbench_cleanup (spaetz)

8 files changed, 96 insertions, 43 deletions

diff --git a/mediagoblin/decorators.py b/mediagoblin/decorators.py
index a40f1d5a..804fab7e 100644
--- a/mediagoblin/decorators.py
+++ b/mediagoblin/decorators.py
@@ -20,6 +20,7 @@ from urlparse import urljoin
 from werkzeug.exceptions import Forbidden, NotFound
 from werkzeug.urls import url_quote
 
+from mediagoblin import mg_globals as mgg
 from mediagoblin.db.models import MediaEntry, User
 from mediagoblin.tools.response import redirect, render_404
 
@@ -222,3 +223,14 @@ def get_media_entry_by_id(controller):
         return controller(request, media=media, *args, **kwargs)
 
     return wrapper
+
+
+def get_workbench(func):
+    """Decorator, passing in a workbench as kwarg which is cleaned up afterwards"""
+
+    @wraps(func)
+    def new_func(*args, **kwargs):
+        with mgg.workbench_manager.create() as workbench:
+            return func(*args, workbench=workbench, **kwargs)
+
+    return new_func
diff --git a/mediagoblin/media_types/ascii/processing.py b/mediagoblin/media_types/ascii/processing.py
index 04d1166c..254717eb 100644
--- a/mediagoblin/media_types/ascii/processing.py
+++ b/mediagoblin/media_types/ascii/processing.py
@@ -19,6 +19,7 @@ import Image
 import logging
 
 from mediagoblin import mg_globals as mgg
+from mediagoblin.decorators import get_workbench
 from mediagoblin.processing import create_pub_filepath
 from mediagoblin.media_types.ascii import asciitoimage
 
@@ -38,12 +39,14 @@ def sniff_handler(media_file, **kw):
     return False
 
 
-def process_ascii(entry):
-    '''
-    Code to process a txt file
-    '''
+@get_workbench
+def process_ascii(entry, workbench=None):
+    """Code to process a txt file. Will be run by celery.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits. 
+    """
     ascii_config = mgg.global_config['media_type:mediagoblin.media_types.ascii']
-    workbench = mgg.workbench_manager.create_workbench()
     # Conversions subdirectory to avoid collisions
     conversions_subdir = os.path.join(
         workbench.dir, 'conversions')
diff --git a/mediagoblin/media_types/audio/processing.py b/mediagoblin/media_types/audio/processing.py
index c4ccad49..e12cefe6 100644
--- a/mediagoblin/media_types/audio/processing.py
+++ b/mediagoblin/media_types/audio/processing.py
@@ -19,6 +19,7 @@ from tempfile import NamedTemporaryFile
 import os
 
 from mediagoblin import mg_globals as mgg
+from mediagoblin.decorators import get_workbench
 from mediagoblin.processing import (create_pub_filepath, BadMediaFail,
     FilenameBuilder, ProgressCallback)
 
@@ -42,10 +43,14 @@ def sniff_handler(media_file, **kw):
     return False
 
 
-def process_audio(entry):
-    audio_config = mgg.global_config['media_type:mediagoblin.media_types.audio']
+@get_workbench
+def process_audio(entry, workbench=None):
+    """Code to process uploaded audio. Will be run by celery.
 
-    workbench = mgg.workbench_manager.create_workbench()
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
+    """
+    audio_config = mgg.global_config['media_type:mediagoblin.media_types.audio']
 
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
@@ -143,6 +148,3 @@ def process_audio(entry):
         entry.media_files['thumb'] = ['fake', 'thumb', 'path.jpg']
 
     mgg.queue_store.delete_file(queued_filepath)
-
-    # clean up workbench
-    workbench.destroy_self()
diff --git a/mediagoblin/media_types/image/processing.py b/mediagoblin/media_types/image/processing.py
index bf464069..e6a34ca0 100644
--- a/mediagoblin/media_types/image/processing.py
+++ b/mediagoblin/media_types/image/processing.py
@@ -19,6 +19,7 @@ import os
 import logging
 
 from mediagoblin import mg_globals as mgg
+from mediagoblin.decorators import get_workbench
 from mediagoblin.processing import BadMediaFail, \
     create_pub_filepath, FilenameBuilder
 from mediagoblin.tools.exif import exif_fix_image_orientation, \
@@ -76,11 +77,13 @@ def sniff_handler(media_file, **kw):
     return False
 
 
-def process_image(entry):
-    """
-    Code to process an image
+@get_workbench
+def process_image(entry, workbench=None):
+    """Code to process an image. Will be run by celery.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
     """
-    workbench = mgg.workbench_manager.create_workbench()
     # Conversions subdirectory to avoid collisions
     conversions_subdir = os.path.join(
         workbench.dir, 'conversions')
@@ -147,8 +150,6 @@ def process_image(entry):
             gps_data['gps_' + key] = gps_data.pop(key)
         entry.media_data_init(**gps_data)
 
-    # clean up workbench
-    workbench.destroy_self()
 
 if __name__ == '__main__':
     import sys
diff --git a/mediagoblin/media_types/stl/processing.py b/mediagoblin/media_types/stl/processing.py
index cd949e2a..3089f295 100644
--- a/mediagoblin/media_types/stl/processing.py
+++ b/mediagoblin/media_types/stl/processing.py
@@ -21,6 +21,7 @@ import subprocess
 import pkg_resources
 
 from mediagoblin import mg_globals as mgg
+from mediagoblin.decorators import get_workbench
 from mediagoblin.processing import create_pub_filepath, \
     FilenameBuilder
 
@@ -75,11 +76,13 @@ def blender_render(config):
         env=env)
 
 
-def process_stl(entry):
-    """
-    Code to process an stl or obj model.
+@get_workbench
+def process_stl(entry, workbench=None):
+    """Code to process an stl or obj model. Will be run by celery.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
     """
-    workbench = mgg.workbench_manager.create_workbench()
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
         mgg.queue_store, queued_filepath, 'source')
@@ -164,7 +167,7 @@ def process_stl(entry):
     # Remove queued media file from storage and database
     mgg.queue_store.delete_file(queued_filepath)
     entry.queued_media_file = []
-        
+
     # Insert media file information into database
     media_files_dict = entry.setdefault('media_files', {})
     media_files_dict[u'original'] = model_filepath
@@ -185,6 +188,3 @@ def process_stl(entry):
         "file_type" : ext,
         }
     entry.media_data_init(**dimensions)
-
-    # clean up workbench
-    workbench.destroy_self()
diff --git a/mediagoblin/media_types/video/processing.py b/mediagoblin/media_types/video/processing.py
index 703c4681..22c4355d 100644
--- a/mediagoblin/media_types/video/processing.py
+++ b/mediagoblin/media_types/video/processing.py
@@ -18,6 +18,7 @@ from tempfile import NamedTemporaryFile
 import logging
 
 from mediagoblin import mg_globals as mgg
+from mediagoblin.decorators import get_workbench
 from mediagoblin.processing import \
     create_pub_filepath, FilenameBuilder, BaseProcessingFail, ProgressCallback
 from mediagoblin.tools.translate import lazy_pass_to_ugettext as _
@@ -51,16 +52,17 @@ def sniff_handler(media_file, **kw):
 
     return False
 
-
-def process_video(entry):
+@get_workbench
+def process_video(entry, workbench=None):
     """
     Process a video entry, transcode the queued media files (originals) and
     create a thumbnail for the entry.
+
+    A Workbench() represents a local tempory dir. It is automatically
+    cleaned up when this function exits.
     """
     video_config = mgg.global_config['media_type:mediagoblin.media_types.video']
 
-    workbench = mgg.workbench_manager.create_workbench()
-
     queued_filepath = entry.queued_media_file
     queued_filename = workbench.localized_file(
         mgg.queue_store, queued_filepath,
diff --git a/mediagoblin/tests/test_workbench.py b/mediagoblin/tests/test_workbench.py
index 04a74653..9da8eea0 100644
--- a/mediagoblin/tests/test_workbench.py
+++ b/mediagoblin/tests/test_workbench.py
@@ -19,6 +19,8 @@ import tempfile
 
 
 from mediagoblin import workbench
+from mediagoblin.mg_globals import setup_globals
+from mediagoblin.decorators import get_workbench
 from mediagoblin.tests.test_storage import get_tmp_filestorage
 
 
@@ -28,19 +30,20 @@ class TestWorkbench(object):
             os.path.join(tempfile.gettempdir(), u'mgoblin_workbench_testing'))
 
     def test_create_workbench(self):
-        workbench = self.workbench_manager.create_workbench()
+        workbench = self.workbench_manager.create()
         assert os.path.isdir(workbench.dir)
         assert workbench.dir.startswith(self.workbench_manager.base_workbench_dir)
+        workbench.destroy()
 
     def test_joinpath(self):
-        this_workbench = self.workbench_manager.create_workbench()
+        this_workbench = self.workbench_manager.create()
         tmpname = this_workbench.joinpath('temp.txt')
         assert tmpname == os.path.join(this_workbench.dir, 'temp.txt')
-        this_workbench.destroy_self()
+        this_workbench.destroy()
 
     def test_destroy_workbench(self):
         # kill a workbench
-        this_workbench = self.workbench_manager.create_workbench()
+        this_workbench = self.workbench_manager.create()
         tmpfile_name = this_workbench.joinpath('temp.txt')
         tmpfile = file(tmpfile_name, 'w')
         with tmpfile:
@@ -49,14 +52,14 @@ class TestWorkbench(object):
         assert os.path.exists(tmpfile_name)
 
         wb_dir = this_workbench.dir
-        this_workbench.destroy_self()
+        this_workbench.destroy()
         assert not os.path.exists(tmpfile_name)
         assert not os.path.exists(wb_dir)
 
     def test_localized_file(self):
         tmpdir, this_storage = get_tmp_filestorage()
-        this_workbench = self.workbench_manager.create_workbench()
-        
+        this_workbench = self.workbench_manager.create()
+
         # Write a brand new file
         filepath = ['dir1', 'dir2', 'ourfile.txt']
 
@@ -78,7 +81,7 @@ class TestWorkbench(object):
         filename = this_workbench.localized_file(this_storage, filepath)
         assert filename == os.path.join(
             this_workbench.dir, 'ourfile.txt')
-        
+
         # fake remote file storage, filename_if_copying set
         filename = this_workbench.localized_file(
             this_storage, filepath, 'thisfile')
@@ -91,3 +94,18 @@ class TestWorkbench(object):
             this_storage, filepath, 'thisfile.text', False)
         assert filename == os.path.join(
             this_workbench.dir, 'thisfile.text')
+
+    def test_workbench_decorator(self):
+        """Test @get_workbench decorator and automatic cleanup"""
+        # The decorator needs mg_globals.workbench_manager
+        setup_globals(workbench_manager=self.workbench_manager)
+
+        @get_workbench
+        def create_it(workbench=None):
+            # workbench dir exists?
+            assert os.path.isdir(workbench.dir)
+            return workbench.dir
+
+        benchdir = create_it()
+        # workbench dir has been cleaned up automatically?
+        assert not os.path.isdir(benchdir)
diff --git a/mediagoblin/workbench.py b/mediagoblin/workbench.py
index 2331b551..0d4db52b 100644
--- a/mediagoblin/workbench.py
+++ b/mediagoblin/workbench.py
@@ -119,7 +119,7 @@ class Workbench(object):
 
             return full_dest_filename
 
-    def destroy_self(self):
+    def destroy(self):
         """
         Destroy this workbench!  Deletes the directory and all its contents!
 
@@ -127,18 +127,33 @@ class Workbench(object):
         """
         # just in case
         workbench = os.path.abspath(self.dir)
-
         shutil.rmtree(workbench)
-
         del self.dir
 
+    def __enter__(self):
+        """Make Workbench a context manager so we can use `with Workbench() as bench:`"""
+        return self
+
+    def  __exit__(self, *args):
+        """Clean up context manager, aka ourselves, deleting the workbench"""
+        self.destroy()
+
 
 class WorkbenchManager(object):
     """
     A system for generating and destroying workbenches.
 
-    Workbenches are actually just subdirectories of a temporary storage space
-    for during the processing stage.
+    Workbenches are actually just subdirectories of a (local) temporary
+    storage space for during the processing stage. The preferred way to
+    create them is to use:
+
+        with workbenchmger.create() as workbench:
+            do stuff...
+
+    This will automatically clean up all temporary directories even in
+    case of an exceptions. Also check the
+    @mediagoblin.decorators.get_workbench decorator for a convenient
+    wrapper.
     """
 
     def __init__(self, base_workbench_dir):
@@ -146,7 +161,7 @@ class WorkbenchManager(object):
         if not os.path.exists(self.base_workbench_dir):
             os.makedirs(self.base_workbench_dir)
 
-    def create_workbench(self):
+    def create(self):
         """
         Create and return the path to a new workbench (directory).
         """