1 files changed, 61 insertions, 16 deletions

diff --git a/src/lib/Bcfg2/Reporting/Storage/DjangoORM.py b/src/lib/Bcfg2/Reporting/Storage/DjangoORM.py
index efd9e594c..a8c8ce243 100644
--- a/src/lib/Bcfg2/Reporting/Storage/DjangoORM.py
+++ b/src/lib/Bcfg2/Reporting/Storage/DjangoORM.py
@@ -2,26 +2,54 @@
 The base for the original DjangoORM (DBStats)
 """
 
-from lxml import etree
-from datetime import datetime
+import difflib
 import traceback
+from datetime import datetime
 from time import strptime
-import Bcfg2.Options
-import Bcfg2.DBSettings
-from Bcfg2.Compat import md5
-from Bcfg2.Reporting.Storage.base import StorageBase, StorageError
-from Bcfg2.Server.Plugin.exceptions import PluginExecutionError
-from django.core import management
+from lxml import etree
+import sys
+
+import django
 from django.core.exceptions import ObjectDoesNotExist, MultipleObjectsReturned
 from django.db.models import FieldDoesNotExist
 from django.core.cache import cache
-from django import db
 
-#Used by GetCurrentEntry
-import difflib
-from Bcfg2.Compat import b64decode
-from Bcfg2.Reporting.models import *
+import Bcfg2.Options
+import Bcfg2.DBSettings
+from Bcfg2.Compat import b64decode, md5
 from Bcfg2.Reporting.Compat import transaction
+from Bcfg2.Reporting.Storage.base import StorageBase, StorageError
+from Bcfg2.Server.Plugin.exceptions import PluginExecutionError
+
+
+def load_django_models():
+    """ Load models for Django after option parsing has completed """
+    # pylint: disable=W0602
+    global Interaction, PackageEntry, FilePerms, PathEntry, LinkEntry, \
+        Group, Client, Bundle, TYPE_EXTRA, TYPE_BAD, TYPE_MODIFIED, \
+        FailureEntry, Performance, BaseEntry, ServiceEntry, ActionEntry, \
+        POSIXGroupEntry, POSIXUserEntry, SEBooleanEntry, SEFcontextEntry, \
+        SEInterfaceEntry, SELoginEntry, SEModuleEntry, SENodeEntry, \
+        SEPermissiveEntry, SEPortEntry, SEUserEntry
+    # pylint: enable=W0602
+
+    from Bcfg2.Reporting.models import \
+        Interaction, PackageEntry, FilePerms, PathEntry, LinkEntry, \
+        Group, Client, Bundle, TYPE_EXTRA, TYPE_BAD, TYPE_MODIFIED, \
+        FailureEntry, Performance, BaseEntry, ServiceEntry, ActionEntry, \
+        POSIXGroupEntry, POSIXUserEntry, SEBooleanEntry, SEFcontextEntry, \
+        SEInterfaceEntry, SELoginEntry, SEModuleEntry, SENodeEntry, \
+        SEPermissiveEntry, SEPortEntry, SEUserEntry
+
+
+def get_all_field_names(model):
+    if django.VERSION[0] == 1 and django.VERSION[1] >= 8:
+        return [field.name
+                for field in model._meta.get_fields()
+                if field.auto_created == False and
+                    not (field.is_relation and field.related_model is None)]
+    else:
+        return model._meta.get_all_field_names()
 
 
 class DjangoORM(StorageBase):
@@ -32,6 +60,7 @@ class DjangoORM(StorageBase):
             type=Bcfg2.Options.Types.size,
             help='Reporting file size limit',
             default=1024 * 1024)]
+    options_parsed_hook = staticmethod(load_django_models)
 
     def _import_default(self, entry, state, entrytype=None, defaults=None,
                         mapping=None, boolean=None, xforms=None):
@@ -80,7 +109,7 @@ class DjangoORM(StorageBase):
         for attr in boolean + ["current_exists"]:
             xforms[attr] = boolean_xform
         act_dict = dict(state=state)
-        for fieldname in entrytype._meta.get_all_field_names():
+        for fieldname in get_all_field_names(entrytype):
             if fieldname in ['id', 'hash_key', 'state']:
                 continue
             try:
@@ -284,6 +313,14 @@ class DjangoORM(StorageBase):
                 Group.objects.get_or_create(name=metadata['profile'])
         else:
             profile = None
+
+        flags = {'dry_run': False, 'only_important': False}
+        for flag in stats.findall('./Flags/Flag'):
+            value = flag.get('value', default='false').lower() == 'true'
+            name = flag.get('name')
+            if name in flags:
+                flags[name] = value
+
         inter = Interaction(client=client,
                              timestamp=timestamp,
                              state=stats.get('state', default="unknown"),
@@ -292,7 +329,8 @@ class DjangoORM(StorageBase):
                              good_count=stats.get('good', default="0"),
                              total_count=stats.get('total', default="0"),
                              server=server,
-                             profile=profile)
+                             profile=profile,
+                             **flags)
         inter.save()
         self.logger.debug("Interaction for %s at %s with INSERTED in to db" %
                 (client.id, timestamp))
@@ -363,6 +401,9 @@ class DjangoORM(StorageBase):
                             metric=metric,
                             value=value).save()
 
+        inter.ready = True
+        inter.save()
+
     def import_interaction(self, interaction):
         """Import the data into the backend"""
         try:
@@ -374,8 +415,12 @@ class DjangoORM(StorageBase):
         finally:
             self.logger.debug("%s: Closing database connection" %
                               self.__class__.__name__)
-            db.close_connection()
 
+            if django.VERSION[0] == 1 and django.VERSION[1] >= 7:
+                for connection in django.db.connections.all():
+                    connection.close()
+            else:
+                django.db.close_connection()
 
     def validate(self):
         """Validate backend storage.  Should be called once when loaded"""