couchbase-mcp-server 0.4.0rc1__py3-none-any.whl → 0.5.1__py3-none-any.whl

tools/query.py

@@ -10,22 +10,23 @@ from typing import Any
 from lark_sqlpp import modifies_data, modifies_structure, parse_sqlpp
 from mcp.server.fastmcp import Context
 
+from utils.connection import connect_to_bucket
 from utils.constants import MCP_SERVER_NAME
-from utils.context import ensure_bucket_connection
+from utils.context import get_cluster_connection
 
 logger = logging.getLogger(f"{MCP_SERVER_NAME}.tools.query")
 
 
 def get_schema_for_collection(
-    ctx: Context, scope_name: str, collection_name: str
+    ctx: Context, bucket_name: str, scope_name: str, collection_name: str
 ) -> dict[str, Any]:
     """Get the schema for a collection in the specified scope.
     Returns a dictionary with the collection name and the schema returned by running INFER query on the Couchbase collection.
     """
     schema = {"collection_name": collection_name, "schema": []}
     try:
-        query = f"INFER {collection_name}"
-        result = run_sql_plus_plus_query(ctx, scope_name, query)
+        query = f"INFER `{collection_name}`"
+        result = run_sql_plus_plus_query(ctx, bucket_name, scope_name, query)
         # Result is a list of list of schemas. We convert it to a list of schemas.
         if result:
             schema["schema"] = result[0]
@@ -36,10 +37,13 @@ def get_schema_for_collection(
 
 
 def run_sql_plus_plus_query(
-    ctx: Context, scope_name: str, query: str
+    ctx: Context, bucket_name: str, scope_name: str, query: str
 ) -> list[dict[str, Any]]:
     """Run a SQL++ query on a scope and return the results as a list of JSON objects."""
-    bucket = ensure_bucket_connection(ctx)
+    cluster = get_cluster_connection(ctx)
+
+    bucket = connect_to_bucket(cluster, bucket_name)
+
     app_context = ctx.request_context.lifespan_context
     read_only_query_mode = app_context.read_only_query_mode
     logger.info(f"Running SQL++ queries in read-only mode: {read_only_query_mode}")
@@ -75,3 +79,20 @@ def run_sql_plus_plus_query(
     except Exception as e:
         logger.error(f"Error running query: {e!s}", exc_info=True)
         raise
+
+
+# Don't expose this function to the MCP server until we have a use case
+def run_cluster_query(ctx: Context, query: str, **kwargs: Any) -> list[dict[str, Any]]:
+    """Run a query on the cluster object and return the results as a list of JSON objects."""
+
+    cluster = get_cluster_connection(ctx)
+    results = []
+
+    try:
+        result = cluster.query(query, **kwargs)
+        for row in result:
+            results.append(row)
+        return results
+    except Exception as e:
+        logger.error(f"Error running query: {e}")
+        raise

tools/server.py

@@ -1,17 +1,20 @@
 """
 Tools for server operations.
 
-This module contains tools for getting the server status, testing the connection, and getting the scopes and collections in the bucket.
+This module contains tools for getting the server status, testing the connection, and getting the buckets in the cluster, the scopes and collections in the bucket.
 """
 
+import json
 import logging
 from typing import Any
 
 from mcp.server.fastmcp import Context
 
+from tools.query import run_cluster_query
 from utils.config import get_settings
+from utils.connection import connect_to_bucket
 from utils.constants import MCP_SERVER_NAME
-from utils.context import ensure_bucket_connection
+from utils.context import get_cluster_connection
 
 logger = logging.getLogger(f"{MCP_SERVER_NAME}.tools.server")
 
@@ -26,15 +29,16 @@ def get_server_configuration_status(ctx: Context) -> dict[str, Any]:
     configuration = {
         "connection_string": settings.get("connection_string", "Not set"),
         "username": settings.get("username", "Not set"),
-        "bucket_name": settings.get("bucket_name", "Not set"),
         "read_only_query_mode": settings.get("read_only_query_mode", True),
         "password_configured": bool(settings.get("password")),
+        "ca_cert_path_configured": bool(settings.get("ca_cert_path")),
+        "client_cert_path_configured": bool(settings.get("client_cert_path")),
+        "client_key_path_configured": bool(settings.get("client_key_path")),
     }
 
     app_context = ctx.request_context.lifespan_context
     connection_status = {
         "cluster_connected": app_context.cluster is not None,
-        "bucket_connected": app_context.bucket is not None,
     }
 
     return {
@@ -45,39 +49,46 @@ def get_server_configuration_status(ctx: Context) -> dict[str, Any]:
     }
 
 
-def test_cluster_connection(ctx: Context) -> dict[str, Any]:
-    """Test the connection to Couchbase cluster and bucket.
+def test_cluster_connection(
+    ctx: Context, bucket_name: str | None = None
+) -> dict[str, Any]:
+    """Test the connection to Couchbase cluster and optionally to a bucket.
     This tool verifies the connection to the Couchbase cluster and bucket by establishing the connection if it is not already established.
+    If bucket name is not provided, it will not try to connect to the bucket specified in the MCP server settings.
     Returns connection status and basic cluster information.
     """
     try:
-        bucket = ensure_bucket_connection(ctx)
-
-        # Test basic connectivity by getting bucket name
-        bucket_name = bucket.name
+        cluster = get_cluster_connection(ctx)
+        bucket = None
+        if bucket_name:
+            bucket = connect_to_bucket(cluster, bucket_name)
 
         return {
             "status": "success",
-            "cluster_connected": True,
-            "bucket_connected": True,
+            "cluster_connected": cluster.connected,
+            "bucket_connected": bucket is not None,
             "bucket_name": bucket_name,
-            "message": "Successfully connected to Couchbase cluster and bucket",
+            "message": "Successfully connected to Couchbase cluster",
         }
     except Exception as e:
         return {
             "status": "error",
             "cluster_connected": False,
             "bucket_connected": False,
+            "bucket_name": bucket_name,
             "error": str(e),
-            "message": "Failed to connect to Couchbase",
+            "message": "Failed to connect to Couchbase cluster",
         }
 
 
-def get_scopes_and_collections_in_bucket(ctx: Context) -> dict[str, list[str]]:
+def get_scopes_and_collections_in_bucket(
+    ctx: Context, bucket_name: str
+) -> dict[str, list[str]]:
     """Get the names of all scopes and collections in the bucket.
     Returns a dictionary with scope names as keys and lists of collection names as values.
     """
-    bucket = ensure_bucket_connection(ctx)
+    cluster = get_cluster_connection(ctx)
+    bucket = connect_to_bucket(cluster, bucket_name)
     try:
         scopes_collections = {}
         collection_manager = bucket.collections()
@@ -89,3 +100,81 @@ def get_scopes_and_collections_in_bucket(ctx: Context) -> dict[str, list[str]]:
     except Exception as e:
         logger.error(f"Error getting scopes and collections: {e}")
         raise
+
+
+def get_buckets_in_cluster(ctx: Context) -> list[str]:
+    """Get the names of all the accessible buckets in the cluster."""
+    cluster = get_cluster_connection(ctx)
+    bucket_manager = cluster.buckets()
+    buckets_with_settings = bucket_manager.get_all_buckets()
+
+    buckets = []
+    for bucket in buckets_with_settings:
+        buckets.append(bucket.name)
+
+    return buckets
+
+
+def get_scopes_in_bucket(ctx: Context, bucket_name: str) -> list[str]:
+    """Get the names of all scopes in the given bucket."""
+    cluster = get_cluster_connection(ctx)
+    bucket = connect_to_bucket(cluster, bucket_name)
+    try:
+        scopes = bucket.collections().get_all_scopes()
+        return [scope.name for scope in scopes]
+    except Exception as e:
+        logger.error(f"Error getting scopes in the bucket {bucket_name}: {e}")
+        raise
+
+
+def get_collections_in_scope(
+    ctx: Context, bucket_name: str, scope_name: str
+) -> list[str]:
+    """Get the names of all collections in the given scope and bucket."""
+
+    # Get the collections in the scope using system:all_keyspaces collection
+    query = "SELECT DISTINCT(name) as collection_name FROM system:all_keyspaces where `bucket`=$bucket_name and `scope`=$scope_name"
+    results = run_cluster_query(
+        ctx, query, bucket_name=bucket_name, scope_name=scope_name
+    )
+    return [result["collection_name"] for result in results]
+
+
+def get_cluster_health_and_services(
+    ctx: Context, bucket_name: str | None = None
+) -> dict[str, Any]:
+    """Get cluster health status and list of all running services.
+
+    This tool provides health monitoring by:
+    - Getting health status of all running services with latency information (via ping)
+    - Listing all services running on the cluster with their endpoints
+    - Showing connection status and node information for each service
+
+    If bucket_name is provided, it actively pings services from the perspective of the bucket.
+    Otherwise, it uses cluster-level ping to get the health status of the cluster.
+
+    Returns:
+    - Cluster health status with service-level connection details and latency measurements
+    """
+    try:
+        cluster = get_cluster_connection(ctx)
+
+        if bucket_name:
+            # Ping services from the perspective of the bucket
+            bucket = connect_to_bucket(cluster, bucket_name)
+            result = bucket.ping().as_json()
+        else:
+            # Ping services from the perspective of the cluster
+            result = cluster.ping().as_json()
+
+        return {
+            "status": "success",
+            "data": json.loads(result),
+        }
+    except Exception as e:
+        logger.error(f"Error getting cluster health: {e}")
+        return {
+            "status": "error",
+            "error": str(e),
+            "message": "Failed to get cluster health and services information",
+        }

utils/__init__.py

@@ -7,8 +7,6 @@ This module contains utility functions for configuration, connection, and contex
 # Configuration utilities
 from .config import (
     get_settings,
-    validate_connection_config,
-    validate_required_param,
 )
 
 # Connection utilities
@@ -33,7 +31,12 @@ from .constants import (
 # Context utilities
 from .context import (
     AppContext,
-    ensure_bucket_connection,
+    get_cluster_connection,
+)
+
+# Index utilities
+from .index_utils import (
+    fetch_indexes_from_rest_api,
 )
 
 # Note: Individual modules create their own hierarchical loggers using:
@@ -42,14 +45,14 @@ from .context import (
 __all__ = [
     # Config
     "get_settings",
-    "validate_required_param",
-    "validate_connection_config",
     # Connection
     "connect_to_couchbase_cluster",
     "connect_to_bucket",
     # Context
     "AppContext",
-    "ensure_bucket_connection",
+    "get_cluster_connection",
+    # Index utilities
+    "fetch_indexes_from_rest_api",
     # Constants
     "MCP_SERVER_NAME",
     "DEFAULT_READ_ONLY_MODE",

utils/config.py

@@ -7,32 +7,7 @@ from .constants import MCP_SERVER_NAME
 logger = logging.getLogger(f"{MCP_SERVER_NAME}.utils.config")
 
 
-def validate_required_param(
-    ctx: click.Context, param: click.Parameter, value: str | None
-) -> str:
-    """Validate that a required parameter is not empty."""
-    if not value or value.strip() == "":
-        raise click.BadParameter(f"{param.name} cannot be empty")
-    return value
-
-
 def get_settings() -> dict:
     """Get settings from Click context."""
     ctx = click.get_current_context()
     return ctx.obj or {}
-
-
-def validate_connection_config() -> None:
-    """Validate that all required parameters for the MCP server are available when needed."""
-    settings = get_settings()
-    required_params = ["connection_string", "username", "password", "bucket_name"]
-    missing_params = []
-
-    for param in required_params:
-        if not settings.get(param):
-            missing_params.append(param)
-
-    if missing_params:
-        error_msg = f"Missing required parameters for the MCP server: {', '.join(missing_params)}"
-        logger.error(error_msg)
-        raise ValueError(error_msg)

utils/connection.py

@@ -1,7 +1,8 @@
 import logging
+import os
 from datetime import timedelta
 
-from couchbase.auth import PasswordAuthenticator
+from couchbase.auth import CertificateAuthenticator, PasswordAuthenticator
 from couchbase.cluster import Bucket, Cluster
 from couchbase.options import ClusterOptions
 
@@ -11,15 +12,40 @@ logger = logging.getLogger(f"{MCP_SERVER_NAME}.utils.connection")
 
 
 def connect_to_couchbase_cluster(
-    connection_string: str, username: str, password: str
+    connection_string: str,
+    username: str,
+    password: str,
+    ca_cert_path: str | None = None,
+    client_cert_path: str | None = None,
+    client_key_path: str | None = None,
 ) -> Cluster:
     """Connect to Couchbase cluster and return the cluster object if successful.
+    The connection can be established using the client certificate and key or the username and password. Optionally, the CA root certificate path can also be provided.
+    Either of the path to the client certificate and key or the username and password should be provided.
+    If the client certificate and key are provided, the username and password are not used.
+    If both the client certificate and key and the username and password are provided, the client certificate is used for authentication.
     If the connection fails, it will raise an exception.
     """
 
     try:
         logger.info("Connecting to Couchbase cluster...")
-        auth = PasswordAuthenticator(username, password)
+        if client_cert_path and client_key_path:
+            logger.info("Connecting to Couchbase cluster with client certificate...")
+            if not os.path.exists(client_cert_path) or not os.path.exists(
+                client_key_path
+            ):
+                raise FileNotFoundError(
+                    f"Client certificate files not found at {os.path.basename(client_cert_path)} or {os.path.basename(client_key_path)}."
+                )
+
+            auth = CertificateAuthenticator(
+                cert_path=client_cert_path,
+                key_path=client_key_path,
+                trust_store_path=ca_cert_path,
+            )
+        else:
+            logger.info("Connecting to Couchbase cluster with password...")
+            auth = PasswordAuthenticator(username, password, cert_path=ca_cert_path)
         options = ClusterOptions(auth)
         options.apply_profile("wan_development")
 
@@ -38,7 +64,6 @@ def connect_to_bucket(cluster: Cluster, bucket_name: str) -> Bucket:
     If the operation fails, it will raise an exception.
     """
     try:
-        logger.info(f"Connecting to bucket: {bucket_name}")
         bucket = cluster.bucket(bucket_name)
         return bucket
     except Exception as e:

utils/context.py

@@ -1,11 +1,11 @@
 import logging
 from dataclasses import dataclass
 
-from couchbase.cluster import Bucket, Cluster
+from couchbase.cluster import Cluster
 from mcp.server.fastmcp import Context
 
-from utils.config import get_settings, validate_connection_config
-from utils.connection import connect_to_bucket, connect_to_couchbase_cluster
+from utils.config import get_settings
+from utils.connection import connect_to_couchbase_cluster
 from utils.constants import MCP_SERVER_NAME
 
 logger = logging.getLogger(f"{MCP_SERVER_NAME}.utils.context")
@@ -16,7 +16,6 @@ class AppContext:
     """Context for the MCP server."""
 
     cluster: Cluster | None = None
-    bucket: Bucket | None = None
     read_only_query_mode: bool = True
 
 
@@ -30,51 +29,36 @@ def _set_cluster_in_lifespan_context(ctx: Context) -> None:
         connection_string = settings.get("connection_string")
         username = settings.get("username")
         password = settings.get("password")
+        ca_cert_path = settings.get("ca_cert_path")
+        client_cert_path = settings.get("client_cert_path")
+        client_key_path = settings.get("client_key_path")
+
         cluster = connect_to_couchbase_cluster(
             connection_string,  # type: ignore
             username,  # type: ignore
             password,  # type: ignore
+            ca_cert_path,
+            client_cert_path,
+            client_key_path,
         )
         ctx.request_context.lifespan_context.cluster = cluster
     except Exception as e:
         logger.error(
-            f"Failed to connect to Couchbase: {e} \n Please check your connection string, username and password"
+            "Failed to connect to Couchbase: %s\n"
+            "Verify connection string, and either:\n"
+            "- Username/password are correct, or\n"
+            "- Client certificate and key exist and match server mapping.\n"
+            "If using self-signed or custom CA, set CB_CA_CERT_PATH to the CA file.",
+            e,
         )
         raise
 
 
-def _set_bucket_in_lifespan_context(ctx: Context) -> None:
-    """Set the bucket in the lifespan context.
-    If the bucket is not set, it will try to connect to the bucket using the cluster object in the lifespan context.
+def get_cluster_connection(ctx: Context) -> Cluster:
+    """Return the cluster connection from the lifespan context.
     If the cluster is not set, it will try to connect to the cluster using the connection string, username, and password.
-    If the connection fails, it will raise an exception.
     """
-    settings = get_settings()
-    bucket_name = settings.get("bucket_name")
-
-    # If the bucket is not set, try to connect to the bucket using the cluster object in the lifespan context
-    app_context = ctx.request_context.lifespan_context
-
-    try:
-        # If the cluster is not set, try to connect to the cluster
-        if not app_context.cluster:
-            _set_cluster_in_lifespan_context(ctx)
-        cluster = app_context.cluster
-
-        # Try to connect to the bucket using the cluster object
-        bucket = connect_to_bucket(cluster, bucket_name)  # type: ignore
-        app_context.bucket = bucket
-    except Exception as e:
-        logger.error(
-            f"Failed to connect to bucket: {e} \n Please check your bucket name and credentials."
-        )
-        raise
-
-
-def ensure_bucket_connection(ctx: Context) -> Bucket:
-    """Ensure bucket connection is established and return the bucket object."""
-    validate_connection_config()
     app_context = ctx.request_context.lifespan_context
-    if not app_context.bucket:
-        _set_bucket_in_lifespan_context(ctx)
-    return app_context.bucket
+    if not app_context.cluster:
+        _set_cluster_in_lifespan_context(ctx)
+    return app_context.cluster