From a3eac73fd1bf199aff21b885413949733e92228b Mon Sep 17 00:00:00 2001
From: spengjie <spengjie@sina.com>
Date: Mon, 28 Oct 2019 14:02:08 +0800
Subject: [PATCH] Fix #5734 Celery does not consider authMechanism on mongodb
 backend URLs (#5795)

* Fix #5734 Celery does not consider authMechanism on mongodb backend URLs

* Add unit test: test_get_connection_with_authmechanism

* Add unit test: test_get_connection_with_authmechanism_no_username

* Fix errors in Python 2.7

Remove "," after "**" operator
---
 celery/backends/mongodb.py      |  4 ++++
 t/unit/backends/test_mongodb.py | 37 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/celery/backends/mongodb.py b/celery/backends/mongodb.py
index dd698007241..8d551bca802 100644
--- a/celery/backends/mongodb.py
+++ b/celery/backends/mongodb.py
@@ -157,6 +157,10 @@ def _get_connection(self):
             # don't change self.options
             conf = dict(self.options)
             conf['host'] = host
+            if self.user:
+                conf['username'] = self.user
+            if self.password:
+                conf['password'] = self.password
 
             self._connection = MongoClient(**conf)
 
diff --git a/t/unit/backends/test_mongodb.py b/t/unit/backends/test_mongodb.py
index 8f904d99771..e39d96b0f33 100644
--- a/t/unit/backends/test_mongodb.py
+++ b/t/unit/backends/test_mongodb.py
@@ -5,6 +5,7 @@
 
 import pytest
 from kombu.exceptions import EncodeError
+from pymongo.errors import ConfigurationError
 
 from case import ANY, MagicMock, Mock, mock, patch, sentinel, skip
 from celery import states, uuid
@@ -220,6 +221,42 @@ def test_get_connection_no_connection_mongodb_uri(self):
             )
             assert sentinel.connection == connection
 
+    def test_get_connection_with_authmechanism(self):
+        with patch('pymongo.MongoClient') as mock_Connection:
+            self.app.conf.mongodb_backend_settings = None
+            uri = ('mongodb://'
+                   'celeryuser:celerypassword@'
+                   'localhost:27017/'
+                   'celerydatabase?authMechanism=SCRAM-SHA-256')
+            mb = MongoBackend(app=self.app, url=uri)
+            mock_Connection.return_value = sentinel.connection
+            connection = mb._get_connection()
+            mock_Connection.assert_called_once_with(
+                host=['localhost:27017'],
+                username='celeryuser',
+                password='celerypassword',
+                authmechanism='SCRAM-SHA-256',
+                **mb._prepare_client_options()
+            )
+            assert sentinel.connection == connection
+
+    def test_get_connection_with_authmechanism_no_username(self):
+        with patch('pymongo.MongoClient') as mock_Connection:
+            self.app.conf.mongodb_backend_settings = None
+            uri = ('mongodb://'
+                   'localhost:27017/'
+                   'celerydatabase?authMechanism=SCRAM-SHA-256')
+            mb = MongoBackend(app=self.app, url=uri)
+            mock_Connection.side_effect = ConfigurationError(
+                'SCRAM-SHA-256 requires a username.')
+            with pytest.raises(ConfigurationError):
+                mb._get_connection()
+            mock_Connection.assert_called_once_with(
+                host=['localhost:27017'],
+                authmechanism='SCRAM-SHA-256',
+                **mb._prepare_client_options()
+            )
+
     @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.