Merge branch 'wgen/master'

celery/backends/mongodb.py

@@ -178,6 +178,22 @@ class MongoBackend(BaseDictBackend):
         taskmeta_collection = db[self.mongodb_taskmeta_collection]
         taskmeta_collection.remove({"_id": taskset_id})
 
+    def _forget(self, task_id):
+        """
+        Remove result from MongoDB.
+        
+        :raises celery.exceptions.OperationsError: if the task_id could not be
+                                                   removed.
+        """
+
+        db = self._get_database()
+        taskmeta_collection = db[self.mongodb_taskmeta_collection]
+        
+        # By using safe=True, this will wait until it receives a response from
+        # the server.  Likewise, it will raise an OperationsError if the
+        # response was unable to be completed.
+        taskmeta_collection.remove({"_id": task_id}, safe=True)
+
     def cleanup(self):
         """Delete expired metadata."""
         db = self._get_database()

celery/tests/test_backends/test_mongodb.py

@@ -1,64 +1,257 @@
-from __future__ import absolute_import
-
-import sys
+import uuid
 
+from mock import MagicMock, Mock, patch, sentinel
 from nose import SkipTest
 
+from celery import states
 from celery.backends.mongodb import MongoBackend
-from celery.exceptions import ImproperlyConfigured
 from celery.tests.utils import unittest
-from celery.utils import uuid
-
-
-_no_mongo_msg = "* MongoDB %s. Will not execute related tests."
-_no_mongo_msg_emitted = False
 
 
 try:
-    from pymongo.errors import AutoReconnect
+    import pymongo
 except ImportError:
+    pymongo = None
+
+
+COLLECTION = "taskmeta_celery"
+TASK_ID = str(uuid.uuid1())
+MONGODB_HOST = "localhost"
+MONGODB_PORT = 27017
+MONGODB_USER = "mongo"
+MONGODB_PASSWORD = "1234"
+MONGODB_DATABASE = "testing"
+MONGODB_COLLECTION = "collection1"
+
+
+@patch("celery.backends.mongodb.MongoBackend.decode", Mock())
+@patch("celery.backends.mongodb.MongoBackend.encode", Mock())
+@patch("pymongo.binary.Binary", Mock())
+@patch("datetime.datetime", Mock())
+class TestBackendMongoDb(unittest.TestCase):
+
+    def setUp(self):
+        if pymongo is None:
+            raise SkipTest("pymongo is not installed.")
+
+        self.backend = MongoBackend()
+
+    @patch("pymongo.connection.Connection")
+    def test_get_connection_connection_exists(self, mock_Connection):
+        self.backend._connection = sentinel._connection
+
+        connection = self.backend._get_connection()
+
+        self.assertEquals(sentinel._connection, connection)
+        self.assertFalse(mock_Connection.called)
+
+    @patch("pymongo.connection.Connection")
+    def test_get_connection_no_connection_host(self, mock_Connection):
+        self.backend._connection = None
+        self.backend.mongodb_host = MONGODB_HOST
+        self.backend.mongodb_port = MONGODB_PORT
+        mock_Connection.return_value = sentinel.connection
+
+        connection = self.backend._get_connection()
+        mock_Connection.assert_called_once_with(
+            MONGODB_HOST, MONGODB_PORT)
+        self.assertEquals(sentinel.connection, connection)
+
+    @patch("pymongo.connection.Connection")
+    def test_get_connection_no_connection_mongodb_uri(self, mock_Connection):
+        mongodb_uri = "mongodb://%s:%d" % (MONGODB_HOST, MONGODB_PORT)
+        self.backend._connection = None
+        self.backend.mongodb_host = mongodb_uri
+
+        mock_Connection.return_value = sentinel.connection
+
+        connection = self.backend._get_connection()
+        mock_Connection.assert_called_once_with(mongodb_uri)
+        self.assertEquals(sentinel.connection, connection)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_connection")
+    def test_get_database_no_existing(self, mock_get_connection):
+        # Should really check for combinations of these two, to be complete.
+        self.backend.mongodb_user = MONGODB_USER
+        self.backend.mongodb_password = MONGODB_PASSWORD
+
+        mock_database = Mock()
+        mock_connection = MagicMock(spec=['__getitem__'])
+        mock_connection.__getitem__.return_value = mock_database
+        mock_get_connection.return_value = mock_connection
+
+        database = self.backend._get_database()
+
+        self.assertTrue(database is mock_database)
+        self.assertTrue(self.backend._database is mock_database)
+        mock_database.authenticate.assert_called_once_with(
+            MONGODB_USER, MONGODB_PASSWORD)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_connection")
+    def test_get_database_no_existing_no_auth(self, mock_get_connection):
+        # Should really check for combinations of these two, to be complete.
+        self.backend.mongodb_user = None
+        self.backend.mongodb_password = None
+
+        mock_database = Mock()
+        mock_connection = MagicMock(spec=['__getitem__'])
+        mock_connection.__getitem__.return_value = mock_database
+        mock_get_connection.return_value = mock_connection
+
+        database = self.backend._get_database()
+
+        self.assertTrue(database is mock_database)
+        self.assertFalse(mock_database.authenticate.called)
+        self.assertTrue(self.backend._database is mock_database)
+
+    def test_process_cleanup(self):
+        self.backend._connection = None
+        self.backend.process_cleanup()
+        self.assertEquals(self.backend._connection, None)
+
+        self.backend._connection = "not none"
+        self.backend.process_cleanup()
+        self.assertEquals(self.backend._connection, None)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_store_result(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
+
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+
+        ret_val = self.backend._store_result(
+            sentinel.task_id, sentinel.result, sentinel.status)
+
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        mock_collection.save.assert_called_once()
+        self.assertEquals(sentinel.result, ret_val)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_get_task_meta_for(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
+        mock_collection.find_one.return_value = MagicMock()
+
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+
+        ret_val = self.backend._get_task_meta_for(sentinel.task_id)
+
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        self.assertEquals(
+            ['status', 'date_done', 'traceback', 'result', 'task_id'],
+            ret_val.keys())
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_get_task_meta_for_no_result(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
+        mock_collection.find_one.return_value = None
+
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+
+        ret_val = self.backend._get_task_meta_for(sentinel.task_id)
+
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        self.assertEquals({"status": states.PENDING, "result": None}, ret_val)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_save_taskset(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
+
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+
+        ret_val = self.backend._save_taskset(
+            sentinel.taskset_id, sentinel.result)
+
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        mock_collection.save.assert_called_once()
+        self.assertEquals(sentinel.result, ret_val)
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_restore_taskset(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
+        mock_collection.find_one.return_value = MagicMock()
+
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+
+        ret_val = self.backend._restore_taskset(sentinel.taskset_id)
+
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        mock_collection.find_one.assert_called_once_with(
+            {"_id": sentinel.taskset_id})
+        self.assertEquals(['date_done', 'result', 'task_id'], ret_val.keys())
+
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_delete_taskset(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
+
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
 
-    class AutoReconnect(Exception):  # noqa
-        pass
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
 
+        ret_val = self.backend._delete_taskset(sentinel.taskset_id)
 
-def get_mongo_or_SkipTest():
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(MONGODB_COLLECTION)
+        mock_collection.remove.assert_called_once_with(
+            {"_id": sentinel.taskset_id})
 
-    def emit_no_mongo_msg(reason):
-        global _no_mongo_msg_emitted
-        if not _no_mongo_msg_emitted:
-            sys.stderr.write("\n" + _no_mongo_msg % reason + "\n")
-            _no_mongo_msg_emitted = True
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_forget(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
 
-    try:
-        tb = MongoBackend()
-        try:
-            tb._get_database()
-        except AutoReconnect, exc:
-            emit_no_mongo_msg("not running")
-            raise SkipTest("Can't connect to MongoDB: %s" % (exc, ))
-        return tb
-    except ImproperlyConfigured, exc:
-        if "need to install" in str(exc):
-            emit_no_mongo_msg("pymongo not installed")
-            raise SkipTest("pymongo not installed")
-        emit_no_mongo_msg("not configured")
-        raise SkipTest("MongoDB not configured correctly: %s" % (exc, ))
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
 
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
 
-class TestMongoBackend(unittest.TestCase):
+        self.backend._forget(sentinel.task_id)
 
-    def test_save__restore__delete_taskset(self):
-        tb = get_mongo_or_SkipTest()
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(
+            MONGODB_COLLECTION)
+        mock_collection.remove.assert_called_once_with(
+            {"_id": sentinel.task_id}, safe=True)
 
-        tid = uuid()
-        res = {u"foo": "bar"}
-        self.assertEqual(tb.save_taskset(tid, res), res)
+    @patch("celery.backends.mongodb.MongoBackend._get_database")
+    def test_cleanup(self, mock_get_database):
+        self.backend.mongodb_taskmeta_collection = MONGODB_COLLECTION
 
-        res2 = tb.restore_taskset(tid)
-        self.assertEqual(res2, res)
+        mock_database = MagicMock(spec=['__getitem__', '__setitem__'])
+        mock_collection = Mock()
 
-        tb.delete_taskset(tid)
-        self.assertIsNone(tb.restore_taskset(tid))
+        mock_get_database.return_value = mock_database
+        mock_database.__getitem__.return_value = mock_collection
+        
+        self.backend.cleanup()
 
-        self.assertIsNone(tb.restore_taskset("xxx-nonexisting-id"))
+        mock_get_database.assert_called_once_with()
+        mock_database.__getitem__.assert_called_once_with(
+            MONGODB_COLLECTION)
+        mock_collection.assert_called_once()

contrib/centos/celeryd.init

@@ -1,152 +1,142 @@
-#! /bin/sh
+#!/bin/sh
+### BEGIN INIT INFO
+# Provides: celeryd
+# Required-Start: $network $local_fs $remote_fs
+# Required-Stop: $network $local_fs $remote_fs
+# Default-Start: 2 3 4 5
+# Default-Stop: 0 1 6
+# Short-Description: celery task worker daemon
+### END INIT INFO
 #
-# chkconfig: - 85 15
-# description:  Celery worker daemon
-# processname: celeryd
-# config: /etc/sysconfig/celeryd
-# pidfile: /var/run/celeryd.pid
+# ============================================
+# celeryd - Starts the Celery worker daemon.
+# ============================================
 #
-# To configure celeryd you probably need to tell it where to chdir.
+# :Usage: /etc/init.d/${basename $0} {start|stop|restart|status}
+# :Configuration file: /etc/default/celeryd
 #
-# EXAMPLE CONFIGURATION
-# =====================
+# To implement separate init scripts, do NOT copy this script.  Instead,
+# symlink it.  I.e., if my new application, "little-worker" needs an init, I
+# should just use:
 #
-# this is an example configuration for a Python project:
+#   ln -s /etc/init.d/celeryd /etc/init.d/little-worker
 #
-# /etc/sysconfig/celeryd:
+# You can then configure this by manipulating /etc/sysconfig/little-worker.  It
+# will still read defaults from /etc/defaults/celeryd, but everything can be
+# overriden by sysconfig.
 #
-#   # Where to chdir at start.
-#   CELERYD_CHDIR="/opt/Myproject/"
-#
-#   # Extra arguments to celeryd
-#   CELERYD_OPTS="--time-limit=300"
-#
-#   # Name of the celery config module.#
-#   CELERY_CONFIG_MODULE="celeryconfig"
-#
-# EXAMPLE DJANGO CONFIGURATION
-# ============================
-#
-#   # Where the Django project is.
-#   CELERYD_CHDIR="/opt/Project/"
-#
-#   # Name of the projects settings module.
-#   export DJANGO_SETTINGS_MODULE="settings"
-#
-#   # Path to celeryd
-#   CELERYD="/opt/Project/manage.py"
-#
-#   # Extra arguments to manage.py
-#   CELERYD_OPTS="celeryd"
-#
-# AVAILABLE OPTIONS
-# =================
-#
-#   * CELERYD_OPTS
-#       Additional arguments to celeryd, see `celeryd --help` for a list.
-#
-#   * CELERYD_CHDIR
-#       Path to chdir at start. Default is to stay in the current directory.
-#
-#   * CELERYD_PID_FILE
-#       Full path to the pidfile. Default is /var/run/celeryd.pid.
-#
-#   * CELERYD_LOG_FILE
-#       Full path to the celeryd logfile. Default is /var/log/celeryd.log
-#
-#   * CELERYD_LOG_LEVEL
-#       Log level to use for celeryd. Default is INFO.
-#
-#   * CELERYD
-#       Path to the celeryd program. Default is `celeryd`.
-#       You can point this to an virtualenv, or even use manage.py for django.
-#
-#   * CELERYD_USER
-#       User to run celeryd as. Default is current user.
-#
-#   * CELERYD_GROUP
-#       Group to run celeryd as. Default is current user.
-#
-#   * VIRTUALENV
-#       Full path to the virtualenv environment to activate. Default is none.
-#
-#   * PYTHONPATH
-#       A directory to add to the Python path.
 
-# Source function library.
+# Setting `prog` here allows you to symlink this init script, making it easy
+# to run multiple processes on the system.
+prog="$(basename $0)"
+
+# Source the centos stuff
 . /etc/init.d/functions
 
-if test -f /etc/sysconfig/celeryd; then
-    . /etc/sysconfig/celeryd
-fi
+# Also look at sysconfig; this is where environmental variables should be set
+# on RHEL systems.
+[ -f "/etc/sysconfig/$prog" ] && . /etc/sysconfig/$prog
 
-RETVAL=0
+CELERYD=${CELERYD:-"-m celery.bin.celeryd_detach"}
+CELERYD_MULTI=${CELERYD_MULTI:-"/usr/bin/celeryd-multi"}
+CELERYD_PID_FILE=${CELERYD_PID_FILE:-"/var/run/celeryd/$prog.pid"}
+CELERYD_LOG_FILE=${CELERYD_LOG_FILE:-"/var/log/celeryd/$prog.log"}
+CELERYD_LOG_LEVEL=${CELERYD_LOG_LEVEL:-"INFO"}
 
-DEFAULT_CELERYD="/usr/bin/celeryd"
-CELERYD_LOG_FILE=${CELERYD_LOG_FILE:-${CELERYD_LOGFILE:-"/var/log/celeryd.log"}}
-CELERYD_PID_FILE=${CELERYD_PID_FILE:-${CELERYD_PIDFILE:-"/var/run/celeryd.pid"}}
-CELERYD_LOG_LEVEL=${CELERYD_LOG_LEVEL:-${CELERYD_LOGLEVEL:-"INFO"}}
-CELERYD_USER=${CELERYD_USER:-${CELERYD_USER:-"celeryd"}}
+# This is used to change how Celery loads in the configs.  It does not need to
+# be set to be run.
+export CELERY_LOADER
 
-CELERYD=${CELERYD:-$DEFAULT_CELERYD}
+start_workers () {
+    CELERYD_LOG_DIR=$(dirname $CELERYD_LOG_FILE)
+    CELERYD_PID_DIR=$(dirname $CELERYD_PID_FILE)
+    # Ensure that the directories exist.
+    mkdir -p $CELERYD_LOG_DIR $CELERYD_PID_DIR
 
-export CELERY_LOADER
+    # If we specified a user, and/or a group, chown as needed
+    if [ -n "$CELERYD_USER" ]; then
+        CHOWN_UG="${CELERYD_USER}"
 
-CELERYD_OPTS="$CELERYD_OPTS -f $CELERYD_LOG_FILE -l $CELERYD_LOG_LEVEL --pidfile=$CELERYD_PID_FILE"
+        # If the group is specified, also use that in the chown.
+        [ -n "$CELERYD_GROUP" ] && CHOWN_UG="$CHOWN_UG:$CELERYD_GROUP"
 
-if [ -n "$2" ]; then
-    CELERYD_OPTS="$CELERYD_OPTS $2"
-fi
+        # Execute the chown on the directory only
+        chown $CHOWN_UG $CELERYD_LOG_DIR $CELERYD_PID_DIR
 
-# Append the Django settings module to use, if specified
-if [ -n "$DJANGO_SETTINGS_MODULE" ]; then
-    CELERYD_OPTS="$CELERYD_OPTS --settings=$DJANGO_SETTINGS_MODULE"
-fi
+        CELERYD_OPTS="$CELERYD_OPTS --uid=$CELERYD_USER"
+    fi
 
-start_worker () {
-    echo -n $"Starting celeryd: "
-    daemon --pidfile=$CELERYD_PID_FILE --user=$CELERYD_USER \
-      PYTHONPATH=$PYTHONPATH:$CELERY_PYTHONPATH $CELERYD $CELERYD_OPTS 2>/dev/null &
-    RETVAL=$?
-    sleep 3; echo
-    if [ -n "$VIRTUALENV" ]; then
-        source $VIRTUALENV/bin/activate
+    # If we need to be run from a specific location, cd to it before launch
+    if [ -n "$CELERYD_CHDIR" ]; then
+        cd $CELERYD_CHDIR
     fi
-    [ $RETVAL -eq 0 ] && touch /var/lock/subsys/celeryd
-}
 
-stop_worker () {
-    echo -n $"Stopping celeryd: "
-    killproc -p $CELERYD_PID_FILE $CELERYD 2>/dev/null
+    echo -n $"Starting $prog: "
+    $CELERYD_MULTI start $prog \
+                               --pidfile=$CELERYD_PID_FILE \
+                               --logfile=$CELERYD_LOG_FILE \
+                               --loglevel=$CELERYD_LOG_LEVEL \
+                               --cmd="$CELERYD" \
+			                   --quiet \
+                               $CELERYD_OPTS
     RETVAL=$?
-    sleep 3; echo
-    if [ $RETVAL -eq 0 ]; then
-        rm -f /var/lock/subsys/celeryd
-        rm -f $CELERYD_PID_FILE
+
+    if [ "$RETVAL" == "0" ]; then
+        touch /var/lock/subsys/$prog
+        success
+    else
+        failure
     fi
+    echo
 }
 
-case "$1" in
-  start)
-    start_worker
-    ;;
+stop_workers () {
+    echo -n $"Stopping $prog: "
 
-  stop)
-    stop_worker
-    ;;
+    # If we haven't ended, explicitly kill it!
+    if [ ! -f $CELERYD_PID_FILE ] || [ ! -e /proc/$(cat $CELERYD_PID_FILE) ]; then
+        failure
+	echo
+        return
+    fi
 
-  restart)
-    stop_worker
-    start_worker
-    ;;
+    # First, try to nicely shut it down.
+    $CELERYD_MULTI stop $prog --pidfile=$CELERYD_PID_FILE --quiet
+    RETVAL=$?
 
-  status)
-    status celeryd
-    ;;
+    # SLeep a few seconds. (this was part of the original script; we can't
+    # trust that it will end immediately, or that running the command will
+    # stop it.
+    sleep 3
 
-  *)
-    echo $"Usage: $0 {start|stop|restart|status}"
-    exit 1
+    # If we haven't ended, explicitly kill it!
+    if [ -f $CELERYD_PID_FILE ] && [ -e /proc/$(cat $CELERYD_PID_FILE) ]; then
+        $CELERYD_MULTI stop $prog -KILL --pidfile=$CELERYD_PID_FILE --quiet
+    fi
+
+    if [ "$RETVAL" == "0" ]; then
+        rm -f /var/lock/sybsys/$prog
+        success
+    else
+        failure
+    fi
+    echo
+}
+
+case "$1" in
+    start)
+        start_workers ;;
+    stop)
+        stop_workers ;;
+    status)
+        status -p $CELERYD_PID_FILE $prog ;;
+    restart)
+        stop_workers
+        start_workers ;;
+    *)
+        echo "Usage: /etc/init.d/$prog {start|stop|restart|status}"
+        exit 1
+    ;;
 esac
 
-exit $?
+exit 0

contrib/centos/celeryd.sysconfig

@@ -1,15 +1,21 @@
-CELERY_PYTHONPATH="/srv/platform/lib"
-CELERYD_NODES="w1"
+# Passed into celeryd multi
+#CELERYD="-m celery.bin.celeryd_detach"
 
-# 5 minute hard time limit for tasks + enable events at startup.
-CELERYD_OPTS="--time-limit=300 -E --beat -s /tmp/celerybeat-schedule"
+# Path to the celerd multi
+#CELERYD_MULTI="/usr/bin/celeryd-multi"
 
-# Full path to the Django project directory.
-CELERYD_CHDIR="/srv/platform/lib/atizo"
+# Sets the verbosity of the celeryd logging.
+#CELERYD_LOG_LEVEL="INFO"
 
-CELERYD="$CELERYD_CHDIR/manage.py celeryd"
-CELERYD_LOG_FILE=/var/log/celeryd/celeryd.log
-CELERYD_PID_FILE=/var/run/celeryd/celeryd.pid
+# Define the loader that celeryd should use for loading in configs.
+#CELERY_LOADER=""
 
-CELERYD_USER="celeryd"
-CELERYD_GROUP="celeryd"
+# User and group information for directories
+#CELERYD_USER="celery"
+#CELERYD_GROUP="celery"
+
+# Default arguments to be passed into celeryd.
+#CELERYD_OPTS=""
+
+# Change to this directory first before launching celeryd.
+#CELERYD_CHDIR=""