liumiaomiao hace 9 meses
padre
commit
6ec643f8b6
Se han modificado 2 ficheros con 82 adiciones y 1 borrados
  1. 81 0
      lib/mogodb_helper.py
  2. 1 1
      tools/从mongo导出数据/mongo.py

+ 81 - 0
lib/mogodb_helper.py

@@ -0,0 +1,81 @@
+# coding:utf-8
+from pymongo import MongoClient
+from pymongo.errors import AutoReconnect
+from bson import ObjectId
+import time
+import urllib.parse as parse
+
+
+class MongoDBInterface:
+    def __init__(self, config):
+        self.__host = config.get("ip_port", "")
+        self.__user = config.get("user", "")
+        self.__password = config.get("password", "")
+        self.__database = config.get("db", "")
+        self.connect()
+
+    def connect(self):
+        """
+        连接数据库
+        :return:
+        """
+        # 特殊符号转义
+        self.__user = parse.quote_plus(self.__user)
+        self.__password = parse.quote_plus(self.__password)
+
+        # 连接数据库
+        try:
+            if self.__user:
+                self.client = MongoClient(
+                    "mongodb://{}:{}@{}".format(self.__user, self.__password, self.__host),
+                    unicode_decode_error_handler='ignore', directConnection=True)
+            else:
+                self.client = MongoClient("mongodb://{}".format(self.__host), unicode_decode_error_handler='ignore',
+                                          directConnection=True)
+            self.db = self.client[self.__database]
+        except AutoReconnect as e:
+            self.connect()
+
+    def disconnect(self):
+        self.client.close()
+
+    def find_by_id(self, collection_name, m_id):
+        try:
+            collection = self.db[collection_name]
+            if isinstance(m_id, str):
+                m_id = ObjectId(m_id)
+            documents = collection.find_one({"_id": m_id})
+            return documents
+        except AutoReconnect:
+            time.sleep(5)
+            print("Connection to MongoDB lost. Reconnecting...")
+            self.connect()  # 自动重连
+            return self.find_by_id(collection_name, m_id)
+
+    def update_by_id(self, collection_name, m_id, fields: dict):
+        try:
+            collection = self.db[collection_name]
+            if isinstance(m_id, str):
+                m_id = ObjectId(m_id)
+            documents = collection.update_one({"_id": m_id}, {"$set": fields})
+            return documents
+        except AutoReconnect:
+            print("Connection to MongoDB lost. Reconnecting...")
+            self.connect()  # 自动重连
+            return self.update_by_id(collection_name, m_id, fields)
+
+    def find_relus(self, collection_name, hash_id):
+        collection = self.db[collection_name]
+        row = collection.find_one({"rules_id": hash_id})
+        return row
+
+    def find_rule_by_company(self, collection_name, company, version):
+        collection = self.db[collection_name]
+        row = collection.find_one({"company_name": company, "version": version})
+        if not row:
+            return None
+        return row.get("rules_id")
+
+    def insert2db(self, collection_name, data: dict):
+        collection = self.db[collection_name]
+        collection.insert_one(data)

+ 1 - 1
tools/从mongo导出数据/mongo.py

@@ -1,6 +1,6 @@
 from pymongo import MongoClient
 from bson import ObjectId
-from util.mogodb_helper import MongoDBInterface
+from lib.mogodb_helper import MongoDBInterface
 
 #复制mongo数据源 从一个源到另一个源