admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

MongoDB analytics toolkit (#548)

Browse source 
Design System: https://github.com/ArcadeAI/Design-System/pull/31
Docs: https://github.com/ArcadeAI/docs/pull/439

<img width="832" height="797" alt="Screenshot 2025-09-11 at 2 35 36 PM"
src="https://github.com/user-attachments/assets/198f194f-d41c-41a0-bedf-b77bccd70dfa"
/>
This commit is contained in:
Evan Tahler 2025-09-12 18:41:23 -07:00 committed by GitHub
parent f50e05aa9b
commit 766a262a25
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 2273 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.github/workflows/test-toolkits.yml vendored
View file

@ -22,7 +22,7 @@ jobs:
run: |
# Find all directories in toolkits/ that have a pyproject.toml
TOOLKITS=$(find toolkits -maxdepth 1 -type d -not -name "toolkits" -exec test -f {}/pyproject.toml \; -exec basename {} \; | jq -R -s -c 'split("\n")[:-1]')
TOOLKITS_WITH_GHA_SECRETS='["postgres", "clickhouse"]'
TOOLKITS_WITH_GHA_SECRETS='["postgres", "clickhouse", "mongodb"]'
TOOLKITS_WITHOUT_GHA_SECRETS=$(echo "$TOOLKITS" | jq -c --argjson with "$TOOLKITS_WITH_GHA_SECRETS" '[.[] | select(. as $t | $with | index($t) | not)]')
echo "Found toolkits: $TOOLKITS"
echo "Found toolkits without GHA secrets: $TOOLKITS_WITHOUT_GHA_SECRETS"
@ -104,6 +104,7 @@ jobs:
env:
TEST_POSTGRES_DATABASE_CONNECTION_STRING: ${{ secrets.TEST_POSTGRES_DATABASE_CONNECTION_STRING }} # TODO: dynamically only load the `TEST_${{ matrix.toolkit }}_DATABASE_CONNECTION_STRING secret`
TEST_CLICKHOUSE_DATABASE_CONNECTION_STRING: ${{ secrets.TEST_CLICKHOUSE_DATABASE_CONNECTION_STRING }}
TEST_MONGODB_CONNECTION_STRING: ${{ secrets.TEST_MONGODB_CONNECTION_STRING }}
run: |
# If there's a custom test_setup.sh file, run it
if [ -f tests/test_setup.sh ]; then

53
toolkits/mongodb/Makefile Normal file
View file

@ -0,0 +1,53 @@
.PHONY: help
help:
@echo "🛠️ github Commands:\n"
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}'
.PHONY: install
install: ## Install the uv environment and install all packages with dependencies
@echo "🚀 Creating virtual environment and installing all packages using uv"
@uv sync --active --all-extras --no-sources
@uv run pre-commit install
@echo "✅ All packages and dependencies installed via uv"
.PHONY: install-local
install-local: ## Install the uv environment and install all packages with dependencies with local Arcade sources
@echo "🚀 Creating virtual environment and installing all packages using uv"
@uv sync --active --all-extras
@uv run pre-commit install
@echo "✅ All packages and dependencies installed via uv"
.PHONY: build
build: clean-build ## Build wheel file using poetry
@echo "🚀 Creating wheel file"
uv build
.PHONY: clean-build
clean-build: ## clean build artifacts
@echo "🗑️ Cleaning dist directory"
rm -rf dist
.PHONY: test
test: ## Test the code with pytest
@echo "🚀 Testing code: Running pytest"
@uv run pytest -W ignore -v --cov --cov-config=pyproject.toml --cov-report=xml
.PHONY: coverage
coverage: ## Generate coverage report
@echo "coverage report"
coverage report
@echo "Generating coverage report"
coverage html
.PHONY: bump-version
bump-version: ## Bump the version in the pyproject.toml file by a patch version
@echo "🚀 Bumping version in pyproject.toml"
uv version --bump patch
.PHONY: check
check: ## Run code quality tools.
@echo "🚀 Linting code: Running pre-commit"
@uv run pre-commit run -a
@echo "🚀 Static type checking: Running mypy"
@uv run mypy --config-file=pyproject.toml

0
toolkits/mongodb/arcade_mongodb/__init__.py Normal file
View file

118
toolkits/mongodb/arcade_mongodb/database_engine.py Normal file
View file

@ -0,0 +1,118 @@
from typing import Any, ClassVar
from arcade_tdk.errors import RetryableToolError
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
from pymongo.errors import ServerSelectionTimeoutError
MAX_RECORDS_RETURNED = 1000
TEST_QUERY = {"ping": 1}
class DatabaseEngine:
_instance: ClassVar[None] = None
_clients: ClassVar[dict[str, AsyncIOMotorClient]] = {}
@classmethod
async def get_instance(cls, connection_string: str) -> AsyncIOMotorClient:
key = connection_string
if key not in cls._clients:
cls._clients[key] = AsyncIOMotorClient(connection_string)
# try a simple query to see if the connection is valid
try:
admin_db = cls._clients[key].admin
await admin_db.command(TEST_QUERY)
return cls._clients[key]
except ServerSelectionTimeoutError:
# close and try again
cls._clients[key].close()
cls._clients[key] = AsyncIOMotorClient(connection_string)
try:
admin_db = cls._clients[key].admin
await admin_db.command(TEST_QUERY)
return cls._clients[key]
except Exception as e:
raise RetryableToolError(
f"Connection failed: {e}",
developer_message="Connection to MongoDB failed.",
additional_prompt_content="Check the connection string and try again.",
) from e
@classmethod
async def get_database(cls, connection_string: str, database_name: str) -> Any:
client = await cls.get_instance(connection_string)
class DatabaseContextManager:
def __init__(self, client: AsyncIOMotorClient, database_name: str) -> None:
self.client = client
self.database_name = database_name
self.database = client[database_name]
async def __aenter__(self) -> AsyncIOMotorDatabase:
return self.database
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
# Connection cleanup is handled by the client cache
pass
return DatabaseContextManager(client, database_name)
@classmethod
async def cleanup(cls) -> None:
"""Clean up all cached clients. Call this when shutting down."""
for client in cls._clients.values():
client.close()
cls._clients.clear()
@classmethod
def clear_cache(cls) -> None:
"""Clear the client cache without closing clients. Use with caution."""
cls._clients.clear()
@classmethod
def sanitize_query_params(
cls,
database_name: str,
collection_name: str,
filter_dict: dict[str, Any] | None,
projection: dict[str, Any] | None,
sort: list[dict[str, Any]] | None,
limit: int,
skip: int,
) -> tuple[
str, str, dict[str, Any], dict[str, Any] | None, list[dict[str, Any]] | None, int, int
]:
if not database_name:
raise RetryableToolError(
"Database name is required.",
developer_message="Database name cannot be empty.",
)
if not collection_name:
raise RetryableToolError(
"Collection name is required.",
developer_message="Collection name cannot be empty.",
)
if filter_dict is None:
filter_dict = {}
if limit > MAX_RECORDS_RETURNED:
raise RetryableToolError(
f"Limit is too high. Maximum is {MAX_RECORDS_RETURNED}.",
)
if skip < 0:
raise RetryableToolError(
"Skip must be greater than or equal to 0.",
developer_message="Skip must be greater than or equal to 0.",
)
if limit <= 0:
raise RetryableToolError(
"Limit must be greater than 0.",
developer_message="Limit must be greater than 0.",
)
return database_name, collection_name, filter_dict, projection, sort, limit, skip

0
toolkits/mongodb/arcade_mongodb/tools/__init__.py Normal file
View file

367
toolkits/mongodb/arcade_mongodb/tools/mongodb.py Normal file
View file

@ -0,0 +1,367 @@
import json
from typing import Annotated, Any
from arcade_tdk import ToolContext, tool
from arcade_tdk.errors import RetryableToolError
from ..database_engine import MAX_RECORDS_RETURNED, DatabaseEngine
from .utils import (
_infer_schema_from_docs,
_parse_json_list_parameter,
_parse_json_parameter,
_serialize_document,
)
# class UserStatus(str, Enum):
# """User status enumeration."""
# ACTIVE = "active"
# INACTIVE = "inactive"
# BANNED = "banned"
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def discover_databases(
context: ToolContext,
) -> list[str]:
"""Discover all the databases in the MongoDB instance."""
client = await DatabaseEngine.get_instance(context.get_secret("MONGODB_CONNECTION_STRING"))
databases = await client.list_database_names()
# Filter out admin and config databases by default
databases = [db for db in databases if db not in ["admin", "config", "local"]]
return databases
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def discover_collections(
context: ToolContext,
database_name: Annotated[str, "The database name to discover collections in"],
) -> list[str]:
"""Discover all the collections in the MongoDB database when the list of collections is not known.
ALWAYS use this tool before any other tool that requires a collection name.
"""
async with await DatabaseEngine.get_database(
context.get_secret("MONGODB_CONNECTION_STRING"), database_name
) as db:
collections = await db.list_collection_names()
return list(collections)
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def get_collection_schema(
context: ToolContext,
database_name: Annotated[str, "The database name to get the collection schema of"],
collection_name: Annotated[str, "The collection to get the schema of"],
sample_size: Annotated[
int,
f"The number of documents to sample for schema discovery (default: {MAX_RECORDS_RETURNED})",
] = MAX_RECORDS_RETURNED,
) -> dict[str, Any]:
"""
Get the schema/structure of a MongoDB collection by sampling documents.
Since MongoDB is schema-less, this tool samples a configurable number of documents
to infer the schema structure and data types.
This tool should ALWAYS be used before executing any query. All collections in the query must be discovered first using the <discover_collections> tool.
"""
async with await DatabaseEngine.get_database(
context.get_secret("MONGODB_CONNECTION_STRING"), database_name
) as db:
collection = db[collection_name]
# Sample documents at random to infer schema
# Use MongoDB's $sample aggregation to get random documents
sample_docs = []
async for doc in collection.aggregate([{"$sample": {"size": sample_size}}]):
sample_docs.append(doc)
if not sample_docs:
return {"message": "Collection is empty", "schema": {}}
# Infer schema from sampled documents
schema = _infer_schema_from_docs(sample_docs)
return {
"total_documents_sampled": len(sample_docs),
"sample_size_requested": sample_size,
"schema": schema,
}
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def find_documents(
context: ToolContext,
database_name: Annotated[str, "The database name to query"],
collection_name: Annotated[str, "The collection name to query"],
filter_dict: Annotated[
str | None,
'MongoDB filter/query as JSON string. Leave None for no filter (find all documents). Example: \'{"status": "active", "age": {"$gte": 18}}\'',
] = None,
projection: Annotated[
str | None,
'Fields to include/exclude as JSON string. Use 1 to include, 0 to exclude. Example: \'{"name": 1, "email": 1, "_id": 0}\'. Leave None to include all fields.',
] = None,
sort: Annotated[
list[str] | None,
'Sort criteria as list of JSON strings, each containing \'field\' and \'direction\' keys. Use 1 for ascending, -1 for descending. Example: [\'{"field": "name", "direction": 1}\', \'{"field": "created_at", "direction": -1}\']',
] = None,
limit: Annotated[
int,
f"The maximum number of documents to return. Default: {MAX_RECORDS_RETURNED}.",
] = MAX_RECORDS_RETURNED,
skip: Annotated[int, "The number of documents to skip. Default: 0."] = 0,
) -> list[str]:
"""
Find documents in a MongoDB collection.
ONLY use this tool if you have already loaded the schema of the collection you need to query.
Use the <get_collection_schema> tool to load the schema if not already known.
Returns a list of JSON strings, where each string represents a document from the collection (tools cannot return complex types).
When running queries, follow these rules which will help avoid errors:
* Always specify projection to limit fields returned if you don't need all data.
* Always sort your results by the most important fields first. If you aren't sure, sort by '_id'.
* Use appropriate MongoDB query operators for complex filtering ($gte, $lte, $in, $regex, etc.).
* Be mindful of case sensitivity when querying string fields.
* Use indexes when possible (typically on _id and commonly queried fields).
"""
# Initialize variables to avoid UnboundLocalError in exception handler
parsed_filter = None
parsed_projection = None
parsed_sort = None
try:
# Parse JSON string inputs
parsed_filter = _parse_json_parameter(filter_dict, "filter_dict")
parsed_projection = _parse_json_parameter(projection, "projection")
parsed_sort = _parse_json_list_parameter(sort, "sort")
(
database_name,
collection_name,
parsed_filter,
parsed_projection,
parsed_sort,
limit,
skip,
) = DatabaseEngine.sanitize_query_params(
database_name=database_name,
collection_name=collection_name,
filter_dict=parsed_filter,
projection=parsed_projection,
sort=parsed_sort,
limit=limit,
skip=skip,
)
async with await DatabaseEngine.get_database(
context.get_secret("MONGODB_CONNECTION_STRING"), database_name
) as db:
collection = db[collection_name]
# Build the query
cursor = collection.find(parsed_filter, parsed_projection)
if parsed_sort:
# Convert list of dicts to list of tuples for MongoDB sort
sort_tuples = [(str(item["field"]), int(item["direction"])) for item in parsed_sort]
cursor = cursor.sort(sort_tuples)
cursor = cursor.skip(skip).limit(limit)
# Execute query and collect results
documents = []
async for doc in cursor:
# Convert ObjectId and other non-serializable types to strings
doc = _serialize_document(doc)
documents.append(json.dumps(doc))
return documents
except RetryableToolError:
# Re-raise RetryableToolError as-is to preserve JSON validation messages
raise
except Exception as e:
raise RetryableToolError(
f"Query failed: {e}",
developer_message=f"Query failed with parameters: database_name={database_name}, collection_name={collection_name}, filter_dict={parsed_filter}, projection={parsed_projection}, sort={parsed_sort}, limit={limit}, skip={skip}.",
additional_prompt_content="Load the collection schema <get_collection_schema> or use the <discover_collections> tool to discover the collections and try again.",
retry_after_ms=10,
) from e
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def count_documents(
context: ToolContext,
database_name: Annotated[str, "The database name to query"],
collection_name: Annotated[str, "The collection name to query"],
filter_dict: Annotated[
str | None,
'MongoDB filter/query as JSON string. Leave None for no filter (count all documents). Example: \'{"status": "active"}\'',
] = None,
) -> int:
"""Count documents in a MongoDB collection matching the given filter."""
parsed_filter = None
try:
# Parse JSON string input
parsed_filter = _parse_json_parameter(filter_dict, "filter_dict") or {}
async with await DatabaseEngine.get_database(
context.get_secret("MONGODB_CONNECTION_STRING"), database_name
) as db:
collection = db[collection_name]
count = await collection.count_documents(parsed_filter)
return int(count)
except RetryableToolError:
# Re-raise RetryableToolError as-is to preserve JSON validation messages
raise
except Exception as e:
raise RetryableToolError(
f"Count query failed: {e}",
developer_message=f"Count query failed with parameters: database_name={database_name}, collection_name={collection_name}, filter_dict={parsed_filter}.",
additional_prompt_content="Check the collection name and filter criteria and try again.",
retry_after_ms=10,
) from e
@tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
async def aggregate_documents(
context: ToolContext,
database_name: Annotated[str, "The database name to query"],
collection_name: Annotated[str, "The collection name to query"],
pipeline: Annotated[
list[str],
'MongoDB aggregation pipeline as a list of JSON strings, each representing a stage. Example: [\'{"$match": {"status": "active"}}\', \'{"$group": {"_id": "$category", "count": {"$sum": 1}}}\']',
],
limit: Annotated[
int,
f"The maximum number of results to return from the aggregation. Default: {MAX_RECORDS_RETURNED}.",
] = MAX_RECORDS_RETURNED,
) -> list[str]:
"""
Execute a MongoDB aggregation pipeline on a collection.
ONLY use this tool if you have already loaded the schema of the collection you need to query.
Use the <get_collection_schema> tool to load the schema if not already known.
Returns a list of JSON strings, where each string represents a result document from the aggregation (tools cannot return complex types).
Aggregation pipelines allow for complex data processing including:
* $match - filter documents
* $group - group documents and perform calculations
* $project - reshape documents
* $sort - sort documents
* $limit - limit results
* $lookup - join with other collections
* And many more stages
"""
parsed_pipeline = None
try:
# Parse JSON string inputs
parsed_pipeline = _parse_json_list_parameter(pipeline, "pipeline")
if parsed_pipeline is None:
raise RetryableToolError( # noqa: TRY301
"Pipeline cannot be empty",
developer_message="The pipeline parameter is required and cannot be None",
)
async with await DatabaseEngine.get_database(
context.get_secret("MONGODB_CONNECTION_STRING"), database_name
) as db:
collection = db[collection_name]
# Add limit to pipeline if not already present
pipeline_with_limit = parsed_pipeline.copy()
has_limit = any("$limit" in stage for stage in pipeline_with_limit)
if not has_limit:
pipeline_with_limit.append({"$limit": limit})
# Execute aggregation
cursor = collection.aggregate(pipeline_with_limit)
documents = []
async for doc in cursor:
# Convert ObjectId and other non-serializable types to strings
doc = _serialize_document(doc)
documents.append(json.dumps(doc))
return documents
except RetryableToolError:
# Re-raise RetryableToolError as-is to preserve JSON validation messages
raise
except Exception as e:
raise RetryableToolError(
f"Aggregation query failed: {e}",
developer_message=f"Aggregation query failed with parameters: database_name={database_name}, collection_name={collection_name}, pipeline={parsed_pipeline}, limit={limit}.",
additional_prompt_content="Check the aggregation pipeline syntax and collection schema, then try again.",
retry_after_ms=10,
) from e
# @tool(requires_secrets=["MONGODB_CONNECTION_STRING"])
# async def update_user_status(
# context: ToolContext,
# database_name: Annotated[str, "The database name containing the users collection"],
# collection_name: Annotated[str, "The collection name containing user documents"],
# user_id: Annotated[str, "The _id of the user to update"],
# status: Annotated[UserStatus, "The new status for the user"],
# ) -> dict[str, Any]:
# """
# [CUSTOM TOOL]
# Update the status of a user in the MongoDB collection.
# This tool updates a user document by setting the status field to the specified value.
# The status must be one of: active, inactive, or banned.
# Returns information about the update operation including the number of documents modified.
# """
# try:
# async with await DatabaseEngine.get_database(
# context.get_secret("MONGODB_CONNECTION_STRING"), database_name
# ) as db:
# collection = db[collection_name]
# # cast the user_id to int if it looks like an integer
# if isinstance(user_id, str) and user_id.isdigit():
# user_id = int(user_id)
# result = await collection.update_one(
# {"_id": user_id}, {"$set": {"status": status.value}}
# )
# print(result)
# if result.matched_count == 0:
# return {
# "success": False,
# "message": f"No user found with _id: {user_id}",
# "matched_count": 0,
# "modified_count": 0,
# }
# return {
# "success": True,
# "message": f"User status updated to '{status.value}'",
# "user_id": user_id,
# "new_status": status.value,
# "matched_count": result.matched_count,
# "modified_count": result.modified_count,
# }
# except Exception as e:
# raise RetryableToolError(
# f"Failed to update user status: {e}",
# developer_message=f"Update operation failed with parameters: database_name={database_name}, collection_name={collection_name}, user_id={user_id}, status={status}.",
# additional_prompt_content="Check the database name, collection name, and user ID, then try again.",
# retry_after_ms=10,
# ) from e

281
toolkits/mongodb/arcade_mongodb/tools/utils.py Normal file
View file

@ -0,0 +1,281 @@
import json
from datetime import datetime
from typing import Any
from arcade_tdk.errors import RetryableToolError
from bson import ObjectId
def _validate_no_write_operations(obj: Any, parameter_name: str, path: str = "") -> None:
"""
Recursively validate that an object doesn't contain MongoDB write operations.
Args:
obj: The object to validate
parameter_name: Name of the parameter for error messages
path: Current path in the object (for nested validation)
Raises:
RetryableToolError: If write operations are detected
"""
# MongoDB write/update operators that should be blocked
WRITE_OPERATORS = {
# Update operators
"$set",
"$unset",
"$inc",
"$mul",
"$rename",
"$min",
"$max",
"$currentDate",
"$addToSet",
"$pop",
"$pull",
"$push",
"$pullAll",
"$each",
"$slice",
"$sort",
"$position",
"$bit",
"$isolated",
# Array update operators
"$",
"$[]",
"$[<identifier>]",
# Pipeline update operators
"$addFields",
"$replaceRoot",
"$replaceWith",
# Aggregation stages that can modify (in case they're misused)
"$out",
"$merge",
# Other potentially dangerous operators
"$where", # Can execute JavaScript
}
if isinstance(obj, dict):
for key, value in obj.items():
current_path = f"{path}.{key}" if path else key
# Special check for $where operator which can execute JavaScript (check this first)
if key == "$where":
raise RetryableToolError(
f"JavaScript execution operator '$where' not allowed in {parameter_name}",
developer_message=f"Found '$where' operator at path '{current_path}' in parameter '{parameter_name}'. JavaScript execution is not allowed for security reasons.",
additional_prompt_content=f"The {parameter_name} parameter cannot use the $where operator. Use other query operators instead.",
)
# Check if this key is a write operator
if key in WRITE_OPERATORS:
raise RetryableToolError(
f"Write operation '{key}' not allowed in {parameter_name}",
developer_message=f"Found write operation '{key}' at path '{current_path}' in parameter '{parameter_name}'. Only read operations are allowed.",
additional_prompt_content=f"The {parameter_name} parameter cannot contain write operations like '{key}'. Use only query/read operations such as $match, $gte, $lte, $in, $regex, etc.",
)
# Recursively validate nested objects
_validate_no_write_operations(value, parameter_name, current_path)
elif isinstance(obj, list):
for i, item in enumerate(obj):
current_path = f"{path}[{i}]" if path else f"[{i}]"
_validate_no_write_operations(item, parameter_name, current_path)
def _parse_json_parameter(
json_string: str | None, parameter_name: str, validate_read_only: bool = True
) -> Any | None:
"""
Parse a JSON string parameter with proper error handling and optional write operation validation.
Args:
json_string: The JSON string to parse (can be None)
parameter_name: Name of the parameter for error messages
validate_read_only: Whether to validate that no write operations are present
Returns:
Parsed JSON object or None if json_string is None
Raises:
RetryableToolError: If JSON parsing fails or write operations are detected
"""
if json_string is None:
return None
try:
parsed_obj = json.loads(json_string)
# Validate that no write operations are present
if validate_read_only and parsed_obj is not None:
_validate_no_write_operations(parsed_obj, parameter_name)
except json.JSONDecodeError as e:
raise RetryableToolError(
f"Invalid JSON in {parameter_name}: {e}",
developer_message=f"Failed to parse JSON string for parameter '{parameter_name}': {json_string}. Error: {e}",
additional_prompt_content=f"Please provide valid JSON for the {parameter_name} parameter. Check for proper escaping of quotes and valid JSON syntax.",
) from e
else:
return parsed_obj
def _validate_aggregation_pipeline(pipeline: list[Any], parameter_name: str) -> None:
"""
Validate that an aggregation pipeline only contains read operations.
Args:
pipeline: The aggregation pipeline to validate
parameter_name: Name of the parameter for error messages
Raises:
RetryableToolError: If write operations are detected in the pipeline
"""
# MongoDB aggregation stages that can modify data
WRITE_STAGES = {
"$out",
"$merge", # These stages write to collections
}
# Aggregation stages that are potentially dangerous
DANGEROUS_STAGES = {
"$where", # Can execute JavaScript
}
for i, stage in enumerate(pipeline):
if isinstance(stage, dict):
for stage_name in stage:
if stage_name in WRITE_STAGES:
raise RetryableToolError(
f"Write stage '{stage_name}' not allowed in {parameter_name}",
developer_message=f"Found write stage '{stage_name}' at pipeline index {i} in parameter '{parameter_name}'. Only read operations are allowed.",
additional_prompt_content=f"The {parameter_name} parameter cannot contain write stages like '{stage_name}'. Use only read stages such as $match, $group, $project, $sort, $limit, etc.",
)
if stage_name in DANGEROUS_STAGES:
raise RetryableToolError(
f"Dangerous stage '{stage_name}' not allowed in {parameter_name}",
developer_message=f"Found dangerous stage '{stage_name}' at pipeline index {i} in parameter '{parameter_name}'. JavaScript execution is not allowed for security reasons.",
additional_prompt_content=f"The {parameter_name} parameter cannot use the {stage_name} stage. Use other aggregation stages instead.",
)
# Also validate the stage content for write operations
_validate_no_write_operations(
stage[stage_name], f"{parameter_name}[{i}].{stage_name}"
)
def _parse_json_list_parameter(
json_strings: list[str] | None, parameter_name: str, validate_read_only: bool = True
) -> list[Any] | None:
"""
Parse a list of JSON strings with proper error handling and optional write operation validation.
Args:
json_strings: List of JSON strings to parse (can be None)
parameter_name: Name of the parameter for error messages
validate_read_only: Whether to validate that no write operations are present
Returns:
List of parsed JSON objects or None if json_strings is None
Raises:
RetryableToolError: If JSON parsing fails for any string or write operations are detected
"""
if json_strings is None:
return None
try:
parsed_list = [json.loads(json_str) for json_str in json_strings]
# Validate that no write operations are present
if validate_read_only and parsed_list is not None:
# Special handling for pipeline parameters
if parameter_name == "pipeline":
_validate_aggregation_pipeline(parsed_list, parameter_name)
else:
# For non-pipeline lists, validate each item
for i, item in enumerate(parsed_list):
_validate_no_write_operations(item, f"{parameter_name}[{i}]")
except json.JSONDecodeError as e:
raise RetryableToolError(
f"Invalid JSON in {parameter_name}: {e}",
developer_message=f"Failed to parse JSON string list for parameter '{parameter_name}': {json_strings}. Error: {e}",
additional_prompt_content=f"Please provide valid JSON strings for the {parameter_name} parameter. Each string must be valid JSON with proper escaping of quotes.",
) from e
else:
return parsed_list
def _infer_schema_from_docs(docs: list[dict[str, Any]]) -> dict[str, Any]:
"""Infer schema structure from a list of documents."""
schema: dict[str, Any] = {}
for doc in docs:
_update_schema_with_doc(schema, doc)
# Convert sets to lists for serialization
for key in schema:
if isinstance(schema[key]["types"], set):
schema[key]["types"] = list(schema[key]["types"])
return schema
def _update_schema_with_doc(schema: dict[str, Any], doc: dict[str, Any], prefix: str = "") -> None:
"""Recursively update schema with document structure."""
for key, value in doc.items():
full_key = f"{prefix}.{key}" if prefix else key
if full_key not in schema:
schema[full_key] = {
"types": set(),
"sample_values": [],
"null_count": 0,
"total_count": 0,
}
schema[full_key]["total_count"] += 1
if value is None:
schema[full_key]["null_count"] += 1
schema[full_key]["types"].add("null")
else:
value_type = type(value).__name__
schema[full_key]["types"].add(value_type)
# Store sample values (limit to 3 unique samples)
if (
len(schema[full_key]["sample_values"]) < 3
and value not in schema[full_key]["sample_values"]
):
schema[full_key]["sample_values"].append(value)
# Handle nested objects
if isinstance(value, dict):
_update_schema_with_doc(schema, value, full_key)
elif isinstance(value, list) and value and isinstance(value[0], dict):
# Handle arrays of objects by sampling the first few
for i, item in enumerate(value[:3]): # Sample first 3 array items
if isinstance(item, dict):
_update_schema_with_doc(schema, item, f"{full_key}[{i}]")
def _serialize_document(doc: dict[str, Any]) -> dict[str, Any]:
"""Convert MongoDB document to JSON-serializable format."""
if isinstance(doc, dict):
result = {}
for key, value in doc.items():
result[key] = _serialize_document(value)
return result
elif isinstance(doc, list):
return [_serialize_document(item) for item in doc]
elif isinstance(doc, ObjectId):
return str(doc)
elif isinstance(doc, datetime):
return doc.isoformat()
else:
return doc

190
toolkits/mongodb/evals/eval_mongodb.py Normal file
View file

@ -0,0 +1,190 @@
# RUN ME WITH `uv run arcade evals evals --host api.arcade.dev`
import arcade_mongodb
from arcade_evals import (
BinaryCritic,
EvalRubric,
EvalSuite,
ExpectedToolCall,
SimilarityCritic,
tool_eval,
)
from arcade_mongodb.tools.mongodb import (
aggregate_documents,
count_documents,
discover_collections,
discover_databases,
find_documents,
get_collection_schema,
)
from arcade_tdk import ToolCatalog
# Evaluation rubric
rubric = EvalRubric(
fail_threshold=0.85,
warn_threshold=0.95,
)
catalog = ToolCatalog()
catalog.add_module(arcade_mongodb)
@tool_eval()
def mongodb_eval_suite() -> EvalSuite:
suite = EvalSuite(
name="MongoDB Tools Evaluation",
system_message=(
"You are an AI assistant with access to MongoDB tools. "
"Use them to help the user with their tasks."
),
catalog=catalog,
rubric=rubric,
)
suite.add_case(
name="Discover databases",
user_message="What databases are available in my MongoDB instance?",
expected_tool_calls=[
ExpectedToolCall(func=discover_databases, args={}),
],
rubric=rubric,
)
suite.add_case(
name="Discover collections",
user_message="What collections are in the 'admin' database?",
expected_tool_calls=[
ExpectedToolCall(func=discover_collections, args={"database_name": "admin"}),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=1.0),
],
)
suite.add_case(
name="Get collection schema (single tool call)",
user_message="Get the schema of the 'system.users' collection in the 'admin' database.",
expected_tool_calls=[
ExpectedToolCall(
func=get_collection_schema,
args={"database_name": "admin", "collection_name": "system.users"},
),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=0.5),
BinaryCritic(critic_field="collection_name", weight=0.5),
],
)
suite.add_case(
name="Find documents (direct call)",
user_message="Find documents in the 'startup_log' collection of the 'local' database, limited to 5 results.",
additional_messages=[
{
"role": "user",
"content": "You can call find_documents directly without discovering collections first for this test.",
}
],
expected_tool_calls=[
ExpectedToolCall(
func=find_documents,
args={
"database_name": "local",
"collection_name": "startup_log",
"limit": 5,
},
),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=0.33),
BinaryCritic(critic_field="collection_name", weight=0.33),
BinaryCritic(critic_field="limit", weight=0.34),
],
)
suite.add_case(
name="Count documents",
user_message="Count all documents in the 'startup_log' collection of the 'local' database.",
additional_messages=[
{
"role": "user",
"content": "You can call count_documents directly without discovering collections first for this test.",
}
],
expected_tool_calls=[
ExpectedToolCall(
func=count_documents,
args={
"database_name": "local",
"collection_name": "startup_log",
},
),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=0.5),
BinaryCritic(critic_field="collection_name", weight=0.5),
],
)
suite.add_case(
name="Count documents with filter",
user_message="Count documents in the 'startup_log' collection of the 'local' database where the level is 'INFO'.",
additional_messages=[
{
"role": "user",
"content": "You can call count_documents directly without discovering collections first for this test.",
}
],
expected_tool_calls=[
ExpectedToolCall(
func=count_documents,
args={
"database_name": "local",
"collection_name": "startup_log",
"filter_dict": '{"level": "INFO"}',
},
),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=0.25),
BinaryCritic(critic_field="collection_name", weight=0.25),
SimilarityCritic(critic_field="filter_dict", weight=0.5),
],
)
suite.add_case(
name="Aggregate documents",
user_message="Group documents in the 'startup_log' collection of the 'local' database by level and count them.",
additional_messages=[
{
"role": "user",
"content": "You can call aggregate_documents directly without discovering collections first for this test.",
}
],
expected_tool_calls=[
ExpectedToolCall(
func=aggregate_documents,
args={
"database_name": "local",
"collection_name": "startup_log",
"pipeline": [
'{"$group": {"_id": "$level", "count": {"$sum": 1}}}',
],
},
),
],
rubric=rubric,
critics=[
BinaryCritic(critic_field="database_name", weight=0.2),
BinaryCritic(critic_field="collection_name", weight=0.2),
SimilarityCritic(critic_field="pipeline", weight=0.6),
],
)
return suite

62
toolkits/mongodb/pyproject.toml Normal file
View file

@ -0,0 +1,62 @@
[build-system]
requires = [ "hatchling",]
build-backend = "hatchling.build"
[project]
name = "arcade_mongodb"
version = "0.1.0"
description = "Tools to query and explore a MongoDB database"
requires-python = ">=3.10"
dependencies = [
"arcade-tdk>=2.0.0,<3.0.0",
"pymongo>=4.10.1",
"pydantic>=2.11.7",
"motor>=3.6.0",
]
[[project.authors]]
name = "evantahler"
email = "support@arcade.dev"
[project.optional-dependencies]
dev = [
"arcade-ai[evals]>=2.0.0,<3.0.0",
"arcade-serve>=2.0.0,<3.0.0",
"pytest>=8.3.0,<8.4.0",
"pytest-cov>=4.0.0,<4.1.0",
"pytest-mock>=3.11.1,<3.12.0",
"pytest-asyncio>=0.24.0,<0.25.0",
"mypy>=1.5.1,<1.6.0",
"pre-commit>=3.4.0,<3.5.0",
"tox>=4.11.1,<4.12.0",
"ruff>=0.7.4,<0.8.0",
]
# Use local path sources for arcade libs when working locally
[tool.uv.sources]
arcade-ai = { path = "../../", editable = true }
arcade-serve = { path = "../../libs/arcade-serve/", editable = true }
arcade-tdk = { path = "../../libs/arcade-tdk/", editable = true }
[tool.mypy]
files = [ "arcade_mongodb/**/*.py",]
python_version = "3.10"
disallow_untyped_defs = "True"
disallow_any_unimported = "True"
no_implicit_optional = "True"
check_untyped_defs = "True"
warn_return_any = "True"
warn_unused_ignores = "True"
show_error_codes = "True"
ignore_missing_imports = "True"
[tool.pytest.ini_options]
testpaths = [ "tests",]
asyncio_default_fixture_loop_scope = "function"
[tool.coverage.report]
skip_empty = true
[tool.hatch.build.targets.wheel]
packages = [ "arcade_mongodb",]

0
toolkits/mongodb/tests/__init__.py Normal file
View file

45
toolkits/mongodb/tests/conftest.py Normal file
View file

@ -0,0 +1,45 @@
import os
import shutil
import subprocess
from os import environ
import pytest_asyncio
from arcade_mongodb.database_engine import DatabaseEngine
TEST_MONGODB_CONNECTION_STRING = (
environ.get("TEST_MONGODB_CONNECTION_STRING") or "mongodb://localhost:27017"
)
@pytest_asyncio.fixture(autouse=True)
async def restore_database():
"""Restore the database from the dump before each test."""
dump_file = f"{os.path.dirname(__file__)}/dump.js"
# Execute the MongoDB dump script to restore test data
mongosh_path = shutil.which("mongosh")
if not mongosh_path:
raise RuntimeError("mongosh executable not found in PATH")
result = subprocess.run(
[mongosh_path, TEST_MONGODB_CONNECTION_STRING, dump_file],
check=True,
capture_output=True,
text=True,
)
if result.returncode != 0:
print(f"Error loading test data: {result.stderr}")
raise RuntimeError(f"Failed to load test data: {result.stderr}")
yield # This allows tests to run
# Optional cleanup could go here if needed
@pytest_asyncio.fixture(autouse=True)
async def cleanup_engines():
"""Clean up database engines after each test to prevent connection leaks."""
yield
await DatabaseEngine.cleanup()

378
toolkits/mongodb/tests/dump.js Normal file
View file

@ -0,0 +1,378 @@
// MongoDB test data dump - equivalent to PostgreSQL dump.sql
// This script sets up test data for the MongoDB toolkit
// Switch to test database
use('test_database');
// Clear existing data
db.users.drop();
db.messages.drop();
// Create users collection with data
db.users.insertMany([
{
_id: 1,
name: 'Alice',
email: 'alice@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$tMg1Rd3IEDnp3iFKrqsF4Dsbw6/Cbf6seRB/H5bhaPg$zZj5yn4x3D3O3mDHcW2aczQNiYfAs3cw21XMEIgkF0E',
created_at: new Date('2024-09-01T20:49:38.759Z'),
updated_at: new Date('2024-09-02T03:49:39.927Z'),
status: 'active'
},
{
_id: 2,
name: 'Bob',
email: 'bob@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$CvOMK1WUd99R7kYXpiBPNYw4OQP53pYIgeMnwz92mrE$HPthId4phMoPT1TWuCRHHCr9BSQA8XoUkQuB1HZsqTY',
created_at: new Date('2024-09-02T17:49:23.377Z'),
updated_at: new Date('2024-09-02T17:49:23.377Z'),
status: 'active'
},
{
_id: 3,
name: 'Charlie',
email: 'charlie@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$paCAAD1HVZkncP/WvecuUO6zFXp2/8BISpgr5rXRxps$M5kBFc9JHHGNw9SXnPu2ggpJY0mFFCska7TXMrllndo',
created_at: new Date('2024-09-03T10:30:15.123Z'),
updated_at: new Date('2024-09-03T10:30:15.123Z'),
status: 'active'
},
{
_id: 4,
name: 'Diana',
email: 'diana@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$xyz123ABC456DEF789GHI$SampleHashForDiana123',
created_at: new Date('2024-09-04T14:20:30.654Z'),
updated_at: new Date('2024-09-04T14:20:30.654Z'),
status: 'active'
},
{
_id: 5,
name: 'Evan',
email: 'evan@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$evanHash123$EvanPasswordHash456',
created_at: new Date('2024-09-05T09:15:45.987Z'),
updated_at: new Date('2024-09-05T09:15:45.987Z'),
status: 'active'
},
{
_id: 6,
name: 'Fiona',
email: 'fiona@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$fionaHash456$FionaPasswordHash789',
created_at: new Date('2024-09-06T16:45:12.345Z'),
updated_at: new Date('2024-09-06T16:45:12.345Z'),
status: 'active'
},
{
_id: 7,
name: 'George',
email: 'george@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$georgeHash789$GeorgePasswordHash012',
created_at: new Date('2024-09-07T11:30:25.876Z'),
updated_at: new Date('2024-09-07T11:30:25.876Z'),
status: 'active'
},
{
_id: 8,
name: 'Helen',
email: 'helen@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$helenHash012$HelenPasswordHash345',
created_at: new Date('2024-09-08T13:25:40.234Z'),
updated_at: new Date('2024-09-08T13:25:40.234Z'),
status: 'active'
},
{
_id: 9,
name: 'Ian',
email: 'ian@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$ianHash345$IanPasswordHash678',
created_at: new Date('2024-09-09T08:40:55.765Z'),
updated_at: new Date('2024-09-09T08:40:55.765Z'),
status: 'active'
},
{
_id: 10,
name: 'Julia',
email: 'julia@example.com',
password_hash: '$argon2id$v=19$m=65536,t=2,p=1$juliaHash678$JuliaPasswordHash901',
created_at: new Date('2024-09-10T15:55:18.123Z'),
updated_at: new Date('2024-09-10T15:55:18.123Z'),
status: 'active'
}
]);
// Create messages collection with data
db.messages.insertMany([
// User 1 (Alice) - 3 messages
{
_id: 1,
body: 'Hello everyone!',
user_id: 1,
created_at: new Date('2025-01-10T10:00:00.000Z'),
updated_at: new Date('2025-01-10T10:00:00.000Z')
},
{
_id: 2,
body: 'How is everyone doing today?',
user_id: 1,
created_at: new Date('2025-01-10T11:30:00.000Z'),
updated_at: new Date('2025-01-10T11:30:00.000Z')
},
{
_id: 3,
body: 'Great to see you all here!',
user_id: 1,
created_at: new Date('2025-01-10T14:15:00.000Z'),
updated_at: new Date('2025-01-10T14:15:00.000Z')
},
// User 2 (Bob) - 2 messages
{
_id: 4,
body: 'Hi Alice! Doing well, thanks for asking.',
user_id: 2,
created_at: new Date('2025-01-10T11:35:00.000Z'),
updated_at: new Date('2025-01-10T11:35:00.000Z')
},
{
_id: 5,
body: 'Anyone up for a game later?',
user_id: 2,
created_at: new Date('2025-01-10T16:20:00.000Z'),
updated_at: new Date('2025-01-10T16:20:00.000Z')
},
// User 3 (Charlie) - 3 messages
{
_id: 6,
body: 'Count me in for the game!',
user_id: 3,
created_at: new Date('2025-01-10T16:25:00.000Z'),
updated_at: new Date('2025-01-10T16:25:00.000Z')
},
{
_id: 7,
body: 'What time works for everyone?',
user_id: 3,
created_at: new Date('2025-01-10T16:30:00.000Z'),
updated_at: new Date('2025-01-10T16:30:00.000Z')
},
{
_id: 8,
body: 'I can play around 8 PM',
user_id: 3,
created_at: new Date('2025-01-10T17:00:00.000Z'),
updated_at: new Date('2025-01-10T17:00:00.000Z')
},
// User 4 (Diana) - 2 messages
{
_id: 9,
body: '8 PM works for me too!',
user_id: 4,
created_at: new Date('2025-01-10T17:05:00.000Z'),
updated_at: new Date('2025-01-10T17:05:00.000Z')
},
{
_id: 10,
body: 'What game should we play?',
user_id: 4,
created_at: new Date('2025-01-10T17:10:00.000Z'),
updated_at: new Date('2025-01-10T17:10:00.000Z')
},
// User 5 (Evan) - 13 messages (including 10 additional ones)
{
_id: 11,
body: 'I suggest we try the new arcade game!',
user_id: 5,
created_at: new Date('2025-01-10T17:15:00.000Z'),
updated_at: new Date('2025-01-10T17:15:00.000Z')
},
{
_id: 12,
body: 'It has great multiplayer features',
user_id: 5,
created_at: new Date('2025-01-10T17:20:00.000Z'),
updated_at: new Date('2025-01-10T17:20:00.000Z')
},
{
_id: 13,
body: 'Perfect timing for a weekend session',
user_id: 5,
created_at: new Date('2025-01-10T18:00:00.000Z'),
updated_at: new Date('2025-01-10T18:00:00.000Z')
},
{
_id: 26,
body: 'Just finished setting up the game server!',
user_id: 5,
created_at: new Date('2025-01-10T20:00:00.000Z'),
updated_at: new Date('2025-01-10T20:00:00.000Z')
},
{
_id: 27,
body: 'Everyone should be able to connect now',
user_id: 5,
created_at: new Date('2025-01-10T20:05:00.000Z'),
updated_at: new Date('2025-01-10T20:05:00.000Z')
},
{
_id: 28,
body: 'I added some custom maps too',
user_id: 5,
created_at: new Date('2025-01-10T20:10:00.000Z'),
updated_at: new Date('2025-01-10T20:10:00.000Z')
},
{
_id: 29,
body: 'The graphics look amazing on this new version',
user_id: 5,
created_at: new Date('2025-01-10T20:15:00.000Z'),
updated_at: new Date('2025-01-10T20:15:00.000Z')
},
{
_id: 30,
body: 'Hope you all enjoy the new features',
user_id: 5,
created_at: new Date('2025-01-10T20:20:00.000Z'),
updated_at: new Date('2025-01-10T20:20:00.000Z')
},
{
_id: 31,
body: 'I also set up a leaderboard system',
user_id: 5,
created_at: new Date('2025-01-10T20:25:00.000Z'),
updated_at: new Date('2025-01-10T20:25:00.000Z')
},
{
_id: 32,
body: 'We can track high scores now',
user_id: 5,
created_at: new Date('2025-01-10T20:30:00.000Z'),
updated_at: new Date('2025-01-10T20:30:00.000Z')
},
{
_id: 33,
body: 'The game supports up to 8 players simultaneously',
user_id: 5,
created_at: new Date('2025-01-10T20:35:00.000Z'),
updated_at: new Date('2025-01-10T20:35:00.000Z')
},
{
_id: 34,
body: 'I tested it earlier and it runs smoothly',
user_id: 5,
created_at: new Date('2025-01-10T20:40:00.000Z'),
updated_at: new Date('2025-01-10T20:40:00.000Z')
},
{
_id: 35,
body: 'Cannot wait to see everyone online tonight!',
user_id: 5,
created_at: new Date('2025-01-10T20:45:00.000Z'),
updated_at: new Date('2025-01-10T20:45:00.000Z')
},
// User 6 (Fiona) - 2 messages
{
_id: 14,
body: 'Sounds like fun! I love arcade games.',
user_id: 6,
created_at: new Date('2025-01-10T18:05:00.000Z'),
updated_at: new Date('2025-01-10T18:05:00.000Z')
},
{
_id: 15,
body: 'Should I bring snacks?',
user_id: 6,
created_at: new Date('2025-01-10T18:10:00.000Z'),
updated_at: new Date('2025-01-10T18:10:00.000Z')
},
// User 7 (George) - 3 messages
{
_id: 16,
body: 'Snacks are always welcome!',
user_id: 7,
created_at: new Date('2025-01-10T18:15:00.000Z'),
updated_at: new Date('2025-01-10T18:15:00.000Z')
},
{
_id: 17,
body: 'I can bring some drinks',
user_id: 7,
created_at: new Date('2025-01-10T18:20:00.000Z'),
updated_at: new Date('2025-01-10T18:20:00.000Z')
},
{
_id: 18,
body: 'This is going to be awesome',
user_id: 7,
created_at: new Date('2025-01-10T19:00:00.000Z'),
updated_at: new Date('2025-01-10T19:00:00.000Z')
},
// User 8 (Helen) - 2 messages
{
_id: 19,
body: 'I agree! Cannot wait for the game night.',
user_id: 8,
created_at: new Date('2025-01-10T19:05:00.000Z'),
updated_at: new Date('2025-01-10T19:05:00.000Z')
},
{
_id: 20,
body: 'Should we set up a Discord call?',
user_id: 8,
created_at: new Date('2025-01-10T19:10:00.000Z'),
updated_at: new Date('2025-01-10T19:10:00.000Z')
},
// User 9 (Ian) - 3 messages
{
_id: 21,
body: 'Discord would be perfect for voice chat',
user_id: 9,
created_at: new Date('2025-01-10T19:15:00.000Z'),
updated_at: new Date('2025-01-10T19:15:00.000Z')
},
{
_id: 22,
body: 'I will create a server for us',
user_id: 9,
created_at: new Date('2025-01-10T19:20:00.000Z'),
updated_at: new Date('2025-01-10T19:20:00.000Z')
},
{
_id: 23,
body: 'Link will be shared in a few minutes',
user_id: 9,
created_at: new Date('2025-01-10T19:25:00.000Z'),
updated_at: new Date('2025-01-10T19:25:00.000Z')
},
// User 10 (Julia) - 2 messages
{
_id: 24,
body: 'Thanks Ian! You are the best.',
user_id: 10,
created_at: new Date('2025-01-10T19:30:00.000Z'),
updated_at: new Date('2025-01-10T19:30:00.000Z')
},
{
_id: 25,
body: 'See you all at 8 PM!',
user_id: 10,
created_at: new Date('2025-01-10T19:35:00.000Z'),
updated_at: new Date('2025-01-10T19:35:00.000Z')
},{
_id: 99,
body: 'You are a mean jerk, you shithead!',
user_id: 10,
created_at: new Date('2025-01-10T19:35:00.000Z'),
updated_at: new Date('2025-01-10T19:35:00.000Z')
}
]);
// Create indexes for better performance (equivalent to PostgreSQL indexes)
db.users.createIndex({ "name": 1 }, { unique: true });
db.users.createIndex({ "email": 1 }, { unique: true });
db.messages.createIndex({ "user_id": 1 });
db.messages.createIndex({ "created_at": 1 });
print("MongoDB test data setup completed successfully!");
print("Users collection: " + db.users.countDocuments());
print("Messages collection: " + db.messages.countDocuments());

222
toolkits/mongodb/tests/test_json_validation.py Normal file
View file

@ -0,0 +1,222 @@
import pytest
from arcade_core.errors import ToolExecutionError
from arcade_mongodb.tools.mongodb import aggregate_documents, count_documents, find_documents
from arcade_tdk import ToolContext, ToolSecretItem
from arcade_tdk.errors import RetryableToolError
from .conftest import TEST_MONGODB_CONNECTION_STRING
@pytest.fixture
def mock_context():
context = ToolContext()
context.secrets = []
context.secrets.append(
ToolSecretItem(key="MONGODB_CONNECTION_STRING", value=TEST_MONGODB_CONNECTION_STRING)
)
return context
@pytest.mark.asyncio
async def test_invalid_json_in_filter_dict(mock_context) -> None:
"""Test that invalid JSON in filter_dict returns a reasonable error message."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active",}', # Invalid JSON - trailing comma
limit=1,
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in filter_dict" in error_message
# Check that the developer message contains helpful information
assert "filter_dict" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_invalid_json_in_projection(mock_context) -> None:
"""Test that invalid JSON in projection returns a reasonable error message."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
projection='{"name": 1, "email": 1,}', # Invalid JSON - trailing comma
limit=1,
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in projection" in error_message
# Check that the error message is helpful
assert "projection" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_invalid_json_in_sort(mock_context) -> None:
"""Test that invalid JSON in sort returns a reasonable error message."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
sort=['{"field": "name", "direction": 1,}'], # Invalid JSON - trailing comma
limit=1,
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in sort" in error_message
# Check that the error message is helpful
assert "sort" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_invalid_json_in_count_filter(mock_context) -> None:
"""Test that invalid JSON in count_documents filter returns a reasonable error message."""
with pytest.raises(RetryableToolError) as exc_info:
await count_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active",}', # Invalid JSON - trailing comma
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in filter_dict" in error_message
# Check that the error message is helpful
assert "filter_dict" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_invalid_json_in_pipeline(mock_context) -> None:
"""Test that invalid JSON in aggregation pipeline returns a reasonable error message."""
with pytest.raises(RetryableToolError) as exc_info:
await aggregate_documents(
mock_context,
database_name="test_database",
collection_name="users",
pipeline=['{"$match": {"status": "active",}}'], # Invalid JSON - trailing comma
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in pipeline" in error_message
# Check that the error message is helpful
assert "pipeline" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_malformed_json_string(mock_context) -> None:
"""Test various malformed JSON strings return reasonable error messages."""
test_cases = [
('{"unclosed": "string}', "Unterminated string"),
('{"missing_quotes": value}', "Expecting"),
('{missing_closing_brace: "value"}', "Expecting"),
('[{"array": "with"}, {"missing": }]', "Expecting"),
]
for invalid_json, expected_error_fragment in test_cases:
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict=invalid_json,
limit=1,
)
# Check that this is a JSON validation error
error_message = str(exc_info.value)
assert "Invalid JSON in filter_dict" in error_message
# Check that specific error details are included when expected
if expected_error_fragment:
assert (
expected_error_fragment in error_message
or expected_error_fragment in exc_info.value.developer_message
)
# Ensure helpful context is provided
assert "filter_dict" in exc_info.value.developer_message
assert "JSON" in exc_info.value.additional_prompt_content
assert "escaping" in exc_info.value.additional_prompt_content
# Check that the original JSON error is in the cause chain
assert exc_info.value.__cause__ is not None
@pytest.mark.asyncio
async def test_valid_json_does_not_error(mock_context) -> None:
"""Test that valid JSON does not raise JSON parsing errors."""
# This should not raise a JSON parsing error (might raise other errors, but not JSON-related)
try:
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active"}',
projection='{"name": 1, "_id": 0}',
sort=['{"field": "name", "direction": 1}'],
limit=1,
)
# If we get here, JSON parsing succeeded
assert isinstance(result, list)
except (ToolExecutionError, RetryableToolError) as e:
# If we get an error, it should not be about JSON parsing
# Check both the outer error and any nested error
error_message = str(e)
nested_message = str(e.__cause__) if e.__cause__ else ""
assert "Invalid JSON" not in error_message
assert "Invalid JSON" not in nested_message
@pytest.mark.asyncio
async def test_duplicate_keys_are_valid_json(mock_context) -> None:
"""Test that duplicate keys in JSON are valid (Python JSON allows this)."""
# This should NOT raise a JSON parsing error because duplicate keys are valid JSON
try:
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"duplicate": "key", "duplicate": "key"}', # Valid JSON - last value wins
limit=1,
)
# If we get here, JSON parsing succeeded (might get empty results, but no JSON error)
assert isinstance(result, list)
except (ToolExecutionError, RetryableToolError) as e:
# If we get an error, it should not be about JSON parsing
error_message = str(e)
nested_message = str(e.__cause__) if e.__cause__ else ""
assert "Invalid JSON" not in error_message
assert "Invalid JSON" not in nested_message

294
toolkits/mongodb/tests/test_mongodb.py Normal file
View file

@ -0,0 +1,294 @@
import json
import pytest
from arcade_mongodb.database_engine import DatabaseEngine
from arcade_mongodb.tools.mongodb import (
# UserStatus,
aggregate_documents,
count_documents,
discover_collections,
discover_databases,
find_documents,
get_collection_schema,
# update_user_status,
)
from arcade_tdk import ToolContext, ToolSecretItem
from arcade_tdk.errors import RetryableToolError
from .conftest import TEST_MONGODB_CONNECTION_STRING
@pytest.fixture
def mock_context():
context = ToolContext()
context.secrets = []
context.secrets.append(
ToolSecretItem(key="MONGODB_CONNECTION_STRING", value=TEST_MONGODB_CONNECTION_STRING)
)
return context
@pytest.mark.asyncio
async def test_discover_databases(mock_context) -> None:
databases = await discover_databases(mock_context)
assert isinstance(databases, list)
# Should not include system databases like admin, config, local
for db in databases:
assert db not in ["admin", "config", "local"]
@pytest.mark.asyncio
async def test_discover_collections(mock_context) -> None:
collections = await discover_collections(mock_context, "test_database")
assert "users" in collections
assert "messages" in collections
@pytest.mark.asyncio
async def test_get_collection_schema(mock_context) -> None:
schema_result = await get_collection_schema(
mock_context, "test_database", "users", sample_size=10
)
assert "schema" in schema_result
assert "total_documents_sampled" in schema_result
assert schema_result["total_documents_sampled"] == 10 # We have 10 users
schema = schema_result["schema"]
assert "_id" in schema
assert "name" in schema
assert "email" in schema
assert "password_hash" in schema
assert "status" in schema
assert "created_at" in schema
assert "updated_at" in schema
@pytest.mark.asyncio
async def test_find_documents_basic(mock_context) -> None:
# Find all users
result = await find_documents(
mock_context, database_name="test_database", collection_name="users", limit=10
)
assert len(result) == 10
# Parse JSON strings to check contents
docs = [json.loads(doc_str) for doc_str in result]
assert all("name" in doc for doc in docs)
assert all("email" in doc for doc in docs)
@pytest.mark.asyncio
async def test_find_documents_with_filter(mock_context) -> None:
# Find active users
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active"}',
limit=10,
)
assert len(result) == 10 # All users in dump are active
docs = [json.loads(doc_str) for doc_str in result]
assert all(doc["status"] == "active" for doc in docs)
@pytest.mark.asyncio
async def test_find_documents_with_projection(mock_context) -> None:
# Find users with only name and email
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
projection='{"name": 1, "email": 1, "_id": 0}',
limit=10,
)
assert len(result) == 10
docs = [json.loads(doc_str) for doc_str in result]
for doc in docs:
assert "name" in doc
assert "email" in doc
assert "_id" not in doc
assert "password_hash" not in doc
@pytest.mark.asyncio
async def test_find_documents_with_sort(mock_context) -> None:
# Find users sorted by _id descending
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
sort=['{"field": "_id", "direction": -1}'],
limit=3,
)
assert len(result) == 3
docs = [json.loads(doc_str) for doc_str in result]
ids = [doc["_id"] for doc in docs]
assert ids == [10, 9, 8] # Descending order
@pytest.mark.asyncio
async def test_count_documents(mock_context) -> None:
# Count all users
count = await count_documents(
mock_context, database_name="test_database", collection_name="users"
)
assert count == 10
# Count active users
active_count = await count_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active"}',
)
assert active_count == 10
@pytest.mark.asyncio
async def test_aggregate_documents(mock_context) -> None:
# Aggregate to count users by status
pipeline = ['{"$group": {"_id": "$status", "count": {"$sum": 1}}}', '{"$sort": {"count": -1}}']
result = await aggregate_documents(
mock_context, database_name="test_database", collection_name="users", pipeline=pipeline
)
assert len(result) == 1 # Only active users
# Should be sorted by count descending
doc = json.loads(result[0])
assert doc["_id"] == "active"
assert doc["count"] == 10
@pytest.mark.asyncio
async def test_find_documents_with_skip_and_limit(mock_context) -> None:
# Test pagination
result1 = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
sort=['{"field": "name", "direction": 1}'],
limit=2,
skip=0,
)
result2 = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
sort=['{"field": "name", "direction": 1}'],
limit=2,
skip=2,
)
assert len(result1) == 2
assert len(result2) == 2
docs1 = [json.loads(doc_str) for doc_str in result1]
docs2 = [json.loads(doc_str) for doc_str in result2]
assert docs1[0]["name"] == "Alice"
assert docs1[1]["name"] == "Bob"
assert docs2[0]["name"] == "Charlie"
assert docs2[1]["name"] == "Diana"
@pytest.mark.asyncio
async def test_error_handling_invalid_database(mock_context) -> None:
# Test with non-existent database - should not error but return empty results
collections = await discover_collections(mock_context, "nonexistent_database")
assert collections == []
@pytest.mark.asyncio
async def test_error_handling_invalid_collection(mock_context) -> None:
# Test with non-existent collection
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="nonexistent_collection",
limit=10,
)
assert result == []
@pytest.mark.asyncio
async def test_sanitize_query_params() -> None:
# Test parameter validation
with pytest.raises(RetryableToolError) as e:
DatabaseEngine.sanitize_query_params("", "users", {}, None, None, 10, 0)
assert "Database name is required" in str(e.value)
with pytest.raises(RetryableToolError) as e:
DatabaseEngine.sanitize_query_params("test_db", "", {}, None, None, 10, 0)
assert "Collection name is required" in str(e.value)
with pytest.raises(RetryableToolError) as e:
DatabaseEngine.sanitize_query_params(
"test_db", "users", {}, None, None, 2000, 0
) # Too high limit
assert "Limit is too high" in str(e.value)
# @pytest.mark.asyncio
# async def test_update_user_status_success(mock_context) -> None:
# """Test successful user status update."""
# # First, find a user to update
# users = await find_documents(
# mock_context, database_name="test_database", collection_name="users", limit=1
# )
# assert len(users) > 0
# user_doc = json.loads(users[0])
# user_id = user_doc["_id"]
# # Update user status to inactive
# result = await update_user_status(
# mock_context,
# database_name="test_database",
# collection_name="users",
# user_id=user_id,
# status=UserStatus.INACTIVE,
# )
# assert result["success"] is True
# assert result["user_id"] == user_id
# assert result["new_status"] == "inactive"
# assert result["matched_count"] == 1
# assert result["modified_count"] == 1
# # Verify the update by finding the user again
# # Convert user_id to int since the test data uses integer IDs
# user_id_int = int(user_id)
# updated_users = await find_documents(
# mock_context,
# database_name="test_database",
# collection_name="users",
# filter_dict=f'{{"_id": {user_id_int}}}',
# limit=1,
# )
# assert len(updated_users) == 1
# updated_user = json.loads(updated_users[0])
# assert updated_user["status"] == "inactive"
# @pytest.mark.asyncio
# async def test_update_user_status_user_not_found(mock_context) -> None:
# """Test updating status for non-existent user."""
# result = await update_user_status(
# mock_context,
# database_name="test_database",
# collection_name="users",
# user_id="nonexistent_user_id",
# status=UserStatus.BANNED,
# )
# assert result["success"] is False
# assert "No user found with _id" in result["message"]
# assert result["matched_count"] == 0
# assert result["modified_count"] == 0

12
toolkits/mongodb/tests/test_setup.sh Executable file
View file

@ -0,0 +1,12 @@
#!/bin/bash
# install mongosh to load sample data
sudo apt-get update
sudo apt-get install -y wget gnupg
wget -qO - https://www.mongodb.org/static/pgp/server-6.0.asc | sudo apt-key add -
echo "deb [ arch=amd64,arm64 ] https://repo.mongodb.org/apt/ubuntu jammy/mongodb-org/6.0 multiverse" | sudo tee /etc/apt/sources.list.d/mongodb-org-6.0.list
sudo apt-get update
sudo apt-get install -y mongodb-mongosh
# Run mongodb container
docker run -d --name some-mongodb-server -p 27017:27017 mongo

249
toolkits/mongodb/tests/test_write_validation.py Normal file
View file

@ -0,0 +1,249 @@
import pytest
from arcade_mongodb.tools.mongodb import aggregate_documents, count_documents, find_documents
from arcade_tdk import ToolContext, ToolSecretItem
from arcade_tdk.errors import RetryableToolError
from .conftest import TEST_MONGODB_CONNECTION_STRING
@pytest.fixture
def mock_context():
context = ToolContext()
context.secrets = []
context.secrets.append(
ToolSecretItem(key="MONGODB_CONNECTION_STRING", value=TEST_MONGODB_CONNECTION_STRING)
)
return context
@pytest.mark.asyncio
async def test_filter_dict_blocks_set_operation(mock_context) -> None:
"""Test that $set operation in filter_dict is blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"$set": {"status": "modified"}}', # Write operation
limit=1,
)
error_message = str(exc_info.value)
assert "Write operation '$set' not allowed in filter_dict" in error_message
assert "$set" in exc_info.value.developer_message
assert "Only read operations are allowed" in exc_info.value.developer_message
@pytest.mark.asyncio
async def test_filter_dict_blocks_update_operations(mock_context) -> None:
"""Test that various update operations in filter_dict are blocked."""
update_ops = ["$inc", "$unset", "$push", "$pull", "$rename", "$currentDate"]
for op in update_ops:
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict=f'{{"{op}": {{"field": "value"}}}}',
limit=1,
)
error_message = str(exc_info.value)
assert f"Write operation '{op}' not allowed in filter_dict" in error_message
@pytest.mark.asyncio
async def test_projection_blocks_write_operations(mock_context) -> None:
"""Test that write operations in projection are blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
projection='{"$set": {"modified": true}, "name": 1}', # Write operation in projection
limit=1,
)
error_message = str(exc_info.value)
assert "Write operation '$set' not allowed in projection" in error_message
@pytest.mark.asyncio
async def test_sort_blocks_write_operations(mock_context) -> None:
"""Test that write operations in sort are blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
sort=['{"field": "name", "direction": 1, "$inc": {"counter": 1}}'], # Write op in sort
limit=1,
)
error_message = str(exc_info.value)
assert "Write operation '$inc' not allowed in sort[0]" in error_message
@pytest.mark.asyncio
async def test_count_filter_blocks_write_operations(mock_context) -> None:
"""Test that write operations in count filter are blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await count_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active", "$unset": {"password": ""}}', # Write operation
)
error_message = str(exc_info.value)
assert "Write operation '$unset' not allowed in filter_dict" in error_message
@pytest.mark.asyncio
async def test_aggregation_pipeline_blocks_out_stage(mock_context) -> None:
"""Test that $out stage in aggregation pipeline is blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await aggregate_documents(
mock_context,
database_name="test_database",
collection_name="users",
pipeline=[
'{"$match": {"status": "active"}}',
'{"$out": "output_collection"}', # Write stage
],
)
error_message = str(exc_info.value)
assert "Write stage '$out' not allowed in pipeline" in error_message
@pytest.mark.asyncio
async def test_aggregation_pipeline_blocks_merge_stage(mock_context) -> None:
"""Test that $merge stage in aggregation pipeline is blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await aggregate_documents(
mock_context,
database_name="test_database",
collection_name="users",
pipeline=[
'{"$match": {"status": "active"}}',
'{"$merge": {"into": "target_collection"}}', # Write stage
],
)
error_message = str(exc_info.value)
assert "Write stage '$merge' not allowed in pipeline" in error_message
@pytest.mark.asyncio
async def test_where_operator_blocked(mock_context) -> None:
"""Test that $where operator is blocked for security reasons."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"$where": "this.name == \'admin\'"}', # JavaScript execution
limit=1,
)
error_message = str(exc_info.value)
assert "JavaScript execution operator '$where' not allowed in filter_dict" in error_message
assert (
"JavaScript execution is not allowed for security reasons"
in exc_info.value.developer_message
)
@pytest.mark.asyncio
async def test_nested_write_operations_blocked(mock_context) -> None:
"""Test that nested write operations are blocked."""
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": "active", "nested": {"$set": {"field": "value"}}}', # Nested write op
limit=1,
)
error_message = str(exc_info.value)
assert "Write operation '$set' not allowed in filter_dict" in error_message
assert "nested.$set" in exc_info.value.developer_message # Should show the path
@pytest.mark.asyncio
async def test_valid_read_operations_allowed(mock_context) -> None:
"""Test that valid read operations are allowed."""
# These should not raise write operation errors
try:
# Test query operators
result = await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict='{"status": {"$in": ["active", "inactive"]}, "name": {"$regex": "^A"}}',
projection='{"name": 1, "email": 1, "_id": 0}',
sort=['{"field": "name", "direction": 1}'],
limit=1,
)
assert isinstance(result, list)
# Test aggregation pipeline with read-only stages
pipeline_result = await aggregate_documents(
mock_context,
database_name="test_database",
collection_name="users",
pipeline=[
'{"$match": {"status": "active"}}',
'{"$group": {"_id": "$status", "count": {"$sum": 1}}}',
'{"$sort": {"count": -1}}',
],
)
assert isinstance(pipeline_result, list)
except RetryableToolError as e:
# If we get an error, it should not be about write operations
error_message = str(e)
nested_message = str(e.__cause__) if e.__cause__ else ""
assert "Write operation" not in error_message
assert "Write stage" not in error_message
assert "Write operation" not in nested_message
assert "Write stage" not in nested_message
@pytest.mark.asyncio
async def test_array_write_operations_blocked(mock_context) -> None:
"""Test that array write operations are blocked."""
array_write_ops = ["$addToSet", "$pop", "$pull", "$push", "$pullAll"]
for op in array_write_ops:
with pytest.raises(RetryableToolError) as exc_info:
await find_documents(
mock_context,
database_name="test_database",
collection_name="users",
filter_dict=f'{{"{op}": {{"tags": "new_tag"}}}}',
limit=1,
)
error_message = str(exc_info.value)
assert f"Write operation '{op}' not allowed in filter_dict" in error_message
@pytest.mark.asyncio
async def test_aggregation_stage_content_validated(mock_context) -> None:
"""Test that content within aggregation stages is also validated for write operations."""
with pytest.raises(RetryableToolError) as exc_info:
await aggregate_documents(
mock_context,
database_name="test_database",
collection_name="users",
pipeline=[
'{"$match": {"status": "active", "$set": {"modified": true}}}' # Write op inside $match
],
)
error_message = str(exc_info.value)
assert "Write operation '$set' not allowed in pipeline[0].$match" in error_message