Support multiple buckets (#47)

* Add support for the MCP server to work with multiple buckets

- Introduced `connect_to_bucket` utility to establish bucket connections using the cluster context.
- Updated relevant functions to accept `bucket_name` as an argument where necessary, improving flexibility in handling Couchbase operations.
- Enhanced documentation and logging for better clarity on connection status and error handling.

* Update README.md to enhance feature list and clarify environment variable usage

- Added new features for retrieving bucket, scope, and collection lists in the MCP server.
- Removed references to `CB_BUCKET_NAME` from environment variable sections

* Remove `validate_connection_config` function and its references from the utils module

* Add the bucket configuration back in to the MCP server

* Add bucket name retrieval to server configuration and update connection logic

- Included 'bucket_name' in the server configuration status.
- Updated the connection logic to resolve the bucket name before attempting to connect, improving error handling and clarity in connection status.

* Use Query based retrieval for getting the collections in a scope

* Fix typo in docstring of run_cluster_query function in query.py

* Handle review comments

- Typo in docstring.
- Enhanced error handling by removing redundant checks for empty strings.
- Improved the retrieval of bucket names in connection status responses.

* Update README.md to clarify environment variable usage and enhance feature descriptions

- Added note regarding the requirement of Query services for retrieving collections in a specified scope and bucket.
- Included `CB_BUCKET_NAME` in the environment variable section with a description for its usage.

* Update error handling in cluster connection response to always indicate disconnection status

* Update README.md to add a note on database credential requirements for default bucket in MCP server configuration

* Refactor bucket name resolution in KV, Query, and Server modules

- Updated the logic to use `resolved_bucket_name` instead of `bucket_name` for improved clarity and consistency across functions.
- Enhanced error logging in the server module to specify the resolved bucket name when encountering issues.

* Remove bucket_name from MCP server initialization

- Removed the `resolve_bucket_name` function and updated the code to directly use `bucket_name` in KV, Query, and Server modules for improved clarity.
- Adjusted function signatures to require `bucket_name` as a mandatory parameter, enhancing consistency in bucket handling.
- Updated error logging to reflect the direct use of `bucket_name` in connection status responses.

* Update README.md to reflect removal of CB_BUCKET_NAME from configuration examples

- Removed references to `CB_BUCKET_NAME` in environment variable examples and usage notes for clarity.
- Adjusted descriptions to ensure consistency with recent changes in MCP server initialization.

* Handle error in case of unspecified bucket_name for test_cluster_connection tool

* Update README.md to remove note on database credential limitations for MCP server configuration

Author: nithishr

README.md

@@ -2,17 +2,18 @@
 
 An [MCP](https://modelcontextprotocol.io/) server implementation of Couchbase that allows LLMs to directly interact with Couchbase clusters.
 
-[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) [![PyPI version](https://badge.fury.io/py/couchbase-mcp-server.svg)](https://pypi.org/project/couchbase-mcp-server/) [![Verified on MseeP](https://mseep.ai/badge.svg)](https://mseep.ai/app/13fce476-0e74-4b1e-ab82-1df2a3204809)
-
-[![Trust Score](https://archestra.ai/mcp-catalog/api/badge/quality/Couchbase-Ecosystem/mcp-server-couchbase)](https://archestra.ai/mcp-catalog/couchbase-ecosystem__mcp-server-couchbase)
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) [![PyPI version](https://badge.fury.io/py/couchbase-mcp-server.svg)](https://pypi.org/project/couchbase-mcp-server/) [![Verified on MseeP](https://mseep.ai/badge.svg)](https://mseep.ai/app/13fce476-0e74-4b1e-ab82-1df2a3204809) [![Trust Score](https://archestra.ai/mcp-catalog/api/badge/quality/Couchbase-Ecosystem/mcp-server-couchbase)](https://archestra.ai/mcp-catalog/couchbase-ecosystem__mcp-server-couchbase)
 
 <a href="https://glama.ai/mcp/servers/@Couchbase-Ecosystem/mcp-server-couchbase">
   <img width="380" height="200" src="https://glama.ai/mcp/servers/@Couchbase-Ecosystem/mcp-server-couchbase/badge" alt="Couchbase Server MCP server" />
 </a>
 
 ## Features
 
+- Get a list of all the buckets in the cluster
 - Get a list of all the scopes and collections in the specified bucket
+- Get a list of all the scopes in the specified bucket
+- Get a list of all the collections in a specified scope and bucket. Note that this tool requires the cluster to have Query service.
 - Get the structure for a collection
 - Get a document by ID from a specified scope and collection
 - Upsert a document by ID to a specified scope and collection
@@ -48,8 +49,7 @@ We publish a pre built [PyPI package](https://pypi.org/project/couchbase-mcp-ser
       "env": {
         "CB_CONNECTION_STRING": "couchbases://connection-string",
         "CB_USERNAME": "username",
-        "CB_PASSWORD": "password",
-        "CB_BUCKET_NAME": "bucket_name"
+        "CB_PASSWORD": "password"
       }
     }
   }
@@ -86,8 +86,7 @@ This is the common configuration for the MCP clients such as Claude Desktop, Cur
       "env": {
         "CB_CONNECTION_STRING": "couchbases://connection-string",
         "CB_USERNAME": "username",
-        "CB_PASSWORD": "password",
-        "CB_BUCKET_NAME": "bucket_name"
+        "CB_PASSWORD": "password"
       }
     }
   }
@@ -105,9 +104,8 @@ The server can be configured using environment variables or command line argumen
 | Environment Variable          | CLI Argument             | Description                                | Default      |
 | ----------------------------- | ------------------------ | ------------------------------------------ | ------------ |
 | `CB_CONNECTION_STRING`        | `--connection-string`    | Connection string to the Couchbase cluster | **Required** |
-| `CB_USERNAME`                 | `--username`             | Username with bucket access                | **Required** |
+| `CB_USERNAME`                 | `--username`             | Username with access to required buckets   | **Required** |
 | `CB_PASSWORD`                 | `--password`             | Password for authentication                | **Required** |
-| `CB_BUCKET_NAME`              | `--bucket-name`          | Name of the bucket to access               | **Required** |
 | `CB_MCP_READ_ONLY_QUERY_MODE` | `--read-only-query-mode` | Prevent data modification queries          | `true`       |
 | `CB_MCP_TRANSPORT`            | `--transport`            | Transport mode: `stdio`, `http`, `sse`     | `stdio`      |
 | `CB_MCP_HOST`                 | `--host`                 | Host for HTTP/SSE transport modes          | `127.0.0.1`  |
@@ -210,7 +208,6 @@ uvx couchbase-mcp-server \
   --connection-string='<couchbase_connection_string>' \
   --username='<database_username>' \
   --password='<database_password>' \
-  --bucket-name='<couchbase_bucket_to_use>' \
   --read-only-query-mode=true \
   --transport=http
 ```
@@ -244,7 +241,6 @@ uvx couchbase-mcp-server \
   --connection-string='<couchbase_connection_string>' \
   --username='<database_username>' \
   --password='<database_password>' \
-  --bucket-name='<couchbase_bucket_to_use>' \
   --read-only-query-mode=true \
   --transport=sse
 ```
@@ -321,7 +317,6 @@ docker run --rm -i \
   -e CB_CONNECTION_STRING='<couchbase_connection_string>' \
   -e CB_USERNAME='<database_user>' \
   -e CB_PASSWORD='<database_password>' \
-  -e CB_BUCKET_NAME='<bucket_name>' \
   -e CB_MCP_TRANSPORT='<http|sse|stdio>' \
   -e CB_MCP_READ_ONLY_QUERY_MODE='<true|false>' \
   -e CB_MCP_PORT=9001 \
@@ -350,8 +345,6 @@ The Docker image can be used in `stdio` transport mode with the following config
         "CB_USERNAME=<database_user>",
         "-e",
         "CB_PASSWORD=<database_password>",
-        "-e",
-        "CB_BUCKET_NAME=<bucket_name>",
         "mcp/couchbase"
       ]
     }
@@ -377,7 +370,7 @@ The Couchbase MCP server can also be used as a managed server in your agentic ap
 ## Troubleshooting Tips
 
 - Ensure the path to your MCP server repository is correct in the configuration if running from source.
-- Verify that your Couchbase connection string, database username, password and bucket name are correct.
+- Verify that your Couchbase connection string, database username, password are correct.
 - If using Couchbase Capella, ensure that the cluster is [accessible](https://docs.couchbase.com/cloud/clusters/allow-ip-address.html) from the machine where the MCP server is running.
 - Check that the database user has proper permissions to access the specified bucket.
 - Confirm that the `uv` package manager is properly installed and accessible. You may need to provide absolute path to `uv`/`uvx` in the `command` field in the configuration.

src/mcp_server.py

@@ -77,11 +77,6 @@ async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
     envvar="CB_PASSWORD",
     help="Couchbase database password (required for operations)",
 )
-@click.option(
-    "--bucket-name",
-    envvar="CB_BUCKET_NAME",
-    help="Couchbase bucket name (required for operations)",
-)
 @click.option(
     "--read-only-query-mode",
     envvar=[
@@ -121,7 +116,6 @@ def main(
     connection_string,
     username,
     password,
-    bucket_name,
     read_only_query_mode,
     transport,
     host,
@@ -133,7 +127,6 @@ def main(
         "connection_string": connection_string,
         "username": username,
         "password": password,
-        "bucket_name": bucket_name,
         "read_only_query_mode": read_only_query_mode,
         "transport": transport,
         "host": host,

src/tools/__init__.py

@@ -19,16 +19,22 @@
 
 # Server tools
 from .server import (
+    get_buckets_in_cluster,
+    get_collections_in_scope,
     get_scopes_and_collections_in_bucket,
+    get_scopes_in_bucket,
     get_server_configuration_status,
     test_cluster_connection,
 )
 
 # List of all tools for easy registration
 ALL_TOOLS = [
+    get_buckets_in_cluster,
     get_server_configuration_status,
     test_cluster_connection,
     get_scopes_and_collections_in_bucket,
+    get_collections_in_scope,
+    get_scopes_in_bucket,
     get_document_by_id,
     upsert_document_by_id,
     delete_document_by_id,
@@ -41,6 +47,9 @@
     "get_server_configuration_status",
     "test_cluster_connection",
     "get_scopes_and_collections_in_bucket",
+    "get_collections_in_scope",
+    "get_scopes_in_bucket",
+    "get_buckets_in_cluster",
     "get_document_by_id",
     "upsert_document_by_id",
     "delete_document_by_id",

src/tools/kv.py

@@ -9,18 +9,25 @@
 
 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.kv")
 
 
 def get_document_by_id(
-    ctx: Context, scope_name: str, collection_name: str, document_id: str
+    ctx: Context,
+    bucket_name: str,
+    scope_name: str,
+    collection_name: str,
+    document_id: str,
 ) -> dict[str, Any]:
     """Get a document by its ID from the specified scope and collection.
     If the document is not found, it will raise an exception."""
-    bucket = ensure_bucket_connection(ctx)
+
+    cluster = get_cluster_connection(ctx)
+    bucket = connect_to_bucket(cluster, bucket_name)
     try:
         collection = bucket.scope(scope_name).collection(collection_name)
         result = collection.get(document_id)
@@ -32,14 +39,16 @@ def get_document_by_id(
 
 def upsert_document_by_id(
     ctx: Context,
+    bucket_name: str,
     scope_name: str,
     collection_name: str,
     document_id: str,
     document_content: dict[str, Any],
 ) -> bool:
     """Insert or update a document by its ID.
     Returns True on success, False on failure."""
-    bucket = ensure_bucket_connection(ctx)
+    cluster = get_cluster_connection(ctx)
+    bucket = connect_to_bucket(cluster, bucket_name)
     try:
         collection = bucket.scope(scope_name).collection(collection_name)
         collection.upsert(document_id, document_content)
@@ -51,11 +60,16 @@ def upsert_document_by_id(
 
 
 def delete_document_by_id(
-    ctx: Context, scope_name: str, collection_name: str, document_id: str
+    ctx: Context,
+    bucket_name: str,
+    scope_name: str,
+    collection_name: str,
+    document_id: str,
 ) -> bool:
     """Delete a document by its ID.
     Returns True on success, False on failure."""
-    bucket = ensure_bucket_connection(ctx)
+    cluster = get_cluster_connection(ctx)
+    bucket = connect_to_bucket(cluster, bucket_name)
     try:
         collection = bucket.scope(scope_name).collection(collection_name)
         collection.remove(document_id)

src/tools/query.py

@@ -10,22 +10,23 @@
 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