Snapshots: improve generality of the importer

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

2 files changed, 51 insertions, 55 deletions

diff --git a/src/lib/Server/Plugins/Snapshots.py b/src/lib/Server/Plugins/Snapshots.py
index 6f33fff4d..40a153034 100644
--- a/src/lib/Server/Plugins/Snapshots.py
+++ b/src/lib/Server/Plugins/Snapshots.py
@@ -24,36 +24,36 @@ class Snapshots(Bcfg2.Server.Plugin.Statistics,
         t1 = time.time()
         entries = dict([('Package', dict()),
                         ('Service', dict()), ('Path', dict())])
-        extra = dict([('Package', list()), ('Service', list()),
-                      ('Path', list())])
-        pdisp = {'Package': 'version',
-                 'Service': 'status'}
+        extra = dict([('Package', dict()), ('Service', dict()),
+                      ('Path', dict())])
+        pdisp = {'Package': ['name', 'type', 'version'],
+                 'Service': ['name', 'type', 'status']}
         for entry in xdata.find('.//Bad'):
             if entry.tag not in pdisp:
                 print "Not Found", entry.tag, entry.get('name')
                 continue
             else:
-                data = [False, False, unicode(entry.get('type')),
-                        unicode(entry.get('%s_%s' % ('current', pdisp[entry.tag]))),
-                        unicode(entry.get('%s_%s' % ('current', pdisp[entry.tag])))]
+                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
         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:
-                    data = [False, False, unicode(entry.get('type')),
-                            unicode(entry.get('%s_%s' % ('current', pdisp[entry.tag]))),
-                            unicode(entry.get(pdisp[entry.tag]))]
+                    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
             else:
                 print entry.tag, entry.get('name')
         for entry in xdata.find('.//Extra'):
             if entry.tag in pdisp:
-                tname = pdisp[entry.tag]
-                data = dict([('name', unicode(entry.get('name'))),
-                              ('type', unicode(entry.get('type'))),
-                              (tname, unicode(entry.get(tname)))])
+                current = dict([(key, unicode(entry.get(key))) for key in pdisp[entry.tag]])
+                extra[entry.tag][entry.get('name')] = current
             else:
                 print "extra", entry.tag, entry.get('name')
         t2 = time.time()
diff --git a/src/lib/Server/Snapshots/model.py b/src/lib/Server/Snapshots/model.py
index d6592bf6e..f1432ffce 100644
--- a/src/lib/Server/Snapshots/model.py
+++ b/src/lib/Server/Snapshots/model.py
@@ -1,7 +1,7 @@
-from sqlalchemy import Table, Column, Integer, Unicode, MetaData, ForeignKey, Boolean, DateTime, create_engine, UnicodeText
+from sqlalchemy import Table, Column, Integer, Unicode, ForeignKey, Boolean, DateTime, create_engine, UnicodeText
 import datetime
 import sqlalchemy.exceptions
-from sqlalchemy.orm import relation, backref, sessionmaker
+from sqlalchemy.orm import relation, backref
 from sqlalchemy.ext.declarative import declarative_base
 
 class Uniquer(object):
@@ -12,6 +12,10 @@ class Uniquer(object):
         except sqlalchemy.exceptions.InvalidRequestError:
             return cls(**kwargs)
 
+    @classmethod
+    def from_record(cls, session, data):
+        return cls.by_value(session, **data)
+
 Base = declarative_base()
 
 class Administrator(Uniquer, Base):
@@ -89,7 +93,20 @@ class Package(Base, Uniquer):
     version = Column(Unicode(16))
     verification_status = Column(Boolean)
 
-class PackageCorrespondence(Base):
+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)
+        if s_dict != e_dict:
+            end = cls.mtype.by_value(mysession, **e_dict)
+        else:
+            end = start
+        return cls(start=start, end=end, modified=mod, correct=corr)
+
+class PackageCorrespondence(Base, CorrespondenceType):
+    mtype = Package
     __tablename__ = 'package_pair'
     id = Column(Integer, primary_key=True)    
     start_id = Column(Integer, ForeignKey('package.id'))
@@ -99,18 +116,6 @@ class PackageCorrespondence(Base):
     modified = Column(Boolean)
     correct = Column(Boolean)
 
-    @classmethod
-    def from_record(cls, mysession, name, record):
-        (mod, corr, ptype, s_vers, e_vers) = record
-        start = Package.by_value(mysession, name=unicode(name), type=ptype,
-                                 version=s_vers)
-        if s_vers != e_vers:
-            end = Package.by_value(mysession, name=unicode(name), type=ptype,
-                                     version=e_vers)
-        else:
-            end = start
-        return cls(start=start, end=end, modified=mod, correct=corr)
-
 package_snap = Table('package_snap', Base.metadata,
                      Column('ppair_id', Integer, ForeignKey('package_pair.id')),
                      Column('snapshot_id', Integer, ForeignKey('snapshot.id')))
@@ -122,7 +127,8 @@ class Service(Base, Uniquer):
     type = Column(Unicode(12))
     status = Column(Boolean)
 
-class ServiceCorrespondence(Base):
+class ServiceCorrespondence(Base, CorrespondenceType):
+    mtype = Service
     __tablename__ = 'service_pair'
     id = Column(Integer, primary_key=True)    
     start_id = Column(Integer, ForeignKey('service.id'))
@@ -132,18 +138,6 @@ class ServiceCorrespondence(Base):
     modified = Column(Boolean)
     correct = Column(Boolean)    
 
-    @classmethod
-    def from_record(cls, mysession, name, record):
-        (mod, corr, ptype, s_status, e_status) = record
-        start = Service.by_value(mysession, name=unicode(name), type=ptype,
-                                 status=s_status)
-        if s_status != e_status:
-            end = Service.by_value(mysession, name=unicode(name), type=ptype,
-                                   status=e_status)
-        else:
-            end = start
-        return cls(start=start, end=end, modified=mod, correct=corr)
-
 service_snap = Table('service_snap', Base.metadata,
                      Column('spair_id', Integer, ForeignKey('service_pair.id')),
                      Column('snapshot_id', Integer, ForeignKey('snapshot.id')))
@@ -158,7 +152,8 @@ class File(Base):
     perms = Column(Integer(5))
     contents = Column(UnicodeText)
 
-class FileCorrespondence(Base):
+class FileCorrespondence(Base, CorrespondenceType):
+    mtype = File
     __tablename__ = 'file_pair'
     id = Column(Integer, primary_key=True)    
     start_id = Column(Integer, ForeignKey('file.id'))
@@ -211,21 +206,22 @@ class Snapshot(Base):
     extra_services = relation(Service, secondary=extra_service_snap)
     extra_files = relation(File, secondary=extra_file_snap)
 
+    c_dispatch = dict([('Package', ('packages', PackageCorrespondence)),
+                       ('Service', ('services', ServiceCorrespondence)),
+                       ('Path', ('files', FileCorrespondence))])
+    e_dispatch = dict([('Package', ('extra_packages', Package)),
+                       ('Service', ('extra_services', Service)),
+                       ('Path', ('extra_files', File))])
+
     @classmethod
     def from_data(cls, session, metadata, entries, extra):
         dbm = Metadata.from_metadata(session, metadata)
         snap = cls(client_metadata=dbm, timestamp=datetime.datetime.now())
-        for pkg, pdata in entries['Package'].iteritems():
-            snap.packages.append(\
-                PackageCorrespondence.from_record(session, pkg, pdata))
-        for data in extra['Package']:
-            extra_pkg = Package.by_value(session, **data)
-            snap.extra_packages.append(extra_pkg)
-        for pkg, pdata in entries['Service'].iteritems():
-            snap.services.append(\
-                ServiceCorrespondence.from_record(session, pkg, pdata))
-        for data in extra['Service']:
-            extra_svc = Service.by_value(session, **data)
-            snap.extra_services.append(extra_svc)
+        for (dispatch, data) in [(cls.c_dispatch, entries),
+                                 (cls.e_dispatch, extra)]:
+            for key in dispatch:
+                dest, ecls = dispatch[key]
+                for edata in data[key].values():
+                    getattr(snap, dest).append(ecls.from_record(session, edata))
         return snap