File size: 5,061 Bytes
7ccd501
369b363
7ccd501
 
 
 
 
369b363
 
 
7ccd501
 
 
 
 
 
 
615412d
369b363
615412d
 
 
 
 
 
 
 
 
 
 
 
 
 
369b363
 
615412d
 
 
 
 
 
369b363
615412d
7ccd501
369b363
 
7ccd501
369b363
7ccd501
 
 
369b363
7ccd501
369b363
 
 
7ccd501
 
369b363
 
 
 
 
 
7ccd501
 
 
369b363
 
 
 
 
 
7ccd501
615412d
369b363
615412d
7ccd501
 
 
 
 
615412d
 
7ccd501
615412d
 
7ccd501
 
 
 
 
369b363
7ccd501
369b363
7ccd501
369b363
 
 
7ccd501
 
369b363
7ccd501
369b363
 
7ccd501
369b363
7ccd501
369b363
7ccd501
 
 
 
 
615412d
369b363
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import logging
from typing import Dict, List, Union, Any
from bson import ObjectId
import ast
import json
import re
from pymongo import MongoClient
from threading import Lock
from urllib.parse import urlparse
from pymongo.uri_parser import parse_uri

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO
)
logger = logging.getLogger("mongo_service")

# ==============================================================================
# MONGO CONNECTION MANAGER
# ==============================================================================
class MongoConnectionManager:
    def __init__(self):
        self._clients: Dict[str, MongoClient] = {}
        self._lock = Lock()

    def get_client(self, uri: str) -> MongoClient:
        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]}...")
                    self._clients[uri] = MongoClient(
                        uri, 
                        serverSelectionTimeoutMS=5000, 
                        connectTimeoutMS=5000, 
                        maxPoolSize=50
                    )
        return self._clients[uri]

mongo_manager = MongoConnectionManager()

# ==============================================================================
# HELPER FUNCTIONS
# ==============================================================================
def convert_oid(obj):
    """Recursively convert $oid or ObjectId 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]
    return obj

def sanitize_json_input(query_input: Union[str, Dict, List]) -> Union[Dict, List]:
    """Parses input string/dict/list into Python objects, handling ObjectIds."""
    if isinstance(query_input, (dict, list)): return query_input

    query_str = str(query_input).strip()
    match = re.search(r"```(?:json|javascript)?\s*(.*?)\s*```", query_str, re.DOTALL)
    if match: query_str = match.group(1).strip()

    # Sanitize JS specific wrappers
    query_str = re.sub(r"ObjectId\(['\"]([a-fA-F0-9]+)['\"]\)", r"'\1'", query_str)
    query_str = re.sub(r"ISODate\(['\"](.*?)['\"]\)", r"'\1'", query_str)
    
    try:
        return json.loads(query_str)
    except:
        try:
            return ast.literal_eval(query_str)
        except Exception as e:
            logger.error(f"Failed to parse query: {e}")
            raise ValueError(f"Parsing error: {str(e)}")

# ==============================================================================
# MAIN EXECUTION LOGIC
# ==============================================================================
def execute_mongo_operation(
    mongo_uri: str, 
    db_name: str, 
    collection_name: str, 
    query: Union[Dict, List],
    limited: bool = False,
    limit_rows: int = 20
):
    client = mongo_manager.get_client(mongo_uri)
    
    try:
        db = client[db_name]
        collection = db[collection_name]
        results = []
        
        # --- Pipeline (List) ---
        if isinstance(query, list):
            pipeline = query
            if limited:
                if not (pipeline and "$limit" in pipeline[-1]):
                    pipeline.append({"$limit": limit_rows})
            cursor = collection.aggregate(pipeline, allowDiskUse=True)
            results = list(cursor)
            
        # --- Find (Dict) ---
        elif isinstance(query, dict):
            cursor = collection.find(query)
            if limited: cursor = cursor.limit(limit_rows)
            results = list(cursor)
            
        else:
            raise ValueError("Query must be a Pipeline (List) or Filter (Dict)")

        return results
    except Exception as e:
        logger.error(f"DB Execution Error: {e}")
        raise e
    
def extract_database_name(uri: str) -> str:
    """
    Robustly extracts database name from a MongoDB URI.
    1. Tries PyMongo's strict parser.
    2. Falls back to standard URL path extraction (like your TS code).
    """
    if not uri:
        return None

    # Method 1: PyMongo Strict Parser
    try:
        parsed = parse_uri(uri)
        if parsed.get('database'):
            return parsed['database']
    except Exception:
        pass # Fallback if PyMongo complains about the URI format

    # Method 2: Generic URL Parsing (Mimics your TypeScript logic)
    try:
        # Standard URL parsing
        result = urlparse(uri)
        
        # The database name is usually the 'path' part (e.g., /api_platform)
        path = result.path
        
        # Remove leading slash if present
        if path.startswith('/'):
            path = path[1:]
            
        # If empty, return None
        return path if path else None
    except Exception:
        return None