hobo/hobo/multitenant/threads.py

65 lines
1.8 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 functools
import threading
_Thread_start = threading.Thread.start
_Thread__bootstrap_inner = threading.Thread._bootstrap_inner
def _new_start(self):
from django.db import connection
tenant = getattr(connection, 'tenant', None)
self.tenant = tenant
return _Thread_start(self)
def wrap_run(self, func):
if getattr(func, '_wrapped', False):
return func
@functools.wraps(func)
def wrapper():
tenant = getattr(self, 'tenant', None)
if tenant is not None:
from django.db import connection
old_tenant = connection.tenant
connection.set_tenant(self.tenant)
try:
func()
finally:
connection.set_tenant(old_tenant)
connection.close()
else:
func()
wrapper._wrapped = True
return wrapper
def _new__bootstrap_inner(self):
self.run = wrap_run(self, self.run)
_Thread__bootstrap_inner(self)
def install_tenant_aware_threads():
threading.Thread.start = _new_start
threading.Thread._bootstrap_inner = _new__bootstrap_inner