From 1425718d9e0c17df6259db0bb0bdec640c08ab0a Mon Sep 17 00:00:00 2001
From: Matthew Jones <mat@matburt.net>
Date: Fri, 2 Oct 2015 15:57:50 -0400
Subject: [PATCH] Improve the efficiency of mongo connections

- Remove the connect action in the fact __init__ module
- Monkeypatch mongoengine.Document's _get_db to force a new connection
  or raise an error if we don't have the proper settings.
---
 awx/fact/__init__.py    | 26 --------------------------
 awx/fact/models/fact.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 26 deletions(-)

diff --git a/awx/fact/__init__.py b/awx/fact/__init__.py
index db6bd7fb41..d9f8f41d12 100644
--- a/awx/fact/__init__.py
+++ b/awx/fact/__init__.py
@@ -1,31 +1,5 @@
 # Copyright (c) 2015 Ansible, Inc.
 # All Rights Reserved.
 
-from __future__ import absolute_import
-
 import logging
-from django.conf import settings
-
-from mongoengine import connect
-from mongoengine.connection import get_db, ConnectionError
-from .utils.dbtransform import register_key_transform
-
 logger = logging.getLogger('awx.fact')
-
-# Connect to Mongo
-try:
-    # Sanity check: If we have intentionally invalid settings, then we
-    # know we cannot connect.
-    if settings.MONGO_HOST == NotImplemented:
-        raise ConnectionError
-
-    # Attempt to connect to the MongoDB database.
-    connect(settings.MONGO_DB,
-            host=settings.MONGO_HOST,
-            port=int(settings.MONGO_PORT),
-            username=settings.MONGO_USERNAME,
-            password=settings.MONGO_PASSWORD,
-            tz_aware=settings.USE_TZ)
-    register_key_transform(get_db())
-except ConnectionError:
-    logger.info('Failed to establish connect to MongoDB')
diff --git a/awx/fact/models/fact.py b/awx/fact/models/fact.py
index e6efb9d782..0342ee94ba 100644
--- a/awx/fact/models/fact.py
+++ b/awx/fact/models/fact.py
@@ -1,12 +1,47 @@
 # Copyright (c) 2015 Ansible, Inc.
 # All Rights Reserved
 
+from mongoengine import connect
 from mongoengine.base import BaseField
 from mongoengine import Document, DateTimeField, ReferenceField, StringField, IntField
+from mongoengine.connection import get_db, ConnectionError
+from awx.fact.utils.dbtransform import register_key_transform, KeyTransform
 from awx.fact.utils.dbtransform import KeyTransform
 
+from mongoengine.connection import get_db, DEFAULT_CONNECTION_NAME
+
+from django.conf import settings
+
+import logging
+logger = logging.getLogger('awx.fact.models.fact')
+
+
 key_transform = KeyTransform([('.', '\uff0E'), ('$', '\uff04')])
 
+@classmethod
+def _get_db_monkeypatched(cls):
+    """ Override the default _get_db mechanism to start a connection to the database """
+    # Connect to Mongo
+    try:
+        # Sanity check: If we have intentionally invalid settings, then we
+        # know we cannot connect.
+        if settings.MONGO_HOST == NotImplemented:
+            raise ConnectionError
+
+        # Attempt to connect to the MongoDB database.
+        connect(settings.MONGO_DB,
+                host=settings.MONGO_HOST,
+                port=int(settings.MONGO_PORT),
+                username=settings.MONGO_USERNAME,
+                password=settings.MONGO_PASSWORD,
+                tz_aware=settings.USE_TZ)
+        register_key_transform(get_db())
+    except ConnectionError:
+        logger.info('Failed to establish connect to MongoDB')
+    return get_db(cls._meta.get("db_alias", "default"))
+
+Document._get_db = _get_db_monkeypatched
+
 class TransformField(BaseField):
     def to_python(self, value):
         return key_transform.transform_outgoing(value, None)