Snapshots: Finish off importer (for old-style statistics)

git-svn-id: https://svn.mcs.anl.gov/repos/bcfg/trunk/bcfg2@5108 ce84e21b-d406-0410-9b95-82705330c041

2 files changed, 73 insertions, 33 deletions

diff --git a/src/lib/Server/Plugins/Snapshots.py b/src/lib/Server/Plugins/Snapshots.py
index ea2cb9095..d4efc2be3 100644
--- a/src/lib/Server/Plugins/Snapshots.py
+++ b/src/lib/Server/Plugins/Snapshots.py
@@ -1,4 +1,6 @@
 import lxml.etree
+import binascii
+import difflib
 import sqlalchemy
 import sqlalchemy.orm
 import Bcfg2.Server.Plugin
@@ -7,6 +9,40 @@ from Bcfg2.Server.Snapshots.model import Snapshot
 import time
 
 ftypes = ['ConfigFile', 'SymLink', 'Directory']
+datafields = {'Package': ['version'],
+              'Service': ['status'],
+              'ConfigFile': ['owner', 'group', 'perms']}
+
+def build_snap_ent(entry):
+    basefields = []
+    if entry.tag in ['Package', 'Service']:
+        basefields += ['type']
+    desired = dict([(key, unicode(entry.get(key))) for key in basefields])
+    state = dict([(key, unicode(entry.get(key))) for key in basefields])
+    desired.update([(key, unicode(entry.get(key))) for key in \
+                 datafields[entry.tag]])
+    if entry.tag == 'ConfigFile':
+        if entry.text == None:
+            desired['contents'] = None
+        else:
+            if entry.get('encoding', 'ascii') == 'ascii':
+                desired['contents'] = unicode(entry.text)
+            else:
+                desired['contents'] = unicode(binascii.a2b_base64(entry.text))
+
+        if 'current_bfile' in entry.attrib:
+            state['contents'] = unicode(binascii.a2b_base64( \
+                entry.get('current_bfile')))
+        elif 'current_bdiff' in entry.attrib:
+            diff = binascii.a2b_base64(entry.get('current_bdiff'))
+            state['contents'] = unicode( \
+                '\n'.join(difflib.restore(diff.split('\n'), 1)))
+                             
+    state.update([(key, unicode(entry.get('current_' + key, entry.get(key)))) \
+                  for key in datafields[entry.tag]])
+    if entry.tag == 'ConfigFile' and entry.get('exists', 'true') == 'false':
+        state = None
+    return [desired, state]
 
 class Snapshots(Bcfg2.Server.Plugin.Statistics,
                 Bcfg2.Server.Plugin.Plugin):
@@ -28,40 +64,37 @@ class Snapshots(Bcfg2.Server.Plugin.Statistics,
                         ('Service', dict()), ('Path', dict())])
         extra = dict([('Package', dict()), ('Service', dict()),
                       ('Path', dict())])
-        pdisp = {'Package': ['name', 'type', 'version'],
-                 'Service': ['name', 'type', 'status']}
-
+        bad = []
         for entry in xdata.find('.//Bad'):
-            if entry.tag not in pdisp:
-                print "Not Found", entry.tag, entry.get('name')
-                continue
+            data = [False, False, unicode(entry.get('name'))] \
+                   + build_snap_ent(entry)
+            if entry.tag in ftypes:
+                etag = 'Path'
             else:
-                edata = dict([(key, unicode(entry.get('current_%s' % key))) \
-                              for key in pdisp[entry.tag]])
-                data = [False, False, edata, edata]
-                entries[entry.tag][entry.get('name')] = data
+                etag = entry.tag
+            entries[etag][entry.get('name')] = data
         for entry in xdata.find('.//Modified'):
-            if entry.tag in pdisp:
-                if entry.get('name') in entries[entry.tag]:
-                    entries[entry.tag][entry.get('name')][0] = True
-                else:
-                    current = dict([(key, unicode(entry.get('current_%s' % key))) \
-                                    for key in pdisp[entry.tag]])
-                    desired = dict([(key, unicode(entry.get(key))) \
-                                    for key in pdisp[entry.tag]])
-                    data = [False, False, current, desired]
-                    entries[entry.tag][entry.get('name')] = data
+            if entry.tag in ftypes:
+                etag = 'Path'
+            else:
+                etag = entry.tag
+            if entry.get('name') in entries[etag]:
+                data = [True, False, unicode(entry.get('name'))] + \
+                       build_snap_ent(entry)
             else:
-                print entry.tag, entry.get('name')
+                data = [True, False, unicode(entry.get('name'))] + \
+                       build_snap_ent(entry)
         for entry in xdata.find('.//Extra'):
-            if entry.tag in pdisp:
-                current = dict([(key, unicode(entry.get(key))) for key in pdisp[entry.tag]])
-                extra[entry.tag][entry.get('name')] = current
+            if entry.tag in datafields:
+                current = dict([(key, unicode(entry.get(key))) \
+                                for key in datafields[entry.tag]])
+                extra[entry.tag][unicode(entry.get('name'))] = current
             else:
                 print "extra", entry.tag, entry.get('name')
         t2 = time.time()
         snap = Snapshot.from_data(self.session, metadata, entries, extra)
-        self.session.save(snap)
+        self.session.add(snap)
         self.session.commit()
         t3 = time.time()
+        print "Storage took", t3-t2
         return True
diff --git a/src/lib/Server/Snapshots/model.py b/src/lib/Server/Snapshots/model.py
index f1432ffce..ea9c6b811 100644
--- a/src/lib/Server/Snapshots/model.py
+++ b/src/lib/Server/Snapshots/model.py
@@ -1,15 +1,19 @@
-from sqlalchemy import Table, Column, Integer, Unicode, ForeignKey, Boolean, DateTime, create_engine, UnicodeText
+from sqlalchemy import Table, Column, Integer, Unicode, ForeignKey, Boolean, DateTime, UnicodeText
 import datetime
 import sqlalchemy.exceptions
 from sqlalchemy.orm import relation, backref
 from sqlalchemy.ext.declarative import declarative_base
 
 class Uniquer(object):
+    force_rt = True
     @classmethod
     def by_value(cls, session, **kwargs):
-        try:
-            return session.query(cls).filter_by(**kwargs).one()
-        except sqlalchemy.exceptions.InvalidRequestError:
+        if cls.force_rt:
+            try:
+                return session.query(cls).filter_by(**kwargs).one()
+            except sqlalchemy.exceptions.InvalidRequestError:
+                return cls(**kwargs)
+        else:
             return cls(**kwargs)
 
     @classmethod
@@ -97,10 +101,13 @@ class CorrespondenceType(object):
     mtype = Package
     @classmethod
     def from_record(cls, mysession, record):
-        (mod, corr, s_dict, e_dict) = record
-        start = cls.mtype.by_value(mysession, **s_dict)
+        (mod, corr, name, s_dict, e_dict) = record
+        if not s_dict:
+            start=None
+        else:
+            start = cls.mtype.by_value(mysession, name=name, **s_dict)
         if s_dict != e_dict:
-            end = cls.mtype.by_value(mysession, **e_dict)
+            end = cls.mtype.by_value(mysession, name=name, **e_dict)
         else:
             end = start
         return cls(start=start, end=end, modified=mod, correct=corr)
@@ -142,7 +149,7 @@ service_snap = Table('service_snap', Base.metadata,
                      Column('spair_id', Integer, ForeignKey('service_pair.id')),
                      Column('snapshot_id', Integer, ForeignKey('snapshot.id')))
 
-class File(Base):
+class File(Base, Uniquer):
     __tablename__ = 'file'
     id = Column(Integer, primary_key=True)        
     name = Column(UnicodeText)