some more mongo stuff (sorry that i pushed it here :()

toontown/building/DistributedBuildingMgrAI.py

@@ -1,4 +1,6 @@
 from direct.directnotify.DirectNotifyGlobal import *
+import cPickle
+from pymongo.errors import AutoReconnect
 
 from otp.ai.AIBaseGlobal import *
 from toontown.building import DistributedBuildingAI
@@ -19,6 +21,7 @@ class DistributedBuildingMgrAI:
         self.dnaStore = dnaStore
         self.trophyMgr = trophyMgr
         self.__buildings = {}
+        self.tableName = 'buildings_%s' % self.branchId
         self.findAllLandmarkBuildings()
 
     def cleanup(self):
@@ -105,10 +108,10 @@ class DistributedBuildingMgrAI:
         return (blocks, hqBlocks, gagshopBlocks, petshopBlocks, kartshopBlocks)
 
     def findAllLandmarkBuildings(self):
-        backups = simbase.backups.load('block-info', (self.air.districtId, self.branchId), default={})
+        buildings = self.load()
         (blocks, hqBlocks, gagshopBlocks, petshopBlocks, kartshopBlocks) = self.getDNABlockLists()
         for blockNumber in blocks:
-            self.newBuilding(blockNumber, backup=backups.get(blockNumber, None))
+            self.newBuilding(blockNumber, buildings.get(blockNumber, None))
         for blockNumber in hqBlocks:
             self.newHQBuilding(blockNumber)
         for blockNumber in gagshopBlocks:
@@ -118,25 +121,21 @@ class DistributedBuildingMgrAI:
         for block in kartshopBlocks:
             self.newKartShopBuilding(block)
 
-    def newBuilding(self, blockNumber, backup=None):
-        building = DistributedBuildingAI.DistributedBuildingAI(
-            self.air, blockNumber, self.branchId, self.trophyMgr)
+    def newBuilding(self, blockNumber, blockData = None):
+        building = DistributedBuildingAI.DistributedBuildingAI(self.air, blockNumber, self.branchId, self.trophyMgr)
         building.generateWithRequired(self.branchId)
-        if backup is not None:
-            state = backup.get('state', 'toon')
-            if ((state == 'suit') and simbase.air.wantCogbuildings) or (
-                (state == 'cogdo') and simbase.air.wantCogdominiums):
-                building.track = backup.get('track', 'c')
-                building.difficulty = backup.get('difficulty', 1)
-                building.numFloors = backup.get('numFloors', 1)
-                building.updateSavedBy(backup.get('savedBy'))
-                building.becameSuitTime = backup.get('becameSuitTime', time.time())
-                if (state == 'suit') and simbase.air.wantCogbuildings:
-                    building.setState('suit')
-                elif (state == 'cogdo') and simbase.air.wantCogdominiums:
+        if blockData:
+            building.track = blockData.get('track', 'c')
+            building.realTrack = blockData.get('track', 'c')
+            building.difficulty = int(blockData.get('difficulty', 1))
+            building.numFloors = int(blockData.get('numFloors', 1))
+            building.numFloors = max(1, min(5, building.numFloors))
+            building.becameSuitTime = blockData.get('becameSuitTime', time.time())
+            if blockData['state'] == 'suit':
+                building.setState('suit')
+            elif blockData['state'] == 'cogdo':
+                if simbase.air.wantCogdominiums:
                     building.setState('cogdo')
-                else:
-                    building.setState('toon')
             else:
                 building.setState('toon')
         else:
@@ -181,17 +180,59 @@ class DistributedBuildingMgrAI:
         return building
 
     def save(self):
+        if self.air.dbConn:
+            buildings = []
+            for i in self.__buildings.values():
+                if isinstance(i, HQBuildingAI.HQBuildingAI):
+                    continue
+                buildings.append(i.getPickleData())
+
+            street = {'ai': self.air.districtId, 'branch': self.branchId}
+            try:
+                self.air.dbGlobalCursor.streets.update(street,
+                                                        {'$setOnInsert': street,
+                                                        '$set': {'buildings': buildings}},
+                                                        upsert=True)
+            except AutoReconnect: # Something happened to our DB, but we can reconnect and retry.
+                taskMgr.doMethodLater(config.GetInt('mongodb-retry-time', 2), self.save, 'retrySave', extraArgs=[])
+        
+        else:
+            self.saveDev()
+            
+    def saveDev(self):
         backups = {}
         for blockNumber in self.getSuitBlocks():
             building = self.getBuilding(blockNumber)
-            backup = {
-                'state': building.fsm.getCurrentState().getName(),
-                'block': building.block,
-                'track': building.track,
-                'difficulty': building.difficulty,
-                'numFloors': building.numFloors,
-                'savedBy': building.savedBy,
-                'becameSuitTime': building.becameSuitTime
-            }
-            backups[blockNumber] = backup
+            backups[blockNumber] = building.getPickleData()
         simbase.backups.save('block-info', (self.air.districtId, self.branchId), backups)
+        
+    def load(self):
+        if self.air.dbConn:
+            blocks = {}
+
+            # Ensure that toontown.streets is indexed. Doing this at loading time
+            # is a fine way to make sure that we won't upset players with a
+            # lagspike while we wait for the backend to handle the index request.
+            self.air.dbGlobalCursor.streets.ensure_index([('ai', 1),
+                                                        ('branch', 1)])
+
+            street = {'ai': self.air.districtId, 'branch': self.branchId}
+            try:
+                doc = self.air.dbGlobalCursor.streets.find_one(street)
+            except AutoReconnect: # We're failing over - normally we'd wait to retry, but this is on AI startup so we might want to retry (or refactor the bldgMgr so we can sanely retry).
+                return blocks
+
+            if not doc:
+                return blocks
+
+            for building in doc.get('buildings', []):
+                blocks[int(building['block'])] = building
+
+            return blocks
+        
+        else:
+            blocks = simbase.backups.load('block-info', (self.air.districtId, self.branchId), default={})
+            return blocks
+            
+
+   

toontown/distributed/ToontownInternalRepository.py

@@ -14,10 +14,12 @@ class ToontownInternalRepository(AstronInternalRepository):
         AstronInternalRepository.__init__(
             self, baseChannel, serverId=serverId, dcFileNames=dcFileNames,
             dcSuffix=dcSuffix, connectMethod=connectMethod, threadedNet=threadedNet)
+        
+        self.wantMongo = config.GetBool('want-mongo', False)
     
     def handleConnected(self):
         self.__messenger = ToontownNetMessengerAI(self)
-        if config.GetBool('want-mongo', False):
+        if self.wantMongo:
             import pymongo
             self.dbConn = pymongo.MongoClient(config.GetString('mongodb-url', 'localhost'))
             self.dbGlobalCursor = self.dbConn.toontownstride

toontown/racing/LeaderboardMgrAI.py

@@ -6,13 +6,30 @@ class LeaderboardMgrAI:
 
     def __init__(self, air):
         self.air = air
-        self.database = simbase.backups.load('leaderboard', (self.air.districtId,), default=({}))
+        if self.air.dbConn:
+            self.air.dbGlobalCursor.leaderboards.ensure_index([('ai', 1)])
+            shard = {'ai': self.air.districtId}
+            doc = self.air.dbGlobalCursor.leaderboards.find_one(shard)
+            if not doc:
+                self.database = ({})
+            else:
+                self.database = doc.get('database', ({}))
+            
+        else:
+            self.database = simbase.backups.load('leaderboard', (self.air.districtId,), default=({}))
 
     def getDatabase(self):
         return self.database
 
     def saveDatabase(self):
-        simbase.backups.save('leaderboard', (self.air.districtId,), self.database)
+        if self.air.dbConn:
+            shard = {'ai': self.air.districtId}
+            self.air.dbGlobalCursor.leaderboards.update(shard,
+                                                        {'$setOnInsert': shard,
+                                                        '$set': {'database': self.database}},
+                                                        upsert = True)
+        else:
+            simbase.backups.save('leaderboard', (self.air.districtId,), self.database)
         messenger.send('goofyLeaderboardChange')
 
     def trimList(self, list):