Merge pull request #115 from gisce/add-sudo

Add sudo support for oorq

oorq/decorators.py

@@ -20,6 +20,7 @@
     DB_CURSOR_SAVEPOINT
 )
 from autoworker import AutoWorker
+from ctx import sudo
 
 
 JobToProcess = namedtuple('JobToProcess', ['job', 'queue', 'at_front'])
@@ -136,6 +137,10 @@ def f_job(*args, **kwargs):
                     conf_attrs, dbname, uid, osv_object, fname
                 ) + args[3:]
                 job_kwargs = kwargs
+                if sudo:
+                    job_kwargs['sudo'] = {
+                        'uid': sudo.uid, 'gid': sudo.gid
+                    }
                 if self.on_commit and async_mode:
                     job = Job.create(
                         execute,
@@ -231,6 +236,10 @@ def f_job(*args, **kwargs):
                         conf_attrs, dbname, uid, osv_object, fname
                     ] + args[3:]
                     job_kwargs = kwargs
+                    if sudo:
+                        job_kwargs['sudo'] = {
+                            'uid': sudo.uid, 'gid': sudo.gid
+                        }
                     at_front = self.at_front
                     if self.on_commit:
                         job = Job.create(

oorq/tasks.py

@@ -14,6 +14,14 @@
 from .utils import get_failed_queue
 
 
+class DummySudo(object):
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
 def make_chunks(ids, n_chunks=None, size=None):
     """Do chunks from ids.
 
@@ -53,6 +61,7 @@ def execute(conf_attrs, dbname, uid, obj, method, *args, **kw):
     import service
     import sql_db
     from ctx import _context_stack
+    from service.security import Sudo
     # Reset the pool with config connections as limit
     sql_db._Pool = sql_db.ConnectionPool(int(tools.config['db_maxconn']))
     osv_ = osv.osv.osv_pool()
@@ -73,7 +82,9 @@ def execute(conf_attrs, dbname, uid, obj, method, *args, **kw):
         return
     if _context_stack.top is None:
         _context_stack.push({})
-    res = osv_.execute(dbname, uid, obj, method, *args, **kw)
+    context = 'sudo' in kw and Sudo(**kw.pop('sudo')) or DummySudo()
+    with context:
+        res = osv_.execute(dbname, uid, obj, method, *args, **kw)
     _context_stack.pop()
     logger.info('Time elapsed: %s' % (datetime.now() - start))
     sql_db.close_db(dbname)
@@ -97,6 +108,7 @@ def isolated_execute(conf_attrs, dbname, uid, obj, method, *args, **kw):
     import workflow
     import report
     import service
+    from service.security import Sudo
     import sql_db
     osv_ = osv.osv.osv_pool()
     pooler.get_db_and_pool(dbname)
@@ -113,11 +125,13 @@ def isolated_execute(conf_attrs, dbname, uid, obj, method, *args, **kw):
     # Ensure args is a list to modify
     args = list(args)
     ids = args[0]
+    context = 'sudo' in kw and Sudo(**kw.pop('sudo')) or DummySudo()
     for exe_id in ids:
         try:
             logger.info('Executing id %s' % exe_id)
             args[0] = [exe_id]
-            res = osv_.execute(dbname, uid, obj, method, *args, **kw)
+            with context:
+                res = osv_.execute(dbname, uid, obj, method, *args, **kw)
             all_res.append(res)
         except:
             logger.error('Executing id %s failed' % exe_id)

oorq/tests/test_oorq/partner.py

@@ -4,6 +4,7 @@
 from oorq.decorators import create_jobs_group
 from oorq.oorq import AsyncMode
 from rq.job import Job
+from service.security import Sudo
 
 
 class ResPartner(osv.osv):
@@ -34,7 +35,7 @@ def test_no_enqueue_on_rollback(self, cursor, uid, ids, vals, context=None):
         self.write_async(cursor, uid, ids, vals, context)
         raise osv.except_osv('Error', 'Test error!')
 
-    @job(async=True, queue='default', on_commit=True)
+    @job(queue='default', on_commit=True)
     def write_async(self, cr, user, ids, vals, context=None):
         #TODO: process before updating resource
         res = super(ResPartner, self).write(cr, user, ids, vals, context)
@@ -65,7 +66,6 @@ def dependency_job(self, cursor, uid, ids, vals, context=None):
 
     @job(queue='dependency', on_commit=True)
     def dependency_job_on_commit(self, cursor, uid, ids, vals, context=None):
-        print "First job"
         import time
         self.write_async(cursor, uid, ids, vals, context=context)
         print("I'm working and not affected for the subjob")
@@ -90,5 +90,11 @@ def write_not_async(self, cursor, uid, ids, values, context=None):
             assert isinstance(res, Job)
         return res.result
 
+    def test_sudo(self, cursor, uid, ids, values, context=None):
+        if context is None:
+            context = {}
+        with Sudo(uid, 'base.group_user'):
+            self.write_async(cursor, uid, ids, values, context)
+
 
 ResPartner()