fix bug in set_allocations
[sfa.git] / sfa / storage / model.py
index 780853c..5495800 100644 (file)
@@ -114,6 +114,10 @@ class RegRecord (Base,AlchemyObj):
         result += ">"
         return result
 
+    # shortcut - former implem. was record-based
+    def get (self, field, default):
+        return getattr(self,field,default)
+
     @validates ('gid')
     def validate_gid (self, key, gid):
         if gid is None:                     return
@@ -123,6 +127,8 @@ class RegRecord (Base,AlchemyObj):
     def validate_datetime (self, key, incoming):
         if isinstance (incoming, datetime):     return incoming
         elif isinstance (incoming, (int,float)):return datetime.fromtimestamp (incoming)
+        else: logger.info("Cannot validate datetime for key %s with input %s"%\
+                              (key,incoming))
 
     @validates ('date_created')
     def validate_date_created (self, key, incoming): return self.validate_datetime (key, incoming)
@@ -144,6 +150,20 @@ class RegRecord (Base,AlchemyObj):
         now=datetime.now()
         self.last_updated=now
 
+#################### cross-relations tables
+# authority x user (pis) association
+authority_pi_table = \
+    Table ( 'authority_pi', Base.metadata,
+            Column ('authority_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
+            Column ('pi_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
+            )
+# slice x user (researchers) association
+slice_researcher_table = \
+    Table ( 'slice_researcher', Base.metadata,
+            Column ('slice_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
+            Column ('researcher_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
+            )
+
 ##############################
 # all subclasses define a convenience constructor with a default value for type, 
 # and when applicable a way to define local fields in a kwd=value argument
@@ -152,6 +172,13 @@ class RegAuthority (RegRecord):
     __tablename__       = 'authorities'
     __mapper_args__     = { 'polymorphic_identity' : 'authority' }
     record_id           = Column (Integer, ForeignKey ("records.record_id"), primary_key=True)
+    #### extensions come here
+    reg_pis             = relationship \
+        ('RegUser',
+         secondary=authority_pi_table,
+         primaryjoin=RegRecord.record_id==authority_pi_table.c.authority_id,
+         secondaryjoin=RegRecord.record_id==authority_pi_table.c.pi_id,
+         backref='reg_authorities_as_pi')
     
     def __init__ (self, **kwds):
         # fill in type if not previously set
@@ -163,13 +190,15 @@ class RegAuthority (RegRecord):
     def __repr__ (self):
         return RegRecord.__repr__(self).replace("Record","Authority")
 
-####################
-# slice x user (researchers) association
-slice_researcher_table = \
-    Table ( 'slice_researcher', Base.metadata,
-            Column ('slice_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
-            Column ('researcher_id', Integer, ForeignKey ('records.record_id'), primary_key=True),
-            )
+    def update_pis (self, pi_hrns):
+        # don't ruin the import of that file in a client world
+        from sfa.storage.alchemy import dbsession
+        # strip that in case we have <researcher> words </researcher>
+        pi_hrns = [ x.strip() for x in pi_hrns ]
+        request = dbsession.query (RegUser).filter(RegUser.hrn.in_(pi_hrns))
+        logger.info ("RegAuthority.update_pis: %d incoming pis, %d matches found"%(len(pi_hrns),request.count()))
+        pis = dbsession.query (RegUser).filter(RegUser.hrn.in_(pi_hrns)).all()
+        self.reg_pis = pis
 
 ####################
 class RegSlice (RegRecord):
@@ -182,7 +211,7 @@ class RegSlice (RegRecord):
          secondary=slice_researcher_table,
          primaryjoin=RegRecord.record_id==slice_researcher_table.c.slice_id,
          secondaryjoin=RegRecord.record_id==slice_researcher_table.c.researcher_id,
-         backref="reg_slices_as_researcher")
+         backref='reg_slices_as_researcher')
 
     def __init__ (self, **kwds):
         if 'type' not in kwds: kwds['type']='slice'
@@ -191,6 +220,28 @@ class RegSlice (RegRecord):
     def __repr__ (self):
         return RegRecord.__repr__(self).replace("Record","Slice")
 
+    def update_researchers (self, researcher_hrns):
+        # don't ruin the import of that file in a client world
+        from sfa.storage.alchemy import dbsession
+        # strip that in case we have <researcher> words </researcher>
+        researcher_hrns = [ x.strip() for x in researcher_hrns ]
+        request = dbsession.query (RegUser).filter(RegUser.hrn.in_(researcher_hrns))
+        logger.info ("RegSlice.update_researchers: %d incoming researchers, %d matches found"%(len(researcher_hrns),request.count()))
+        researchers = dbsession.query (RegUser).filter(RegUser.hrn.in_(researcher_hrns)).all()
+        self.reg_researchers = researchers
+
+    # when dealing with credentials, we need to retrieve the PIs attached to a slice
+    def get_pis (self):
+        # don't ruin the import of that file in a client world
+        from sfa.storage.alchemy import dbsession
+        from sfa.util.xrn import get_authority
+        authority_hrn = get_authority(self.hrn)
+        auth_record = dbsession.query(RegAuthority).filter_by(hrn=authority_hrn).first()
+        return auth_record.reg_pis
+        
+    @validates ('expires')
+    def validate_expires (self, key, incoming): return self.validate_datetime (key, incoming)
+
 ####################
 class RegNode (RegRecord):
     __tablename__       = 'nodes'
@@ -260,6 +311,61 @@ class RegKey (Base):
         result += ">"
         return result
 
+class SliverAllocation(Base,AlchemyObj):
+    __tablename__       = 'sliver_allocation'
+    sliver_id           = Column(String, primary_key=True)
+    allocation_state    = Column(String)
+
+    def __init__(self, **kwds):
+        if 'sliver_id' in kwds:
+            self.sliver_id = kwds['sliver_id']
+        if 'allocation_state' in kwds:
+            self.allocation_state = kwds['allocation_state']
+
+    def __repr__(self):
+        result = "<sliver_allocation sliver_id=%s allocation_state=%s" % \
+                  (self.sliver_id, self.allocation_state)
+        return result
+
+    @validates('allocation_state')
+    def validate_allocation_state(self, key, state):
+        allocation_states = ['geni_unallocated', 'geni_allocated', 'geni_provisioned']
+        assert state in allocation_states
+        return state
+
+    @staticmethod    
+    def set_allocations(sliver_ids, state):
+        from sfa.storage.alchemy import dbsession
+        if not isinstance(sliver_ids, list):
+            sliver_ids = [sliver_ids]
+        sliver_state_updated = {}
+        constraint = SliverAllocation.sliver_id.in_(sliver_ids)
+        sliver_allocations = dbsession.query (SliverAllocation).filter(constraint)
+        sliver_ids_found = []
+        for sliver_allocation in sliver_allocations:
+            sliver_allocation.allocation_state = state
+            sliver_ids_found.append(sliver_allocation.sliver_id)
+
+        # Some states may not have been updated becuase no sliver allocation state record
+        # exists for the sliver. Insert new allocation records for these slivers and set
+        # it to geni_allocated.
+        sliver_ids_not_found = set(sliver_ids).difference(sliver_ids_found)
+        for sliver_id in sliver_ids_not_found:
+            record = SliverAllocation(sliver_id=sliver_id, allocation_state=state)
+            dbsession.add(record)
+        dbsession.commit()
+
+    @staticmethod
+    def delete_allocations(sliver_ids):
+        from sfa.storage.alchemy import dbsession
+        if not isinstance(sliver_ids, list):
+            sliver_ids = [sliver_ids]
+        constraint = SliverAllocation.sliver_id.in_(sliver_ids)
+        sliver_allocations = dbsession.query(SliverAllocation).filter(constraint)
+        for sliver_allocation in sliver_allocations:
+            dbsession.delete(sliver_allocation)
+        dbsession.commit()
+
 ##############################
 # although the db needs of course to be reachable for the following functions
 # the schema management functions are here and not in alchemy