diff --git a/.github/workflows/test-toolkits.yml b/.github/workflows/test-toolkits.yml index ed05bf57..ca514166 100644 --- a/.github/workflows/test-toolkits.yml +++ b/.github/workflows/test-toolkits.yml @@ -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 diff --git a/toolkits/mongodb/Makefile b/toolkits/mongodb/Makefile new file mode 100644 index 00000000..7e2c686e --- /dev/null +++ b/toolkits/mongodb/Makefile @@ -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 diff --git a/toolkits/mongodb/arcade_mongodb/__init__.py b/toolkits/mongodb/arcade_mongodb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/toolkits/mongodb/arcade_mongodb/database_engine.py b/toolkits/mongodb/arcade_mongodb/database_engine.py new file mode 100644 index 00000000..744890de --- /dev/null +++ b/toolkits/mongodb/arcade_mongodb/database_engine.py @@ -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 diff --git a/toolkits/mongodb/arcade_mongodb/tools/__init__.py b/toolkits/mongodb/arcade_mongodb/tools/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/toolkits/mongodb/arcade_mongodb/tools/mongodb.py b/toolkits/mongodb/arcade_mongodb/tools/mongodb.py new file mode 100644 index 00000000..68176862 --- /dev/null +++ b/toolkits/mongodb/arcade_mongodb/tools/mongodb.py @@ -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 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 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 or use the 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 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 diff --git a/toolkits/mongodb/arcade_mongodb/tools/utils.py b/toolkits/mongodb/arcade_mongodb/tools/utils.py new file mode 100644 index 00000000..ada7f2a7 --- /dev/null +++ b/toolkits/mongodb/arcade_mongodb/tools/utils.py @@ -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 + "$", + "$[]", + "$[]", + # 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 diff --git a/toolkits/mongodb/evals/eval_mongodb.py b/toolkits/mongodb/evals/eval_mongodb.py new file mode 100644 index 00000000..e1e27511 --- /dev/null +++ b/toolkits/mongodb/evals/eval_mongodb.py @@ -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 diff --git a/toolkits/mongodb/pyproject.toml b/toolkits/mongodb/pyproject.toml new file mode 100644 index 00000000..e2127c6f --- /dev/null +++ b/toolkits/mongodb/pyproject.toml @@ -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",] diff --git a/toolkits/mongodb/tests/__init__.py b/toolkits/mongodb/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/toolkits/mongodb/tests/conftest.py b/toolkits/mongodb/tests/conftest.py new file mode 100644 index 00000000..d635dfca --- /dev/null +++ b/toolkits/mongodb/tests/conftest.py @@ -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() diff --git a/toolkits/mongodb/tests/dump.js b/toolkits/mongodb/tests/dump.js new file mode 100644 index 00000000..b878bf30 --- /dev/null +++ b/toolkits/mongodb/tests/dump.js @@ -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()); diff --git a/toolkits/mongodb/tests/test_json_validation.py b/toolkits/mongodb/tests/test_json_validation.py new file mode 100644 index 00000000..38c2a476 --- /dev/null +++ b/toolkits/mongodb/tests/test_json_validation.py @@ -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 diff --git a/toolkits/mongodb/tests/test_mongodb.py b/toolkits/mongodb/tests/test_mongodb.py new file mode 100644 index 00000000..e70c7fa8 --- /dev/null +++ b/toolkits/mongodb/tests/test_mongodb.py @@ -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 diff --git a/toolkits/mongodb/tests/test_setup.sh b/toolkits/mongodb/tests/test_setup.sh new file mode 100755 index 00000000..23b155df --- /dev/null +++ b/toolkits/mongodb/tests/test_setup.sh @@ -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 diff --git a/toolkits/mongodb/tests/test_write_validation.py b/toolkits/mongodb/tests/test_write_validation.py new file mode 100644 index 00000000..8830dfd1 --- /dev/null +++ b/toolkits/mongodb/tests/test_write_validation.py @@ -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