feat(firestore): support firestore mongodb dialect - #177
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces MongoDB and Firestore support to the context generation guide and evaluation pipeline. It adds a new firestore_driver with custom MCP tools for listing collections and executing queries, integrates these tools dynamically, and implements a FirestoreConfigGenerator for EvalBench configurations. Feedback on these changes highlights a security vulnerability due to disabled SSL verification, bugs in schema flattening and dataset conversion, and logical issues such as ignoring query filters in the native client fallback and hardcoding default collection IDs.
| } | ||
|
|
||
| json_data = json.dumps(payload).encode("utf-8") | ||
| ssl_ctx = ssl._create_unverified_context() |
There was a problem hiding this comment.
Disabling SSL certificate verification via ssl._create_unverified_context() makes the connection vulnerable to Man-in-the-Middle (MitM) attacks. Use ssl.create_default_context() instead to ensure secure communication.
| ssl_ctx = ssl._create_unverified_context() | |
| ssl_ctx = ssl.create_default_context() |
| if isinstance(val, dict): | ||
| value_type = val.get("valueType") or (list(val.keys())[0] if val else "stringValue") | ||
| if value_type == "mapValue": | ||
| map_fields = val.get("mapValue", {}).get("fields", {}) | ||
| columns.extend(_flatten_schema(map_fields, field_name)) | ||
| else: | ||
| data_type = val.get("stringValue") or value_type.replace("Value", "").upper() | ||
| columns.append({"name": field_name, "type": data_type}) | ||
| elif isinstance(val, (dict, list)): | ||
| columns.append({"name": field_name, "type": "JSON"}) |
There was a problem hiding this comment.
The _flatten_schema function assumes any dictionary is a Firestore Value proto (e.g., {"stringValue": "..."}). However, when standard document sampling is used, to_dict() returns native Python dictionaries where nested maps are also native dictionaries (e.g., {"director": "Christopher Nolan"}). This causes native nested dictionaries to be incorrectly flattened as a single column with a type derived from the first key name (e.g., type "DIRECTOR"). We should distinguish between Firestore Value protos and native dictionaries, and recursively flatten native dictionaries.
| if isinstance(val, dict): | |
| value_type = val.get("valueType") or (list(val.keys())[0] if val else "stringValue") | |
| if value_type == "mapValue": | |
| map_fields = val.get("mapValue", {}).get("fields", {}) | |
| columns.extend(_flatten_schema(map_fields, field_name)) | |
| else: | |
| data_type = val.get("stringValue") or value_type.replace("Value", "").upper() | |
| columns.append({"name": field_name, "type": data_type}) | |
| elif isinstance(val, (dict, list)): | |
| columns.append({"name": field_name, "type": "JSON"}) | |
| if isinstance(val, dict): | |
| # Check if it's a Firestore Value proto (REST API format) | |
| is_proto = "valueType" in val or any(k.endswith("Value") for k in val.keys()) | |
| if is_proto: | |
| value_type = val.get("valueType") or (list(val.keys())[0] if val else "stringValue") | |
| if value_type == "mapValue": | |
| map_fields = val.get("mapValue", {}).get("fields", {}) | |
| columns.extend(_flatten_schema(map_fields, field_name)) | |
| else: | |
| data_type = val.get("stringValue") or value_type.replace("Value", "").upper() | |
| columns.append({"name": field_name, "type": data_type}) | |
| else: | |
| # Native nested dictionary: recursively flatten | |
| columns.extend(_flatten_schema(val, field_name)) | |
| elif isinstance(val, list): | |
| columns.append({"name": field_name, "type": "JSON"}) |
| query_obj = json.loads(query) if isinstance(query, str) else query | ||
| col_id = query_obj.get("find") or query_obj.get("collection") or "movies" | ||
| limit = query_obj.get("limit", 10) | ||
|
|
||
| col_ref = client.collection(col_id) | ||
| docs = list(col_ref.limit(limit).stream()) | ||
| results = [doc.to_dict() for doc in docs] |
There was a problem hiding this comment.
The native client fallback in firestore_execute_mongodb completely ignores the query filters, projections, and aggregation stages. It simply extracts the collection name and streams the first N documents. This will return incorrect/unfiltered results during evaluation or agent execution. If standard Firestore doesn't support MQL, we should either raise an error or implement actual execution via pymongo if a connection string is provided.
| if isinstance(golden_sql_raw, dict): | ||
| # Ensure dialect key matches the generator's expected dialect (e.g. mongodb) | ||
| golden_sql_dict = {} | ||
| for d_key, val in golden_sql_raw.items(): | ||
| golden_sql_dict[dialect] = val |
There was a problem hiding this comment.
When golden_sql_raw is a dictionary, the loop blindly overwrites golden_sql_dict[dialect] with the value of every key. If the target dialect is already present in the dictionary, it might be overwritten by another dialect's query depending on the iteration order. We should explicitly check if the target dialect exists in golden_sql_raw first.
| if isinstance(golden_sql_raw, dict): | |
| # Ensure dialect key matches the generator's expected dialect (e.g. mongodb) | |
| golden_sql_dict = {} | |
| for d_key, val in golden_sql_raw.items(): | |
| golden_sql_dict[dialect] = val | |
| if isinstance(golden_sql_raw, dict): | |
| # Ensure dialect key matches the generator's expected dialect (e.g. mongodb) | |
| golden_sql_dict = {} | |
| if dialect in golden_sql_raw: | |
| golden_sql_dict[dialect] = golden_sql_raw[dialect] | |
| elif golden_sql_raw: | |
| golden_sql_dict[dialect] = next(iter(golden_sql_raw.values())) |
| self.project = params.get("project") | ||
| self.database = params.get("database", "nl2sql-mflix") | ||
| self.connection_string = params.get("connection_string") | ||
| self.collection_ids = params.get("collection_ids") or params.get("table_ids") or ["orders", "products", "customers"] |
There was a problem hiding this comment.
Hardcoding default collection IDs ["orders", "products", "customers"] leaks domain-specific assumptions into the generic FirestoreConfigGenerator class. If no collections are specified, it is better to default to None or an empty list to avoid querying incorrect collections on other databases (like nl2sql-mflix).
| self.collection_ids = params.get("collection_ids") or params.get("table_ids") or ["orders", "products", "customers"] | |
| self.collection_ids = params.get("collection_ids") or params.get("table_ids") |
g-lynnzee
left a comment
g-lynnzee
left a comment
There was a problem hiding this comment.
Please remove files that are not intended for this PR.
| @@ -0,0 +1,236 @@ | |||
| import json | |||
There was a problem hiding this comment.
We talked about how this ideally belongs in firestore MCP? Why can't you make the code change there?
There was a problem hiding this comment.
Upstreaming Firestore support directly to MCP Toolbox is currently in progress in googleapis/mcp-toolbox#3826. In the interim, this local driver allows us to run Firestore MongoDB evaluations and unblocks context engineering workflows until the new toolbox version is released and pinned here.
| query_context_dict = MessageToDict( | ||
| query_context._pb, preserving_proto_field_name=True | ||
| ) | ||
| if hasattr(self, "build_custom_query_context"): |
There was a problem hiding this comment.
I dont' think you'll need this given https://github.com/Google...ntext-enrichment/pull/192
There was a problem hiding this comment.
Updated! Cleaned up BaseDBConfigGenerator to conform to the REST dictionary contract from #192, removing the custom query context extraction override.
| ) -> gda.DatasourceReferences: | ||
| return gda.DatasourceReferences() | ||
|
|
||
| def build_custom_query_context(self, context_set_id: str) -> dict[str, Any]: |
There was a problem hiding this comment.
Updated! Conformed FirestoreConfigGenerator.build_datasource_reference to return a dict[str, Any] matching the base class REST structure from #192, and removed hardcoded default collection IDs.
|
|
||
| def _convert_dataset(dataset_path: str, dialect: str) -> str: | ||
| """Reads simplified dataset and converts to EvalBench standard format.""" | ||
| """Reads simplified or standard EvalBench dataset and converts/normalizes to EvalBench standard format.""" |
There was a problem hiding this comment.
Please refactor with use of helper functions for readability.
This isn't related specifically to firestore. should we put it into its own PR?
There was a problem hiding this comment.
Refactored _convert_dataset to use clean, modular helper functions (_normalize_golden_sql, _convert_simplified_entry, _convert_standard_entry) and fixed validation.
pl04351820
changed the title
Summary of Updates & Fixes:
|
00addcc to
16cda8f
Compare
2 participants
Firestore Enterprise Edition (MongoDB Compatible API) Support with context bootstraping
The PR adds the architectural changes for dynamic tool loading with built-in Firestore tools.
With this enhancement, developers and data agents can:
ContextSetartifacts (Templates, Facets, Value Searches) for Firestore collections with MongoDB Compatible API.queryData).System Architecture & High-Level Flow
flowchart TD subgraph Local MCP Server [db-context-enrichment FastMCP] DynamicLoader[load_dynamic_mcp_tools] -- Reads --> ToolsYAML[autoctx/tools.yaml] DynamicLoader -- Registers --> LocalDriver[custom_tools/firestore_driver.py] end subgraph Evaluation & Generation Lifecycle NLQ[Natural Language Question] --> LocalCtx[Local ContextSet Injector] LocalCtx --> GDAPayload[GDA QueryData Payload Construction] GDAPayload --> GDAAPI[GDA REST API /queryData] GDAAPI --> GenMQL[Generated MongoDB MQL Query] GenMQL --> LocalDriver LocalDriver -- Executes MQL --> FirestoreDB[(Firestore Enterprise Edition)] FirestoreDB -- Returns Results --> Scorer[EvalBench Scorer & LLMRater] endKey Lifecycle Phases
ContextSet(Templates, Facets) are loaded and formatted into prompt guidelines.firestore_referencedatasource metadata and dispatched to GDA REST API.db.collection.aggregate(...)ordb.collection.find(...)) is executed against the target Firestore database, and results are evaluated by SetMatcher & LLMRater.Built-in Firestore MCP tool
Rather than depending on an external Toolbox MCP proxy or server binary, tool loading and execution for Firestore Enterprise Edition are handled in-process via a built-in dynamic tool loader (load_dynamic_mcp_tools()) in src/google/cloud/db_context_enrichment/main.py.
On startup, this loader scans autoctx/tools.yaml for tool definitions annotated with provider: local_mcp and dynamically constructs and registers native execution and schema inspection tools (such as firestore-list-collections and firestore-execute-mongodb) directly onto the FastMCP server instance using the local driver in custom_tools/firestore_driver.py.
This self-contained architecture allows data agents and evaluation runners to inspect collections and execute MongoDB MQL queries directly against Firestore Enterprise Edition via pymongo or google-cloud-firestore without requiring external MCP infrastructure.
Firestore Remote MCP tool could eventually replace the built-in tools.
Context Engineering & NoSQL Dialect Specification
NoSQL document databases differ significantly from relational SQL. This specification defines the NoSQL-specific rules for authoring
ContextSetitems.Templates (
references/template/mongodb.md)Templates map end-to-end natural language questions to executable NoSQL MQL queries.
{ "nl_query": "What is the total revenue for completed orders?", "sql": "db.orders.aggregate([{ $match: { status: 'completed' } }, { $group: { _id: null, totalRevenue: { $sum: '$total_amount' } } }, { $project: { _id: 0, totalRevenue: 1 } }])", "intent": "Total revenue for completed orders", "manifest": "Total revenue for orders with a given status", "parameterized": { "parameterized_sql": "db.orders.aggregate([{ $match: { status: '$1' } }, { $group: { _id: null, totalRevenue: { $sum: '$total_amount' } } }, { $project: { _id: 0, totalRevenue: 1 } }])", "parameterized_intent": "Total revenue for orders with status $1" } }Facets (
references/facet/mongodb.md)Facets provide modular, reusable filter snippets and field-value predicates.
{ "sql_snippet": "orders.payment_method = 'credit_card'", "intent": "credit card payment method is literal string 'credit_card'", "manifest": "Credit card payment method", "parameterized": { "parameterized_sql_snippet": "orders.payment_method = '$1'", "parameterized_intent": "payment method is $1" } }Key NoSQL Authoring Guidelines
items.price,customer.satisfaction,orders.status) instead of SQL table joins.{ $unwind: "$items" }stages when calculating metrics over array elements.'credit_card','In store'with space vs'In-store').ISODate('YYYY-MM-DDTHH:MM:SSZ')for date range matching.EvalBench Pipeline & Integration
In
evalbench/generators/models/query_data_api.py, local ContextSet files are parsed and appended to the prompt payload:In
db_generators/firestore.py,build_custom_query_context()passescontext_set_pathand routes execution through the PyMongo / Firestore runner.Verification & Evaluation Results
The implementation was validated using iterative hill-climbing context optimization across two benchmark datasets: