spooler: add decorator to simplify use of spooler (#52221)

The new tenantspool decorator keep the function signature using pickle
serialization (through the pass_arguments=True parameter of @spool),
setup and teardown Django db connections and restore the current tenant.

combo/apps/dataviz/forms.py

@@ -19,10 +19,10 @@ from collections import OrderedDict
 
 from django import forms
 from django.conf import settings
-from django.db import transaction, connection
+from django.db import transaction
 from django.db.models import Q
 
-from combo.utils import requests, cache_during_request
+from combo.utils import requests, cache_during_request, spooler
 
 from .models import ChartCell, ChartNgCell
 
@@ -49,13 +49,7 @@ class ChartForm(forms.ModelForm):
 
 @cache_during_request
 def trigger_statistics_list_refresh():
-    if 'uwsgi' in sys.modules:
-        from combo.utils.spooler import refresh_statistics_list
-
-        tenant = getattr(connection, 'tenant', None)
-        transaction.on_commit(
-            lambda: refresh_statistics_list.spool(domain=getattr(tenant, 'domain_url', None))
-        )
+    transaction.on_commit(spooler.refresh_statistics_list)
 
 
 class ChartNgForm(forms.ModelForm):

combo/apps/dataviz/models.py

@@ -21,7 +21,7 @@ from datetime import date
 from requests.exceptions import RequestException
 
 from django.urls import reverse
-from django.db import models, connection, transaction
+from django.db import models, transaction
 from django.utils import timezone
 from django.utils.encoding import force_text
 from django.utils.translation import ugettext_lazy as _, ungettext, gettext
@@ -34,7 +34,7 @@ import pygal.util
 
 from combo.data.models import CellBase
 from combo.data.library import register_cell_class
-from combo.utils import get_templated_url, requests
+from combo.utils import get_templated_url, requests, spooler
 
 
 class UnsupportedDataSet(Exception):
@@ -295,15 +295,7 @@ class ChartNgCell(CellBase):
         )
 
     def get_chart(self, width=None, height=None, raise_if_not_cached=False):
-        if 'uwsgi' in sys.modules:
-            from combo.utils.spooler import refresh_statistics_data
-
-            tenant = getattr(connection, 'tenant', None)
-            transaction.on_commit(
-                lambda: refresh_statistics_data.spool(
-                    cell_pk=str(self.pk), domain=getattr(tenant, 'domain_url', None)
-                )
-            )
+        transaction.on_commit(lambda: spooler.refresh_statistics_data(cell_pk=self.pk))
         response = self.get_statistic_data(raise_if_not_cached)
         response.raise_for_status()
         response = response.json()

combo/utils/spooler.py

@@ -14,36 +14,88 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from django.db import connection
+import contextlib
+from functools import wraps
+import logging
+import sys
 
-from uwsgidecorators import spool
+from django.db import close_old_connections, connection
 
-from combo.apps.dataviz.models import ChartNgCell
-from combo.apps.dataviz.utils import update_available_statistics
-from combo.utils import requests
+USE_UWSGI = 'uwsgi' in sys.modules
 
 
-def set_connection(domain):
+logger = logging.getLogger(__name__)
+
+
+def ensure_db(func):
+    """Emulate Django"s setup/teardown of database connections before/after
+    each request"""
+
+    @wraps(func)
+    def f(*args, **kwargs):
+        close_old_connections()
+        try:
+            return func(*args, **kwargs)
+        finally:
+            close_old_connections()
+
+    return f
+
+
+@contextlib.contextmanager
+def tenant_context(domain):
     from hobo.multitenant.middleware import TenantMiddleware
+    from tenant_schemas.utils import tenant_context
 
     tenant = TenantMiddleware.get_tenant_by_hostname(domain)
-    connection.set_tenant(tenant)
+    with tenant_context(tenant):
+        yield
 
 
-@spool
-def refresh_statistics_list(args):
-    if args.get('domain'):
-        # multitenant installation
-        set_connection(args['domain'])
+def tenantspool(func):
+    """Wrap a function with uwsgidecorators.spool storing and restoring the
+    current tenant."""
+    if not USE_UWSGI:
+        return func
+
+    from uwsgidecorators import spool
+
+    @ensure_db
+    @wraps(func)
+    def spooler_func(*args, **kwargs):
+        with contextlib.ExitStack() as stack:
+            if 'domain' in kwargs:
+                stack.enter_context(tenant_context(kwargs.pop('domain')))
+            try:
+                func(*args, **kwargs)
+            except Exception:
+                logger.exception('spooler: exception during %s(%s, %s)' % (func.__name__, args, kwargs))
+            else:
+                logger.info('spooler: success of %s(%s, %s)' % (func.__name__, args, kwargs))
+
+    # pass arguments as pickles
+    base_spooler = spool(pass_arguments=True)(spooler_func)
+
+    @wraps(func)
+    def spooler(*args, **kwargs):
+        domain = getattr(getattr(connection, 'tenant', None), 'domain_url', None)
+        if domain is not None:
+            kwargs['domain'] = domain
+        return base_spooler(*args, **kwargs)
+
+    return spooler
+
+
+@tenantspool
+def refresh_statistics_list():
+    from combo.apps.dataviz.utils import update_available_statistics
 
     update_available_statistics()
 
 
-@spool
-def refresh_statistics_data(args):
-    if args.get('domain'):
-        # multitenant installation
-        set_connection(args['domain'])
+@tenantspool
+def refresh_statistics_data(cell_pk):
+    from combo.apps.dataviz.models import ChartNgCell
 
-    cell = ChartNgCell.objects.get(pk=args['cell_pk'])
+    cell = ChartNgCell.objects.get(pk=cell_pk)
     cell.get_statistic_data(invalidate_cache=True)

tests/test_spooler.py

@@ -0,0 +1,52 @@
+import importlib
+import pickle
+import sys
+
+import mock
+import pytest
+
+from combo.utils import spooler
+
+
+@pytest.fixture
+def uwsgi(monkeypatch):
+    uwsgi = mock.Mock()
+    uwsgi.opt = ['spooler']
+    uwsgi.SPOOL_OK = 1
+    monkeypatch.setitem(sys.modules, 'uwsgi', uwsgi)
+    # reload with uwsgi impl
+    importlib.reload(spooler)
+    yield uwsgi
+    monkeypatch.delitem(sys.modules, 'uwsgi')
+    importlib.reload(spooler)
+
+
+def test_tenantspool(uwsgi):
+    import uwsgidecorators
+
+    assert spooler.USE_UWSGI
+
+    mocked_f = mock.Mock()
+
+    @spooler.tenantspool
+    def function_to_spool(*args, **kwargs):
+        mocked_f(*args, **kwargs)
+
+    # check function is registered under its name, and not the name of the wrapper function
+    assert 'function_to_spool' in uwsgidecorators.spooler_functions
+
+    # check the spool request is properly serialized with tenant information
+    with mock.patch('combo.utils.spooler.connection') as connection:
+        connection.tenant.domain_url = 'example.com'
+        function_to_spool(1, b=3)
+    assert pickle.loads(uwsgi.spool.call_args[0][0][b'args']) == (1,)
+    assert pickle.loads(uwsgi.spool.call_args[0][0][b'kwargs']) == {'b': 3, 'domain': 'example.com'}
+
+    # check db is initialized, tenant restored and function called
+    with mock.patch('combo.utils.spooler.close_old_connections') as close_old_connections:
+        with mock.patch('combo.utils.spooler.tenant_context') as tenant_context:
+            uwsgidecorators.manage_spool_request(uwsgi.spool.call_args[0][0])
+            assert close_old_connections.call_count == 2
+            assert tenant_context.call_args[0][0] == 'example.com'
+            assert mocked_f.call_args[0] == (1,)
+            assert mocked_f.call_args[1] == {'b': 3}

tox.ini

@@ -33,6 +33,7 @@ deps =
   django-ratelimit<3
   git+https://git.entrouvert.org/debian/django-ckeditor.git
   black: pre-commit
+  uwsgidecorators
 commands =
   ./getlasso3.sh
   python manage.py compilemessages