Merge pull request #72 from netAir/master

完成对mongodb的支持

Config.ini

@@ -1,6 +1,6 @@
 [DB]
 ;Configure the database information
-;type: SSDB/REDIS if use redis, only modify the host port，the type should be SSDB
+;type: SSDB/REDIS/MONGODB if use redis, only modify the host port，the type should be SSDB
 type = SSDB
 host = localhost
 port = 8888

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

@@ -1,63 +1,74 @@
 # coding: utf-8
-
-__author__ = 'Maps'
+"""
+-------------------------------------------------
+   File Name：    MongodbClient.py
+   Description :  封装mongodb操作
+   Author :       JHao netAir
+   date：          2017/3/3
+-------------------------------------------------
+   Change Activity:
+                   2017/3/3:
+                   2017/9/26:完成对mongodb的支持
+-------------------------------------------------
+"""
+__author__ = 'Maps netAir'
 
 from pymongo import MongoClient
-import random
-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 changeTable(self, name):
+        self.name = name
 
-	def get(self):
-		proxy = self.getAll()
-		return random.choice(proxy) if proxy else None
+    def get(self, proxy):
+        data = self.db[self.name].find_one({'proxy': proxy})
+        return data['num'] if data != None else None
 
+    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 put(self, value):
-		if self.db[self.name].find_one({'proxy': value}):
-			return None
-		else:
-			self.db[self.name].insert({'proxy': value})
+    def pop(self):
+        data = list(self.db[self.name].aggregate([{'$sample': {'size': 1}}]))
+        if data:
+            data = data[0]
+            value = data['proxy']
+            self.delete(value)
+            return {'proxy': value, 'value': data['num']}
+        return None
 
+    def delete(self, value):
+        self.db[self.name].remove({'proxy': value})
 
-	def pop(self):
-		value = self.get()
-		if value:
-			self.delete(value)
-		return value
+    def getAll(self):
+        return {p['proxy']: p['num'] for p in self.db[self.name].find()}
 
+    def clean(self):
+        self.client.drop_database('proxy')
 
-	def delete(self, value):
-		self.db[self.name].remove({'proxy': value})
+    def delete_all(self):
+        self.db[self.name].remove()
 
+    def update(self, key, value):
+        self.db[self.name].update({'proxy': key}, {'$inc': {'num': value}})
 
-	def getAll(self):
-		return [p['proxy'] for p in self.db[self.name].find()]
+    def exists(self, key):
+        return True if self.db[self.name].find_one({'proxy': key}) != None else False
 
-
-	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')
+    print(db.pop())