Produits Entr'ouvert » Hobo

# 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

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'