feat: add functionality to export Data Explorer filters as PyMongo query string

ChingEnLin · ChingEnLin · commit c21717c9a374 · 2026-03-12T11:37:59.000+01:00
diff --git a/backend/models/data_documents.py b/backend/models/data_documents.py
@@ -18,6 +18,9 @@ class DataDocumentsRequest(BaseModel):
     filter: Optional[DataDocumentsFilter] = None
     filters: Optional[List[DataDocumentsFilter]] = None
 
+class DataDocumentsQueryResponse(BaseModel):
+    query_code: str
+
 
 class DataDocumentsResponse(BaseModel):
     documents: List[Dict[str, Any]]
diff --git a/backend/routes/data_documents.py b/backend/routes/data_documents.py
@@ -4,6 +4,7 @@
 from models.data_documents import (
     DataDocumentsRequest,
     DataDocumentsResponse,
+    DataDocumentsQueryResponse,
     FindByIdRequest,
     FindByIdResponse,
     UpdateDocumentRequest,
@@ -17,6 +18,8 @@
 from services.data_documents_service import (
     find_document_by_id,
     fetch_documents,
+    build_mongo_query,
+    generate_mongo_query_string,
     update_document,
     get_single_document,
     insert_document,
@@ -49,6 +52,31 @@ def get_documents(
         filters=[f.model_dump() for f in body.filters] if body.filters else None,
     )
 
+@router.post("/documents/query_code", response_model=DataDocumentsQueryResponse)
+def get_documents_query_code(
+    body: DataDocumentsRequest = Body(...), authorization: str = Header(...)
+):
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid token format")
+    user_token = authorization.replace("Bearer ", "")
+    access_token = exchange_token_obo(user_token)
+    connection_string = get_connection_string(body.account_id, access_token)
+    
+    # Needs to initialize MongoClient briefly to access find_one if 'all' keys are used in the algorithm
+    from pymongo import MongoClient
+    client = MongoClient(connection_string)
+    db = client[body.database_name]
+    collection = db[body.collection_name]
+    
+    query = build_mongo_query(
+        collection=collection,
+        filter=body.filter.model_dump() if body.filter else None,
+        filters=[f.model_dump() for f in body.filters] if body.filters else None,
+    )
+    
+    query_str = generate_mongo_query_string(body.collection_name, query)
+    return DataDocumentsQueryResponse(query_code=query_str)
+
 
 @router.put("/documents", response_model=dict)
 def put_update_document(
diff --git a/backend/services/data_documents_service.py b/backend/services/data_documents_service.py
@@ -14,6 +14,41 @@
 
 ALL_DOCUMENTS_CACHES = [_find_by_id_cache]
 
+def _format_value_for_pymongo(v: any) -> str:
+    if isinstance(v, dict):
+        items = []
+        for key, val in v.items():
+            items.append(f"'{key}': {_format_value_for_pymongo(val)}")
+        return "{" + ", ".join(items) + "}"
+    elif isinstance(v, list):
+        items = [_format_value_for_pymongo(x) for x in v]
+        return "[" + ", ".join(items) + "]"
+    elif isinstance(v, datetime):
+        # E.g. datetime(2026, 3, 12, 10, 0, tzinfo=timezone.utc)
+        return repr(v)
+    elif isinstance(v, ObjectId):
+        return f"ObjectId('{str(v)}')"
+    elif isinstance(v, str):
+        # simple escape
+        if v.startswith("ObjectId(") or v.startswith("datetime("):
+            return v # fallback if it was already stringified somehow
+        
+        escaped = v.replace("'", "\\'")
+        return f"'{escaped}'"
+    elif isinstance(v, bool):
+        return "True" if v else "False"
+    elif v is None:
+        return "None"
+    elif hasattr(v, "pattern"): # Regex object
+        return f"re.compile(r'{v.pattern}', {v.flags})"
+    else:
+        return repr(v)
+
+def generate_mongo_query_string(collection_name: str, query: dict) -> str:
+    query_str = _format_value_for_pymongo(query)
+    # The user specifically requested exactly the db['project'].find({......}) format
+    return f"db['{collection_name}'].find({query_str})"
+
 
 def dict_diff(before, after):
     """
@@ -82,19 +117,7 @@ def json_dumps_safe(obj):
         return None
 
 
-def fetch_documents(
-    connection_string: str,
-    database_name: str,
-    collection_name: str,
-    page: int,
-    limit: int,
-    filter: dict = None,
-    filters: list = None,
-) -> DataDocumentsResponse:
-    client = MongoClient(connection_string)
-    db = client[database_name]
-    collection = db[collection_name]
-
+def build_mongo_query(collection, filter: dict = None, filters: list = None) -> dict:
     query = {}
     and_clauses = []
 
@@ -182,6 +205,24 @@ def get_query_for_filter(f: dict):
             query = and_clauses[0]
         else:
             query = {"$and": and_clauses}
+            
+    return query
+
+
+def fetch_documents(
+    connection_string: str,
+    database_name: str,
+    collection_name: str,
+    page: int,
+    limit: int,
+    filter: dict = None,
+    filters: list = None,
+) -> DataDocumentsResponse:
+    client = MongoClient(connection_string)
+    db = client[database_name]
+    collection = db[collection_name]
+
+    query = build_mongo_query(collection, filter, filters)
 
     total_documents = collection.count_documents(query)
     total_pages = max(1, (total_documents + limit - 1) // limit)
diff --git a/frontend/pages/DataExplorerPage.tsx b/frontend/pages/DataExplorerPage.tsx
@@ -1,7 +1,7 @@
 import React, { useState, useEffect, useCallback, useMemo, useRef } from 'react';
 import { useNavigate } from 'react-router-dom';
 import { SelectedResource, DbInfo, BreadcrumbItem, CosmosDBAccount, CollectionInfo } from '../types';
-import { getDocuments, getCollectionInfo, findDocumentById, getDatabasesForAccount, clearDocumentsCache, getSingleDocument } from '../services/dbService';
+import { getDocuments, getCollectionInfo, findDocumentById, getDatabasesForAccount, clearDocumentsCache, getSingleDocument, getDocumentsQueryCode } from '../services/dbService';
 import { extractSchemaTree, SchemaKeyNode } from '../utils/schemaUtils';
 import MongoIcon from '../components/icons/MongoIcon';
 import { isEqual, omit } from 'lodash';
@@ -351,6 +351,9 @@ const DataExplorerPage: React.FC<DataExplorerPageProps> = ({
   const maxCardWidth = 600;
   const prevPinnedCount = useRef(pinnedDocuments.length);
 
+  // --- Export Query State ---
+  const [copiedQuery, setCopiedQuery] = useState(false);
+
   // --- Theme State ---
   const { theme, toggleTheme } = useTheme();
 
@@ -1352,13 +1355,45 @@ const DataExplorerPage: React.FC<DataExplorerPageProps> = ({
                         )}
                       </div>
                     ))}
-                    <button
-                      onClick={addFilter}
-                      disabled={!selectedCollection || isFetchingSchema}
-                      className="w-full flex justify-center items-center gap-1 p-2 text-sm font-medium text-slate-600 dark:text-slate-300 bg-slate-100 dark:bg-slate-700/50 hover:bg-slate-200 dark:hover:bg-slate-600 border border-slate-300 dark:border-slate-600 rounded-lg transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
-                    >
-                      <span>+ Add Filter</span>
-                    </button>
+                    <div className="flex gap-2">
+                      <button
+                        onClick={addFilter}
+                        disabled={!selectedCollection || isFetchingSchema}
+                        className="flex-1 flex justify-center items-center gap-1 p-2 text-sm font-medium text-slate-600 dark:text-slate-300 bg-slate-100 dark:bg-slate-700/50 hover:bg-slate-200 dark:hover:bg-slate-600 border border-slate-300 dark:border-slate-600 rounded-lg transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
+                      >
+                        <span>+ Add Filter</span>
+                      </button>
+                      <button
+                        onClick={() => {
+                          if (!selectedCollection) return;
+
+                          const validFilters = debouncedFilters.map(f => ({
+                            key: f.key,
+                            value: getCoercedFilterValue(f.value),
+                            operator: f.operator || 'equals',
+                            type: f.type || 'string'
+                          }));
+
+                          getDocumentsQueryCode(selectedCollection, currentResource, undefined, validFilters)
+                            .then((queryCodeStr) => {
+                              navigator.clipboard.writeText(queryCodeStr).then(() => {
+                                setCopiedQuery(true);
+                                setTimeout(() => setCopiedQuery(false), 2000);
+                              });
+                            })
+                            .catch(err => {
+                              console.error("Failed to generate query code:", err);
+                              alert("Failed to generate query code: " + err.message);
+                            });
+                        }}
+                        disabled={!selectedCollection || isFetchingSchema}
+                        className="flex justify-center items-center gap-2 p-2 px-3 text-sm font-medium text-slate-600 dark:text-slate-300 bg-white dark:bg-slate-700 hover:bg-slate-50 dark:hover:bg-slate-600 border border-slate-300 dark:border-slate-600 rounded-lg transition-colors disabled:opacity-50 disabled:cursor-not-allowed"
+                        title="Export current filter criteria as Python script (MongoDB PyMongo)"
+                      >
+                        {copiedQuery ? <CheckIcon className="w-4 h-4 text-green-500" /> : <FileCopyIcon className="w-4 h-4" />}
+                        <span>Export</span>
+                      </button>
+                    </div>
                   </div>
                 </div>
                 <div className="flex-grow overflow-hidden">
diff --git a/frontend/services/dbService.ts b/frontend/services/dbService.ts
@@ -431,6 +431,53 @@ export const getDocuments = async (
   return response.json();
 };
 
+export const getDocumentsQueryCode = async (
+  collectionName: string,
+  resource: SelectedResource,
+  filter?: { key: string, value: any, operator?: string, type?: string },
+  filters?: { key: string, value: any, operator?: string, type?: string }[]
+): Promise<string> => {
+  if (!USE_MSAL_AUTH) {
+    return Promise.resolve("db['" + collectionName + "'].find({\n  // Exporting is only fully supported dynamically in remote mode.\n})");
+  }
+
+  const accessToken = await getAuthenticatedToken();
+
+  const bodyQuery: any = {
+    account_id: resource.accountId,
+    database_name: resource.databaseName,
+    collection_name: collectionName,
+    page: 1,
+    limit: 1,
+  };
+
+  if (filter && ((filter.value !== '' && filter.value !== null && filter.value !== undefined) || filter.operator === 'exists' || filter.operator === 'not_exists')) {
+    bodyQuery.filter = filter;
+  }
+
+  if (filters && filters.length > 0) {
+    bodyQuery.filters = filters.filter(f => (f.value !== '' && f.value !== null && f.value !== undefined) || f.operator === 'exists' || f.operator === 'not_exists');
+  }
+
+  const response = await fetch(`${API_BASE_URL}/data/documents/query_code`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'Authorization': `Bearer ${accessToken}`,
+    },
+    body: JSON.stringify(bodyQuery),
+  });
+
+  if (!response.ok) {
+    const errorData = await response.json().catch(() => ({}));
+    const errorMessage = errorData.detail || errorData.message || `Failed to generate query code. Status: ${response.status}`;
+    throw new Error(errorMessage);
+  }
+
+  const result = await response.json();
+  return result.query_code;
+};
+
 /**
  * Finds a single document by its ID, searching across all provided collections.
  * @param documentId The string representation of the document's ObjectId.
