Added ability to store probe data in database instead of probed.xml

1 files changed, 110 insertions, 61 deletions

diff --git a/src/lib/Bcfg2/Server/Plugins/Probes.py b/src/lib/Bcfg2/Server/Plugins/Probes.py
index 8ef6c8737..3932c44d1 100644
--- a/src/lib/Bcfg2/Server/Plugins/Probes.py
+++ b/src/lib/Bcfg2/Server/Plugins/Probes.py
@@ -1,9 +1,12 @@
-import time
-import lxml.etree
-import operator
 import re
 import os
+import sys
+import time
+import operator
+import lxml.etree
 import Bcfg2.Server
+from django.db import models
+import Bcfg2.Server.Plugin
 
 try:
     import json
@@ -16,20 +19,32 @@ except ImportError:
         has_json = False
 
 try:
-    import syck
-    has_syck = True
+    import syck as yaml
+    has_yaml = True
+    yaml_error = syck.error
 except ImportError:
-    has_syck = False
     try:
         import yaml
+        yaml_error = yaml.YAMLError
         has_yaml = True
     except ImportError:
         has_yaml = False
 
 import Bcfg2.Server.Plugin
 
-specific_probe_matcher = re.compile("(.*/)?(?P<basename>\S+)(.(?P<mode>[GH](\d\d)?)_\S+)")
-probe_matcher = re.compile("(.*/)?(?P<basename>\S+)")
+class ProbesDataModel(models.Model,
+                      Bcfg2.Server.Plugin.PluginDatabaseModel):
+    hostname = models.CharField(max_length=255)
+    probe = models.CharField(max_length=255)
+    timestamp = models.DateTimeField(auto_now=True)
+    data = models.TextField(null=True)
+
+
+class ProbesGroupsModel(models.Model,
+                        Bcfg2.Server.Plugin.PluginDatabaseModel):
+    hostname = models.CharField(max_length=255)
+    group = models.CharField(max_length=255)
+
 
 class ClientProbeDataSet(dict):
     """ dict of probe => [probe data] that records a for each host """
@@ -42,9 +57,9 @@ class ClientProbeDataSet(dict):
 
 
 class ProbeData(str):
-    """ a ProbeData object emulates a str object, but also has .xdata
-    and .json properties to provide convenient ways to use ProbeData
-    objects as XML or JSON data """
+    """ a ProbeData object emulates a str object, but also has .xdata,
+    .json, and .yaml properties to provide convenient ways to use
+    ProbeData objects as XML, JSON, or YAML data """
     def __new__(cls, data):
         return str.__new__(cls, data)
     
@@ -81,22 +96,18 @@ class ProbeData(str):
 
     @property
     def yaml(self):
-        if self._yaml is None:
-            if has_yaml:
-                try:
-                    self._yaml = yaml.load(self.data)
-                except yaml.YAMLError:
-                    pass
-            elif has_syck:
-                try:
-                    self._yaml = syck.load(self.data)
-                except syck.error:
-                    pass
+        if self._yaml is None and has_yaml:
+            try:
+                self._yaml = yaml.load(self.data)
+            except yaml_error:
+                pass
         return self._yaml
 
 
 class ProbeSet(Bcfg2.Server.Plugin.EntrySet):
     ignore = re.compile("^(\.#.*|.*~|\\..*\\.(tmp|sw[px])|probed\\.xml)$")
+    probename = re.compile("(.*/)?(?P<basename>\S+?)(\.(?P<mode>(?:G\d\d)|H)_\S+)?$")
+    bangline = re.compile('^#!\s*(?P<interpreter>.*)$')
 
     def __init__(self, path, fam, encoding, plugin_name):
         fpattern = '[0-9A-Za-z_\-]+'
@@ -105,20 +116,10 @@ class ProbeSet(Bcfg2.Server.Plugin.EntrySet):
                                               Bcfg2.Server.Plugin.SpecificData,
                                               encoding)
         fam.AddMonitor(path, self)
-        self.bangline = re.compile('^#!(?P<interpreter>.*)$')
 
     def HandleEvent(self, event):
-        if event.filename != self.path:
-            if (event.code2str == 'changed' and
-                event.filename.endswith("probed.xml") and
-                event.filename not in self.entries):
-                # for some reason, probed.xml is particularly prone to
-                # getting changed events before created events,
-                # because gamin is the worst ever.  anyhow, we
-                # specifically handle it here to avoid a warning on
-                # every single server startup.
-                self.entry_init(event)
-                return
+        if (event.filename != self.path and
+            not event.filename.endswith("probed.xml")):
             return self.handle_event(event)
 
     def get_probe_data(self, metadata):
@@ -127,9 +128,7 @@ class ProbeSet(Bcfg2.Server.Plugin.EntrySet):
         candidates = self.get_matching(metadata)
         candidates.sort(key=operator.attrgetter('specific'))
         for entry in candidates:
-            rem = specific_probe_matcher.match(entry.name)
-            if not rem:
-                rem = probe_matcher.match(entry.name)
+            rem = self.probename.match(entry.name)
             pname = rem.group('basename')
             if pname not in build:
                 build[pname] = entry
@@ -150,7 +149,8 @@ class ProbeSet(Bcfg2.Server.Plugin.EntrySet):
 
 class Probes(Bcfg2.Server.Plugin.Plugin,
              Bcfg2.Server.Plugin.Probing,
-             Bcfg2.Server.Plugin.Connector):
+             Bcfg2.Server.Plugin.Connector,
+             Bcfg2.Server.Plugin.DatabaseBacked):
     """A plugin to gather information from a client machine."""
     name = 'Probes'
     __author__ = 'bcfg-dev@mcs.anl.gov'
@@ -159,19 +159,32 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
         Bcfg2.Server.Plugin.Plugin.__init__(self, core, datastore)
         Bcfg2.Server.Plugin.Connector.__init__(self)
         Bcfg2.Server.Plugin.Probing.__init__(self)
+        Bcfg2.Server.Plugin.DatabaseBacked.__init__(self)
 
         try:
             self.probes = ProbeSet(self.data, core.fam, core.encoding,
                                    self.name)
         except:
-            raise Bcfg2.Server.Plugin.PluginInitError
+            err = sys.exc_info()[1]
+            raise Bcfg2.Server.Plugin.PluginInitError(err)
 
         self.probedata = dict()
         self.cgroups = dict()
         self.load_data()
 
-    def write_data(self):
+    @property
+    def _use_db(self):
+        return self.core.setup.cfp.getboolean("probes", "use_database",
+                                              default=False)
+
+    def write_data(self, client):
         """Write probe data out for use with bcfg2-info."""
+        if self._use_db:
+            return self._write_data_db(client)
+        else:
+            return self._write_data_xml(client)
+
+    def _write_data_xml(self, _):
         top = lxml.etree.Element("Probed")
         for client, probed in sorted(self.probedata.items()):
             cx = lxml.etree.SubElement(top, 'Client', name=client,
@@ -185,17 +198,45 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
                                    xml_declaration=True,
                                    pretty_print='true')
         try:
-            datafile = open("%s/%s" % (self.data, 'probed.xml'), 'w')
+            datafile = open(os.path.join(self.data, 'probed.xml'), 'w')
             datafile.write(data.decode('utf-8'))
         except IOError:
-            self.logger.error("Failed to write probed.xml")
+            err = sys.exc_info()[1]
+            self.logger.error("Failed to write probed.xml: %s" % err)
+
+    def _write_data_db(self, client):
+        for probe, data in self.probedata[client.hostname].items():
+            pdata = \
+                ProbesDataModel.objects.get_or_create(hostname=client.hostname,
+                                                      probe=probe)[0]
+            if pdata.data != data:
+                pdata.data = data
+                pdata.save()
+        ProbesDataModel.objects.filter(hostname=client.hostname).exclude(probe__in=self.probedata[client.hostname]).delete()
+
+        for group in self.cgroups[client.hostname]:
+            try:
+                ProbesGroupsModel.objects.get(hostname=client.hostname,
+                                              group=group)
+            except ProbesGroupsModel.DoesNotExist:
+                grp = ProbesGroupsModel(hostname=client.hostname,
+                                        group=group)
+                grp.save()
+        ProbesGroupsModel.objects.filter(hostname=client.hostname).exclude(group__in=self.cgroups[client.hostname]).delete()
 
     def load_data(self):
+        if self._use_db:
+            return self._load_data_db()
+        else:
+            return self._load_data_xml()
+            
+    def _load_data_xml(self):
         try:
             data = lxml.etree.parse(os.path.join(self.data, 'probed.xml'),
                                     parser=Bcfg2.Server.XMLParser).getroot()
         except:
-            self.logger.error("Failed to read file probed.xml")
+            err = sys.exc_info()[1]
+            self.logger.error("Failed to read file probed.xml: %s" % err)
             return
         self.probedata = {}
         self.cgroups = {}
@@ -204,12 +245,25 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
                 ClientProbeDataSet(timestamp=client.get("timestamp"))
             self.cgroups[client.get('name')] = []
             for pdata in client:
-                if (pdata.tag == 'Probe'):
+                if pdata.tag == 'Probe':
                     self.probedata[client.get('name')][pdata.get('name')] = \
-                        ProbeData(pdata.get('value'))
-                elif (pdata.tag == 'Group'):
+                        ProbeData(pdata.get("value"))
+                elif pdata.tag == 'Group':
                     self.cgroups[client.get('name')].append(pdata.get('name'))
 
+    def _load_data_db(self):
+        self.probedata = {}
+        self.cgroups = {}
+        for pdata in ProbesDataModel.objects.all():
+            if pdata.hostname not in self.probedata:
+                self.probedata[pdata.hostname] = \
+                    ClientProbeDataSet(timestamp=time.mktime(pdata.timestamp.timetuple()))
+            self.probedata[pdata.hostname][pdata.probe] = ProbeData(pdata.data)
+        for pgroup in ProbesGroupsModel.objects.all():
+            if pgroup.hostname not in self.cgroups:
+                self.cgroups[pgroup.hostname] = []
+            self.cgroups[pgroup.hostname].append(pgroup.group)
+
     def GetProbes(self, meta, force=False):
         """Return a set of probes for execution on client."""
         return self.probes.get_probe_data(meta)
@@ -219,25 +273,24 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
         self.probedata[client.hostname] = ClientProbeDataSet()
         for data in datalist:
             self.ReceiveDataItem(client, data)
-        self.write_data()
+        self.write_data(client)
 
     def ReceiveDataItem(self, client, data):
         """Receive probe results pertaining to client."""
         if client.hostname not in self.cgroups:
             self.cgroups[client.hostname] = []
+        if client.hostname not in self.probedata:
+            self.probedata[client.hostname] = ClientProbeDataSet()
         if data.text == None:
-            self.logger.error("Got null response to probe %s from %s" % \
-                              (data.get('name'), client.hostname))
-            try:
-                self.probedata[client.hostname].update({data.get('name'):
+            self.logger.info("Got null response to probe %s from %s" %
+                             (data.get('name'), client.hostname))
+            self.probedata[client.hostname].update({data.get('name'):
                                                         ProbeData('')})
-            except KeyError:
-                self.probedata[client.hostname] = \
-                    ClientProbeDataSet([(data.get('name'), ProbeData(''))])
             return
         dlines = data.text.split('\n')
-        self.logger.debug("%s:probe:%s:%s" % (client.hostname,
-            data.get('name'), [line.strip() for line in dlines]))
+        self.logger.debug("%s:probe:%s:%s" %
+                          (client.hostname, data.get('name'),
+                           [line.strip() for line in dlines]))
         for line in dlines[:]:
             if line.split(':')[0] == 'group':
                 newgroup = line.split(':')[1].strip()
@@ -245,11 +298,7 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
                     self.cgroups[client.hostname].append(newgroup)
                 dlines.remove(line)
         dobj = ProbeData("\n".join(dlines))
-        try:
-            self.probedata[client.hostname].update({data.get('name'): dobj})
-        except KeyError:
-            self.probedata[client.hostname] = \
-                ClientProbeDataSet([(data.get('name'), dobj)])
+        self.probedata[client.hostname].update({data.get('name'): dobj})
 
     def get_additional_groups(self, meta):
         return self.cgroups.get(meta.hostname, list())