CWTask: add a cnx() context manager to handle CW connection

Handle user and atomic transaction through a single context manager. Deprecate cw_user_cnx() and need_cnx attribute. This is convenient to get a connection on all tasks handlers (on_failure, on_success, ...).

authorPhilippe Pepiot <philippe.pepiot@logilab.fr>
changeset0562b43f1944
branchdefault
phasepublic
hiddenno
parent revision#117ba0fe2252 Move to package layout
child revision#0f0371baa444 [pkg] 0.3.0, #0010193c23ae [pkg] 0.3.0, #fb0ec92bdea1 [pkg] 0.3.0, #ce44788c1cce [pkg] 0.3.0
files modified by this revision
cubicweb_celery/__init__.py
# HG changeset patch
# User Philippe Pepiot <philippe.pepiot@logilab.fr>
# Date 1477397533 -7200
# Tue Oct 25 14:12:13 2016 +0200
# Node ID 0562b43f194489c8a990ec81cad515ab8f213800
# Parent 117ba0fe2252f6e228894bdc1a1fd4f3e0e8076c
CWTask: add a cnx() context manager to handle CW connection

Handle user and atomic transaction through a single context manager.
Deprecate cw_user_cnx() and need_cnx attribute.
This is convenient to get a connection on all tasks handlers (on_failure, on_success, ...).

diff --git a/cubicweb_celery/__init__.py b/cubicweb_celery/__init__.py
@@ -4,13 +4,16 @@
1  """
2 
3  import contextlib
4  import logging
5  import os.path as osp
6 +import warnings
7 
8  from optparse import make_option
9 
10 +from logilab.common.deprecation import deprecated
11 +
12  from celery import Celery
13  from celery import signals, task
14 
15  from cubicweb.server.session import Session
16  from cubicweb.cwconfig import CubicWebConfiguration as cwcfg
@@ -22,16 +25,41 @@
17  class CWTask(task.Task):
18      need_cnx = False
19      abstract = True
20 
21      @contextlib.contextmanager
22 -    def cw_user_cnx(self, user_eid):
23 -        """context manager that provide a a user connection
24 +    def _cnx(self, user=None):
25 +        if user == -1:
26 +            user = None
27 +        with self.app.cwrepo.internal_cnx() as cnx:
28 +            if user is not None:
29 +                if not hasattr(user, "eid"):
30 +                    user = cnx.entity_from_eid(user)
31 +                with Session(user, cnx.repo).new_cnx() as user_cnx:
32 +                    yield user_cnx
33 +            else:
34 +                yield cnx
35 
36 -        Yield a connection for `user_eid` or an internal connection if given
37 -        `user_eid` is None or -1
38 -        """
39 +    @contextlib.contextmanager
40 +    def cnx(self, user=None, atomic=True):
41 +        """context manager that provide an connection as the specified user"""
42 +        with self._cnx(user=user) as cnx:
43 +            if atomic:
44 +                try:
45 +                    yield cnx
46 +                except Exception:
47 +                    cnx.rollback()
48 +                    raise
49 +                else:
50 +                    cnx.commit()
51 +            else:
52 +                yield cnx
53 +
54 +    @contextlib.contextmanager
55 +    @deprecated('use cnx() instead')
56 +    def cw_user_cnx(self, user_eid):
57 +        """DEPRECATED use cnx() instead"""
58          if not self.need_cnx:
59              raise ValueError("use of cw_user_cnx() requires need_cnx")
60          cnx = self.cw_cnx
61          if user_eid in (None, -1):
62              yield cnx
@@ -40,20 +68,17 @@
63              with Session(user, cnx.repo).new_cnx() as user_cnx:
64                  yield user_cnx
65 
66      def __call__(self, *args, **kwargs):
67          if self.need_cnx:
68 -            with self.app.cwrepo.internal_cnx() as cnx:
69 +            warnings.warn('need_cnx attribute should not be specified, use '
70 +                          'cnx() instead',
71 +                          DeprecationWarning, stacklevel=2)
72 +            with self.cnx() as cnx:
73                  try:
74                      self.cw_cnx = cnx
75 -                    ret = super(CWTask, self).__call__(*args, **kwargs)
76 -                except:
77 -                    cnx.rollback()
78 -                    raise
79 -                else:
80 -                    cnx.commit()
81 -                    return ret
82 +                    return super(CWTask, self).__call__(*args, **kwargs)
83                  finally:
84                      self.cw_cnx = None
85          else:
86              return super(CWTask, self).__call__(*args, **kwargs)
87