add uwsgidecorators module (#57019)

hobo/multitenant/uwsgidecorators.py

@@ -0,0 +1,102 @@
+# Copyright (C) 2021 Entr'ouvert
+
+import contextlib
+import logging
+import pickle
+import sys
+
+try:
+    import uwsgi
+except ImportError:
+    uwsgi = None
+
+logger = logging.getLogger(__name__)
+
+spooler_registry = {}
+
+
+@contextlib.contextmanager
+def close_db():
+    if 'django' in sys.modules:
+        from django.db import close_old_connections
+
+        close_old_connections()
+        try:
+            yield None
+        finally:
+            close_old_connections()
+    else:
+        yield
+
+
+@contextlib.contextmanager
+def tenant_context(domain):
+    if domain:
+        from tenant_schemas.utils import tenant_context
+
+        from hobo.multitenant.middleware import TenantMiddleware
+
+        tenant = TenantMiddleware.get_tenant_by_hostname(domain)
+        with tenant_context(tenant):
+            yield
+    else:
+        yield
+
+
+def get_tenant():
+    if 'django.db' not in sys.modules:
+        return ''
+    from django.db import connection
+
+    tenant_model = getattr(connection, 'tenant', None)
+    return getattr(tenant_model, 'domain_url', '')
+
+
+def spool(func):
+    if uwsgi:
+        name = '%s.%s' % (func.__module__, func.__name__)
+        spooler_registry[name] = func
+
+        def spool_function(*args, **kwargs):
+            uwsgi.spool(
+                name=name.encode(),
+                tenant=get_tenant().encode(),
+                body=pickle.dumps({'args': args, 'kwargs': kwargs}),
+            )
+            logger.debug('spooler: spooled function %s', name)
+
+        func.spool = spool_function
+    return func
+
+
+if uwsgi:
+
+    def spooler_function(env):
+        try:
+            try:
+                name = env.get('name').decode()
+                tenant = env.get('tenant', b'').decode()
+                body = env.get('body')
+            except Exception:
+                logger.error('spooler: no name or body found: env.keys()=%s', env.keys())
+                return uwsgi.SPOOL_OK
+            try:
+                params = pickle.loads(body)
+                args = params['args']
+                kwargs = params['kwargs']
+            except Exception:
+                logger.exception('spooler: depickling of body failed')
+                return uwsgi.SPOOL_OK
+            try:
+                function = spooler_registry[name]
+            except KeyError:
+                logger.error('spooler: no function named "%s"', name)
+            # prevent connections to leak between jobs
+            # maintain current tenant when spool is launched
+            with close_db(), tenant_context(tenant):
+                function(*args, **kwargs)
+        except Exception:
+            logger.exception('spooler: function "%s" raised' % name)
+        return uwsgi.SPOOL_OK
+
+    uwsgi.spooler = spooler_function

tests_multitenant/test_uwsgidecorators.py

@@ -0,0 +1,59 @@
+import importlib
+import pickle
+
+import mock
+import pytest
+
+import hobo.multitenant.uwsgidecorators
+
+
+@pytest.fixture
+def uwsgi():
+    import sys
+
+    uwsgi = mock.Mock()
+    uwsgi.SPOOL_OK = -2
+    sys.modules['uwsgi'] = uwsgi
+    importlib.reload(hobo.multitenant.uwsgidecorators)
+    yield uwsgi
+    del sys.modules['uwsgi']
+    importlib.reload(hobo.multitenant.uwsgidecorators)
+
+
+def test_basic():
+    @hobo.multitenant.uwsgidecorators.spool
+    def function(a, b):
+        pass
+
+    function(1, 2)
+    with pytest.raises(AttributeError):
+        function.spool(1, 2)
+
+
+def test_mocked_uwsgi(uwsgi):
+    @hobo.multitenant.uwsgidecorators.spool
+    def function(a, b):
+        pass
+
+    function(1, 2)
+    function.spool(1, 2)
+    assert set(uwsgi.spool.call_args[1].keys()) == {'body', 'tenant', 'name'}
+    assert pickle.loads(uwsgi.spool.call_args[1]['body']) == {'args': (1, 2), 'kwargs': {}}
+    assert uwsgi.spool.call_args[1]['name'] == b'test_uwsgidecorators.function'
+    assert uwsgi.spool.call_args[1]['tenant'] == b''
+
+
+def test_mocked_uwsgi_tenant(uwsgi, tenant):
+    from tenant_schemas.utils import tenant_context
+
+    @hobo.multitenant.uwsgidecorators.spool
+    def function(a, b):
+        pass
+
+    with tenant_context(tenant):
+        function.spool(1, 2)
+
+    assert set(uwsgi.spool.call_args[1].keys()) == {'body', 'tenant', 'name'}
+    assert pickle.loads(uwsgi.spool.call_args[1]['body']) == {'args': (1, 2), 'kwargs': {}}
+    assert uwsgi.spool.call_args[1]['name'] == b'test_uwsgidecorators.function'
+    assert uwsgi.spool.call_args[1]['tenant'] == b'tenant.example.net'