元博网 - 局部MongoDB客户端会话变更为全局会话

ybw/utils/databases.py

@@ -6,22 +6,37 @@ from elasticsearch import Elasticsearch
 
 from config.load import mongo_conf, redis_conf, es_conf, analyze_url
 
-
 # ---------------------------------- mongo ----------------------------------
-def mongo_client(cfg=None):
-    if cfg is None:
-        cfg = mongo_conf
-    return pymongo.MongoClient(host=cfg['host'], port=cfg['port'])
+MONGO_URI_CLIENTS = {}    # a dictionary hold all client with uri as key
+
+
+def mongo_client(cfg=None, host=None, port=None, fork=False, **kwargs):
+    if host is not None and port is not None:
+        uri = f'mongodb://{host}:{port}'
+    else:
+        _cfg = (cfg or mongo_conf)
+        uri = f'mongodb://{_cfg["host"]}:{_cfg["port"]}'
+
+    if fork:
+        return pymongo.MongoClient(uri, **kwargs)
+    global MONGO_URI_CLIENTS
+    matched_client = MONGO_URI_CLIENTS.get(uri)
+    if matched_client is None:
+        new_client = pymongo.MongoClient(uri, **kwargs)
+        if new_client is not None:
+            MONGO_URI_CLIENTS[uri] = new_client
+        return new_client
+    return matched_client
 
 
-def mongo_database(db: str):
-    client = mongo_client()
-    return client[db]
+def mongo_database(name: str, **kw):
+    client = mongo_client(**kw)
+    return client.get_database(name)
 
 
-def mongo_table(db: str, coll: str):
-    client = mongo_client()
-    return client[db][coll]
+def mongo_table(db: str, name: str, **kw):
+    database = mongo_database(db, **kw)
+    return database.get_collection(name)
 
 
 def int2long(param: int):