完成对mongodb的支持

Config.ini

@@ -1,11 +1,17 @@
 [DB]
 ;Configure the database information
-;type: SSDB/REDIS if use redis, only modify the host port，the type should be SSDB
-type = SSDB
+;type: SSDB/REDIS/MONGODB if use redis, only modify the host port，the type should be SSDB
+
+type = MONGODB
 host = localhost
-port = 8888
+port = 27017
 name = proxy
 
+;type = SSDB
+;host = localhost
+;port = 8888
+;name = proxy
+
 [ProxyGetter]
 ;register the proxy getter function
 freeProxyFirst  = 1

DB/DbClient.py

@@ -68,6 +68,8 @@ class DbClient(object):
             __type = "SsdbClient"
         elif "REDIS" == self.config.db_type:
             __type = "RedisClient"
+        elif "MONGODB" == self.config.db_type:
+            __type = "MongodbClient"
         else:
             pass
         assert __type, 'type error, Not support DB type: {}'.format(self.config.db_type)

DB/MongodbClient.py

@@ -8,56 +8,55 @@ import json
 
 
 class MongodbClient(object):
+    def __init__(self, name, host, port):
+        self.name = name
+        self.client = MongoClient(host, port)
+        self.db = self.client.proxy
 
-	def __init__(self, name, host, port):
-		self.name = name
-		self.client = MongoClient(host, port)
-		self.db = self.client.proxy
+    def changeTable(self, name):
+        self.name = name
 
+    def get(self, proxy):
+        data = self.db[self.name].find_one({'proxy': proxy})
+        return data['num'] if data != None else None
 
-	def changeTable(self, name):
-		self.name = name
+    def put(self, proxy, num=1):
+        if self.db[self.name].find_one({'proxy': proxy}):
+            return None
+        else:
+            self.db[self.name].insert({'proxy': proxy, 'num': num})
 
+    def pop(self):
+        value = list(self.db[self.name].aggregate([{'$sample': {'size': 1}}]))[0]['proxy']
+        if value:
+            self.delete(value)
+        return value
 
-	def get(self):
-		proxy = self.getAll()
-		return random.choice(proxy) if proxy else None
+    def delete(self, value):
+        self.db[self.name].remove({'proxy': value})
 
+    def getAll(self):
+        return {p['proxy']: p['num'] for p in self.db[self.name].find()}
 
-	def put(self, value):
-		if self.db[self.name].find_one({'proxy': value}):
-			return None
-		else:
-			self.db[self.name].insert({'proxy': value})
+    def clean(self):
+        self.client.drop_database('proxy')
 
+    def delete_all(self):
+        self.db[self.name].remove()
 
-	def pop(self):
-		value = self.get()
-		if value:
-			self.delete(value)
-		return value
+    def update(self, key, value):
+        self.db[self.name].update({'proxy': key}, {'$inc': {'num': value}})
 
+    def exists(self, key):
+        return True if self.db[self.name].find({'proxy': key}) == None else False
 
-	def delete(self, value):
-		self.db[self.name].remove({'proxy': value})
-
-
-	def getAll(self):
-		return [p['proxy'] for p in self.db[self.name].find()]
-
-
-	def clean(self):
-		self.client.drop_database('proxy')
-
-
-	def delete_all(self):
-		self.db[self.name].remove()
+    def getNumber(self):
+        return self.db[self.name].count()
 
 
 if __name__ == "__main__":
-	db = MongodbClient('first', 'localhost', 27017)
-	db.put('127.0.0.1:1')
-	db2 = MongodbClient('second', 'localhost', 27017)
-	db2.put('127.0.0.1:2')
-	db.clean()
-
+    db = MongodbClient('first', 'localhost', 27017)
+    db.put('127.0.0.1:1')
+    db2 = MongodbClient('second', 'localhost', 27017)
+    db2.put('127.0.0.1:2')
+    db.clean()