File size: 6,412 Bytes
7ccd501 615412d 7ccd501 615412d 7ccd501 615412d 7ccd501 615412d 7ccd501 615412d 7ccd501 615412d 7ccd501 615412d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | import logging
from typing import Dict, List, Union
from bson import ObjectId
import ast
import json
import re
from pymongo import MongoClient
from threading import Lock # Needed for thread-safe pooling
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
level=logging.INFO
)
logger = logging.getLogger("mongo_service")
# ==============================================================================
# MONGO CONNECTION POOL MANAGER
# ==============================================================================
class MongoConnectionManager:
def __init__(self):
self._clients: Dict[str, MongoClient] = {}
self._lock = Lock()
def get_client(self, uri: str) -> MongoClient:
"""
Returns a cached MongoClient instance.
If it doesn't exist, creates one with a connection pool.
"""
# Double-checked locking for performance
if uri not in self._clients:
with self._lock:
if uri not in self._clients:
logger.info(f"Initialize new MongoDB Client for URI: {uri[:20]}...")
# maxPoolSize=50 allows 50 concurrent ops per URI.
# The rest will wait in queue automatically.
self._clients[uri] = MongoClient(
uri,
serverSelectionTimeoutMS=5000,
connectTimeoutMS=5000,
maxPoolSize=50,
minPoolSize=1
)
return self._clients[uri]
# Global Manager Instance
mongo_manager = MongoConnectionManager()
# ==============================================================================
# HELPER FUNCTIONS (Unchanged)
# ==============================================================================
def convert_oid(obj):
"""Recursively convert $oid format or ObjectId objects to strings."""
if isinstance(obj, ObjectId):
return str(obj)
if isinstance(obj, dict):
if set(obj.keys()) == {"$oid"}:
return str(obj["$oid"])
return {k: convert_oid(v) for k, v in obj.items()}
elif isinstance(obj, list):
return [convert_oid(item) for item in obj]
else:
return obj
def parse_query_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
"""Ensures the input is a valid MongoDB query object (Dict or List)."""
if isinstance(query_input, (dict, list)):
return convert_oid(query_input)
query_str = str(query_input).strip()
try:
return json.loads(query_str)
except json.JSONDecodeError:
pass
match = re.search(r"```(?:python|json|javascript)?\s*(.*?)\s*```", query_str, re.DOTALL)
if match:
query_str = match.group(1).strip()
query_str = re.sub(r'^db\.\w+\.\w+\(', '', query_str)
if query_str.endswith(')'):
query_str = query_str[:-1]
try:
parsed = ast.literal_eval(query_str)
return convert_oid(parsed)
except (ValueError, SyntaxError) as e:
logger.error(f"Failed to parse query string: {e}")
raise ValueError(f"Could not parse query string: {str(e)}")
def get_value_ignore_case(d: Dict, keys: List[str], default=None):
for k in keys:
if k in d:
return d[k]
return default
# ==============================================================================
# MAIN EXECUTION LOGIC (Modified)
# ==============================================================================
def execute_mongo_operation(
mongo_uri: str,
db_name: str,
collection_name: str,
query: Union[Dict, List],
limited: bool = False,
limit_rows: int = 20
):
"""
Executes MongoDB operations using the shared connection pool.
"""
# 1. GET CLIENT FROM MANAGER (Do NOT create new MongoClient here)
client = mongo_manager.get_client(mongo_uri)
try:
db = client[db_name]
collection = db[collection_name]
results = []
# --- Aggregation ---
if isinstance(query, list):
# Apply Limit if requested and not already present at the end
if limited:
if not (query and "$limit" in query[-1]):
query.append({"$limit": limit_rows})
cursor = collection.aggregate(query, allowDiskUse=True)
results = list(cursor)
# --- Find / Command ---
elif isinstance(query, dict):
command_keys = {'filter', 'query', '$query', 'projection', 'sort', 'limit', 'skip'}
has_command_keys = bool(set(query.keys()) & command_keys)
cursor = None
if has_command_keys and ('filter' in query or 'query' in query or '$query' in query):
query_filter = get_value_ignore_case(query, ['filter', 'query', '$query'], {})
projection = get_value_ignore_case(query, ['projection', 'fields'], None)
if projection == {}: projection = None
internal_limit = int(get_value_ignore_case(query, ['limit'], 0))
skip = int(get_value_ignore_case(query, ['skip'], 0))
sort_val = get_value_ignore_case(query, ['sort', '$orderby'], None)
cursor = collection.find(query_filter, projection)
if sort_val:
if isinstance(sort_val, dict):
sort_val = list(sort_val.items())
cursor = cursor.sort(sort_val)
if skip > 0: cursor = cursor.skip(skip)
if limited:
cursor = cursor.limit(limit_rows)
elif internal_limit > 0:
cursor = cursor.limit(internal_limit)
else:
cursor = collection.find(query)
if limited:
cursor = cursor.limit(limit_rows)
results = list(cursor)
else:
raise ValueError("Query must be a Dictionary (find) or List (aggregate)")
return results
except Exception as e:
logger.error(f"DB Execution Error: {e}")
raise e
# IMPORTANT: DO NOT CLOSE THE CLIENT
# finally:
# if client: client.close() <-- REMOVED |