From 4ff16548b643f375dac5e54442e849dd523bdc5b Mon Sep 17 00:00:00 2001
From: Chris Meyers <chris.meyers.fsu@gmail.com>
Date: Wed, 9 Mar 2016 15:53:07 -0500
Subject: [PATCH] drop mongo database when complete

---
 awx/main/migrations/_system_tracking.py       | 14 ++++++++++++++
 .../tests/functional/migrations/test_fact.py  | 19 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/awx/main/migrations/_system_tracking.py b/awx/main/migrations/_system_tracking.py
index 0d8ea2cb15..95fdb62c2e 100644
--- a/awx/main/migrations/_system_tracking.py
+++ b/awx/main/migrations/_system_tracking.py
@@ -1,10 +1,23 @@
 
 from awx.fact.models import FactVersion
+from mongoengine.connection import ConnectionError
+from django.conf import settings
+
+def drop_system_tracking_db():
+    try:
+        db = FactVersion._get_db()
+        db.connection.drop_database(settings.MONGO_DB)
+    except ConnectionError:
+        # TODO: Log this. Not a deal-breaker. Just let the user know they
+        # may need to manually drop/delete the database.
+        pass
 
 def migrate_facts(apps, schema_editor):
     Fact = apps.get_model('main', "Fact")
     Host = apps.get_model('main', "Host")
 
+    # TODO: Check to see if mongo connection works and mongo is on.
+
     migrated_count = 0
     not_migrated_count = 0
     for factver in FactVersion.objects.all():
@@ -18,4 +31,5 @@ def migrate_facts(apps, schema_editor):
             # This isn't a hard error. Just something the user would want to know.
             not_migrated_count += 1
 
+    drop_system_tracking_db()
     return (migrated_count, not_migrated_count)
diff --git a/awx/main/tests/functional/migrations/test_fact.py b/awx/main/tests/functional/migrations/test_fact.py
index c1222f03e6..5114bf6a62 100644
--- a/awx/main/tests/functional/migrations/test_fact.py
+++ b/awx/main/tests/functional/migrations/test_fact.py
@@ -8,6 +8,9 @@ from awx.main.models.fact import Fact
 
 from awx.main.migrations import _system_tracking as system_tracking
 
+from awx.fact.models.fact import Fact as FactMongo
+from awx.fact.models.fact import FactVersion, FactHost
+
 def micro_to_milli(micro):
     return micro - (((int)(micro / 1000)) * 1000)
 
@@ -58,3 +61,19 @@ def test_migrate_facts_hostname_does_not_exist(inventories, hosts, hosts_mongo,
         assert len(fact) == 1
         assert fact[0] is not None
    
+@pytest.mark.django_db
+@pytest.mark.mongo_db
+def test_drop_system_tracking_db(inventories, hosts, hosts_mongo, fact_scans):
+    inventory_objs = inventories(1)
+    hosts_mongo(1, inventory_objs)
+    facts_known = fact_scans(1, inventory_objs)
+
+    assert FactMongo.objects.all().count() > 0
+    assert FactVersion.objects.all().count() > 0
+    assert FactHost.objects.all().count() > 0
+
+    system_tracking.drop_system_tracking_db()
+  
+    assert FactMongo.objects.all().count() == 0
+    assert FactVersion.objects.all().count() == 0
+    assert FactHost.objects.all().count() == 0