from __future__ import unicode_literals

from celery import Celery
from mongoengine import connect
from celery.signals import celeryd_init, worker_process_init

from core.config.config import yeti_config


celery_app = Celery("yeti")


class CeleryConfig:
    redis_scheme = "redis"
    if yeti_config.redis.tls:

        redis_scheme = redis_scheme + "s"

    BROKER_URL = "{}://{}:{}/{}".format(
        redis_scheme,
        yeti_config.redis.host,
        yeti_config.redis.port,
        yeti_config.redis.database,
    )
    CELERY_TASK_SERIALIZER = "json"
    CELERY_ACCEPT_CONTENT = ["json"]
    CELERY_IMPORTS = (
        "core.config.celeryimports",
        "core.analytics_tasks",
        "core.exports.export",
        "core.feed",
        "core.investigation",
        "plugins",
    )
    CELERY_TIMEZONE = "UTC"
    CELERYD_POOL_RESTARTS = True
    CELERY_ROUTES = {
        "core.analytics_tasks.single": {"queue": "oneshot"},
        "core.feed.update_feed": {"queue": "feeds"},
        "core.exports.export.execute_export": {"queue": "exports"},
        "core.analytics_tasks.each": {"queue": "analytics"},
        "core.analytics_tasks.schedule": {"queue": "analytics"},
        "core.investigation.import_task": {"queue": "oneshot"},
    }


celery_app.config_from_object(CeleryConfig)


@worker_process_init.connect
def connect_mongo(**kwargs):
    """Connect to mongo and load modules in each worker process."""
    from core.config import celeryimports
    from core.yeti_plugins import get_plugins

    is_tls = False
    if yeti_config.mongodb.tls:
        is_tls = True

    connect(
        yeti_config.mongodb.database,
        host=yeti_config.mongodb.host,
        port=yeti_config.mongodb.port,
        username=yeti_config.mongodb.username,
        password=yeti_config.mongodb.password,
        tls=is_tls,
        connect=False,
    )
    celeryimports.loaded_modules = get_plugins()


@celeryd_init.connect
def unlock_scheduled_entries(**kwargs):
    from core.analytics import ScheduledAnalytics
    from core.feed import Feed
    from core.exports.export import Export
    from core.yeti_plugins import get_plugin_classes

    # Make sure plugin classes are loaded
    get_plugin_classes()

    locked_entries = {
        "analytics": ScheduledAnalytics,
        "exports": Export,
        "feeds": Feed,
    }

    queues = kwargs["options"]["queues"].split(",")

    for queue in queues:
        if queue in locked_entries:
            locked_entries[queue].objects(lock=True).update(
                lock=False, status="Unlocked"
            )

@@ -64,40 +64,8 @@
 ery%0A
-from mongoengine import connect%0A
 from
@@ -121,17 +121,16 @@
 ss_init%0A
-%0A
 from cor
@@ -164,16 +164,154 @@
 _config%0A
+from mongoengine.connection import (DEFAULT_CONNECTION_NAME, connect,%0A                                    disconnect, register_connection)
 %0A%0Acelery
@@ -1409,227 +1409,8 @@
 )%0A%0A%0A
-@worker_process_init.connect%0Adef connect_mongo(**kwargs):%0A    %22%22%22Connect to mongo and load modules in each worker process.%22%22%22%0A    from core.config import celeryimports%0A    from core.yeti_plugins import get_plugins%0A%0A    
 is_t
@@ -1420,20 +1420,16 @@
 = False%0A
-    
 if yeti_
@@ -1448,20 +1448,16 @@
 db.tls:%0A
-    
     is_t
@@ -1471,29 +1471,60 @@
 ue%0A%0A
-    connect(
+register_connection(DEFAULT_CONNECTION_NAME,
 %0A    
-    
+db=
 yeti
@@ -1553,20 +1553,16 @@
 se,%0A    
- 
-   
 host=yet
@@ -1588,20 +1588,16 @@
 st,%0A    
- 
-   
 port=yet
@@ -1615,28 +1615,24 @@
 ngodb.port,%0A
-    
     username
@@ -1666,20 +1666,16 @@
 me,%0A    
-  
-  
 password
@@ -1705,20 +1705,16 @@
 ssword,%0A
-    
     tls=
@@ -1725,20 +1725,16 @@
 ls,%0A    
-   
- 
 connect=
@@ -1736,29 +1736,288 @@
 nect=False,%0A
-    )
+)%0A%0A%0A@worker_process_init.connect%0Adef connect_mongo(**kwargs):%0A    %22%22%22Connect to mongo and load modules in each worker process.%22%22%22%0A    from core.config import celeryimports%0A    from core.yeti_plugins import get_plugins%0A%0A    connect(db=yeti_config.mongodb.database)%0A
 %0A    celeryi
@@ -2172,39 +2172,8 @@
 ics%0A
-    from core.feed import Feed%0A
     
@@ -2211,16 +2211,47 @@
  Export%0A
+    from core.feed import Feed%0A
     from
@@ -2296,16 +2296,62 @@
 lasses%0A%0A
+    connect(db=yeti_config.mongodb.database)%0A%0A
     # Ma
@@ -2762,12 +2762,30 @@
           )%0A
+%0A    disconnect()%0A

fix celery connection issue