hobo/hobo/multitenant/threads.py

146 lines
4.6 KiB
Python

# hobo - portal to configure and deploy applications
# Copyright (C) 2019 Entr'ouvert
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# 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/>.
import threading
from django.utils import six
class TenantAwareThread(threading.Thread):
def __init__(self, *args, **kwargs):
from django.db import connection
if hasattr(connection, 'get_tenant'):
self.tenant = connection.get_tenant()
else:
self.tenant = None
super(TenantAwareThread, self).__init__(*args, **kwargs)
def run(self):
if self.tenant is not None:
from django.db import connection
old_tenant = connection.get_tenant()
connection.set_tenant(self.tenant)
try:
super(TenantAwareThread, self).run()
finally:
connection.set_tenant(old_tenant)
connection.close()
else:
super(TenantAwareThread, self).run()
class _Timer(TenantAwareThread):
"""Call a function after a specified number of seconds:
t = Timer(30.0, f, args=[], kwargs={})
t.start()
t.cancel() # stop the timer's action if it's still waiting
"""
def __init__(self, interval, function, args=[], kwargs={}):
super(_Timer, self).__init__()
self.interval = interval
self.function = function
self.args = args
self.kwargs = kwargs
self.finished = threading.Event()
def cancel(self):
"""Stop the timer if it hasn't finished yet"""
self.finished.set()
def run(self):
self.finished.wait(self.interval)
if not self.finished.is_set():
self.function(*self.args, **self.kwargs)
self.finished.set()
class _MainThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_MainThread, self).__init__(name="MainThread", daemon=False)
self._set_tstate_lock()
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
else:
super(_MainThread, self).__init__(name="MainThread")
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return False
def _exitfunc(self):
self._Thread__stop()
t = threading._pickSomeNonDaemonThread()
if t:
if __debug__:
self._note("%s: waiting for other threads", self)
while t:
t.join()
t = threading._pickSomeNonDaemonThread()
if __debug__:
self._note("%s: exiting", self)
self._Thread__delete()
class _DummyThread(TenantAwareThread):
def __init__(self):
if six.PY3:
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"), daemon=True)
self._started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[self._ident] = self
return
super(_DummyThread, self).__init__(name=threading._newname("Dummy-%d"))
# Thread.__block consumes an OS-level locking primitive, which
# can never be used by a _DummyThread. Since a _DummyThread
# instance is immortal, that's bad, so release this resource.
del self._Thread__block
self._Thread__started.set()
self._set_ident()
with threading._active_limbo_lock:
threading._active[threading._get_ident()] = self
def _set_daemon(self):
return True
def join(self, timeout=None):
assert False, "cannot join a dummy thread"
def install_tenant_aware_threads():
if getattr(threading, 'multitenant', False):
return
# Install tenant aware Thread class
threading.Thread = TenantAwareThread
threading._DummyThread = _DummyThread
threading._MainThread = _MainThread
threading._Timer = _Timer
threading.multitenant = True