from __future__ import absolute_import

try:
    import celery
    from celery import Celery
except ImportError:
    raise ImportError("celery is required to use tenant_schemas_celery")

from django.db import connection

from celery.signals import task_prerun, task_postrun


def get_schema_name_from_task(task, kwargs):
    if celery.VERSION[0] < 4:
        # Pop it from the kwargs since tasks don't except the additional kwarg.
        # This change is transparent to the system.
        return kwargs.pop("_schema_name", None)

    # In some cases (like Redis broker) headers are merged with `task.request`.
    if task.request.headers and "_schema_name" in task.request.headers:
        return task.request.headers.get("_schema_name")
    return task.request.get("_schema_name")


def switch_schema(task, kwargs, **kw):
    """ Switches schema of the task, before it has been run. """
    # Lazily load needed functions, as they import django model functions which
    # in turn load modules that need settings to be loaded and we can't
    # guarantee this module was loaded when the settings were ready.
    from .compat import get_public_schema_name

    old_schema = (connection.schema_name, connection.include_public_schema)
    setattr(task, "_old_schema", old_schema)

    schema = get_schema_name_from_task(task, kwargs) or get_public_schema_name()

    # If the schema has not changed, don't do anything.
    if connection.schema_name == schema:
        return

    if connection.schema_name != get_public_schema_name():
        connection.set_schema_to_public()

    if schema == get_public_schema_name():
        return

    tenant = task.get_tenant_for_schema(schema_name=schema)
    connection.set_tenant(tenant, include_public=True)


def restore_schema(task, **kwargs):
    """ Switches the schema back to the one from before running the task. """
    from .compat import get_public_schema_name

    schema_name = get_public_schema_name()
    include_public = True

    if hasattr(task, "_old_schema"):
        schema_name, include_public = task._old_schema

    # If the schema names match, don't do anything.
    if connection.schema_name == schema_name:
        return

    connection.set_schema(schema_name, include_public=include_public)


task_prerun.connect(
    switch_schema, sender=None, dispatch_uid="tenant_schemas_switch_schema"
)

task_postrun.connect(
    restore_schema, sender=None, dispatch_uid="tenant_schemas_restore_schema"
)


class CeleryApp(Celery):
    registry_cls = 'tenant_schemas_celery.registry:TenantTaskRegistry'
    
    def create_task_cls(self):
        return self.subclass_with_self(
            "tenant_schemas_celery.task:TenantTask",
            abstract=True,
            name="TenantTask",
            attribute="_app",
        )

    def _update_headers(self, kw):
        kw["headers"] = kw.get("headers") or {}
        self._add_current_schema(kw["headers"])

    def _add_current_schema(self, kwds):
        kwds["_schema_name"] = kwds.get("_schema_name", connection.schema_name)

    def send_task(self, name, args=None, kwargs=None, **options):
        if celery.VERSION[0] < 4:
            kwargs = kwargs or {}
            self._add_current_schema(kwargs)

        else:
            # Celery 4.0 introduced strong typing and the `headers` meta dict.
            self._update_headers(options)
        return super(CeleryApp, self).send_task(name, args=args, kwargs=kwargs, **options)

@@ -212,16 +212,29 @@
 nnection
+, connections
 %0A%0Afrom c
@@ -1545,33 +1545,70 @@
 ():%0A        
-connectio
+for conn in connections.all():%0A            con
 n.set_schema
@@ -1740,33 +1740,66 @@
 schema)%0A    
-connectio
+for conn in connections.all():%0A        con
 n.set_tenant
@@ -2273,25 +2273,58 @@
 rn%0A%0A    
-connectio
+for conn in connections.all():%0A        con
 n.set_sc

Add support for django with multiple databases