Merge pull request #33 from neo4j-contrib/aura-update-request-parsing-logic

add cloud provider and region validation to aura MCP

Author: a-s-g93

servers/mcp-neo4j-cloud-aura-api/src/mcp_neo4j_aura_manager/__init__.py

@@ -3,6 +3,8 @@
 import argparse
 import os
 import logging
+import sys 
+
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

servers/mcp-neo4j-cloud-aura-api/src/mcp_neo4j_aura_manager/server.py

@@ -1,9 +1,5 @@
-import argparse
-import asyncio
 import json
 import logging
-import os
-import sys
 import time
 from typing import Any, Dict, List, Optional, Union
 
@@ -19,6 +15,30 @@
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+
+def _validate_region(cloud_provider: str, region: str) -> None:
+    """
+    Validate the region exists for the given cloud provider.
+
+    Args:
+        cloud_provider: The cloud provider to validate the region for
+        region: The region to validate
+
+    Returns:
+        None
+    
+    Raises:
+        ValueError: If the region is not valid for the given cloud provider
+    """
+
+    if cloud_provider == "gcp" and region.count("-") != 1:
+        raise ValueError(f"Invalid region for GCP: {region}. Must follow the format 'region-zonenumber'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
+    elif cloud_provider == "aws" and region.count("-") != 2:
+        raise ValueError(f"Invalid region for AWS: {region}. Must follow the format 'region-zone-number'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
+    elif cloud_provider == "azure" and region.count("-") != 0:
+        raise ValueError(f"Invalid region for Azure: {region}. Must follow the format 'regionzone'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
+
+    
 class AuraAPIClient:
     """Client for interacting with Neo4j Aura API."""
 
@@ -172,12 +192,17 @@ def create_instance(self, tenant_id: str, name: str, memory: int = 1, region: st
             if cloud_provider and cloud_provider not in ["gcp", "aws", "azure"]:
                 raise ValueError("cloud_provider must be one of: gcp, aws, azure")
 
+            if vector_optimized and memory < 4:
+                raise ValueError("vector optimized instances must have at least 4GB memory")
+            
             # If cloning, source_instance_id is required
             if source_instance_id is not None:
                 if not isinstance(source_instance_id, str):
                     raise ValueError("source_instance for clone from instance must be defined")
         else:
             raise ValueError(f"Invalid type {type}")
+        
+        _validate_region(cloud_provider, region)
 
         url = f"{self.BASE_URL}/instances"
         payload = {
@@ -222,6 +247,9 @@ def update_instance(self, instance_id: str, name: Optional[str] = None,
         if vector_optimized is not None:
             payload["vector_optimized"] = str(vector_optimized).lower()
 
+        if payload["vector_optimized"] == "true" and int(payload["memory"]) < 4:
+            raise ValueError("vector optimized instances must have at least 4GB memory")
+        
         print("Update instance payload:")
         print(payload)
         response = requests.patch(url, headers=self._get_headers(), json=payload)
@@ -485,7 +513,7 @@ async def handle_list_tools() -> List[types.Tool]:
                         },
                         "region": {
                             "type": "string",
-                            "description": "Region for the instance (e.g., 'us-east-1')",
+                            "description": "Region for the instance (e.g., 'us-central1')",
                             "default": "us-central1"
                         },
                         "type": {
@@ -495,7 +523,7 @@ async def handle_list_tools() -> List[types.Tool]:
                         },
                         "vector_optimized": {
                             "type": "boolean",
-                            "description": "Whether the instance is optimized for vector operations",
+                            "description": "Whether the instance is optimized for vector operations. Only allowed for instance with more than 4GB memory.",
                             "default": False
                         },
                         "cloud_provider": {

servers/mcp-neo4j-cloud-aura-api/test.sh

@@ -1,4 +1,5 @@
-uv run pytest tests/test_aura_manager.py
 if [ -f .env ]; then
-    uv run --env-file .env pytest tests/test_aura_integration.py
+    uv run --env-file .env pytest tests
+else
+    uv run pytest tests/test_aura_manager.py
 fi

servers/mcp-neo4j-cloud-aura-api/tests/test_aura_integration.py

@@ -263,3 +263,11 @@ def test_create_and_delete_instance_integration(aura_client, test_type):
         assert "status" in instance_details
         print(f"Deleted test instance {instance_id}: {delete_result} {instance_details}")
 
+
+def test_create_instance_vector_optimized_and_memory_less_than_4_should_raise_error(aura_client):
+    with pytest.raises(ValueError):
+        aura_client.create_instance(memory=3, vector_optimized=True, tenant_id="test-tenant-1", name="Test Instance")
+
+def test_update_instance_vector_optimized_and_memory_less_than_4_should_raise_error(aura_client):
+    with pytest.raises(ValueError):
+        aura_client.update_instance(instance_id="test-instance-1", memory=3, vector_optimized=True)

servers/mcp-neo4j-cloud-aura-api/tests/test_utils.py

@@ -0,0 +1,44 @@
+from mcp_neo4j_aura_manager.server import _validate_region
+import pytest
+
+def test_validate_region_aws_valid():
+    # Test GCP regions
+    assert _validate_region("aws", "us-east-1") is None
+    assert _validate_region("aws", "eu-west-1") is None
+    assert _validate_region("aws", "eu-east-1") is None
+
+def test_validate_region_aws_invalid():
+    # Test GCP regions
+    with pytest.raises(ValueError):
+        _validate_region("aws", "us-east1")
+    with pytest.raises(ValueError):
+        _validate_region("aws", "euwest")
+    with pytest.raises(ValueError):
+        _validate_region("aws", "eu-west-1-1-1")
+
+def test_validate_region_gcp_valid():
+    # Test GCP regions
+    assert _validate_region("gcp", "us-central1") is None
+    assert _validate_region("gcp", "europe-west1") is None
+    assert _validate_region("gcp", "us-central2") is None
+
+def test_validate_region_gcp_invalid():
+    # Test GCP regions
+    with pytest.raises(ValueError):
+        _validate_region("gcp", "us-east-1")
+    with pytest.raises(ValueError):
+        _validate_region("gcp", "eu-west-1-1")
+    with pytest.raises(ValueError):
+        _validate_region("gcp", "euwest")
+
+def test_validate_region_azure_valid():
+    # Test Azure regions
+    assert _validate_region("azure", "eastus") is None
+    assert _validate_region("azure", "northeurope") is None
+
+def test_validate_region_azure_invalid():
+    # Test Azure regions
+    with pytest.raises(ValueError):
+        _validate_region("azure", "us-east-1")
+    with pytest.raises(ValueError):
+        _validate_region("azure", "eu-west1")