From f659939149f60328a3192164e1a9ec44193bc3fe Mon Sep 17 00:00:00 2001 From: AayushTyagi1 Date: Wed, 15 Oct 2025 01:28:33 +0530 Subject: [PATCH 1/5] feat: DA-1170 add index advisor support with two new tools --- src/tools/__init__.py | 10 +++ src/tools/index.py | 179 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 src/tools/index.py diff --git a/src/tools/__init__.py b/src/tools/__init__.py index 8ac42a5..83c36f5 100644 --- a/src/tools/__init__.py +++ b/src/tools/__init__.py @@ -4,6 +4,12 @@ This module contains all the MCP tools for Couchbase operations. """ +# Index tools +from .index import ( + create_index_from_recommendation, + get_index_advisor_recommendations, +) + # Key-Value tools from .kv import ( delete_document_by_id, @@ -40,6 +46,8 @@ delete_document_by_id, get_schema_for_collection, run_sql_plus_plus_query, + get_index_advisor_recommendations, + create_index_from_recommendation, ] __all__ = [ @@ -55,6 +63,8 @@ "delete_document_by_id", "get_schema_for_collection", "run_sql_plus_plus_query", + "get_index_advisor_recommendations", + "create_index_from_recommendation", # Convenience "ALL_TOOLS", ] diff --git a/src/tools/index.py b/src/tools/index.py new file mode 100644 index 0000000..75c63f7 --- /dev/null +++ b/src/tools/index.py @@ -0,0 +1,179 @@ +""" +Tools for index operations and optimization. + +This module contains tools for getting index recommendations using the Couchbase Index Advisor +and creating indexes based on those recommendations. +""" + +import logging +from typing import Any + +from mcp.server.fastmcp import Context + +from utils.constants import MCP_SERVER_NAME +from utils.context import get_cluster_connection + +logger = logging.getLogger(f"{MCP_SERVER_NAME}.tools.index") + + +def get_index_advisor_recommendations( + ctx: Context, bucket_name: str, scope_name: str, query: str +) -> dict[str, Any]: + """Get index recommendations from Couchbase Index Advisor for a given SQL++ query. + + The Index Advisor analyzes the query and provides recommendations for optimal indexes. + This tool works with SELECT, UPDATE, DELETE, or MERGE queries. + + Returns a dictionary with: + - current_used_indexes: Array of currently used indexes (if any) + - recommended_indexes: Array of recommended secondary indexes (if any) + - recommended_covering_indexes: Array of recommended covering indexes (if any) + + Each index object contains: + - index: The CREATE INDEX SQL++ command + - statements: Array of statement objects with the query and run count + """ + cluster = get_cluster_connection(ctx) + + try: + # Escape single quotes in the query by doubling them for SQL++ string literal + escaped_query = query.replace("'", "''") + + # Build the ADVISOR query with fully qualified keyspace + advisor_query = f"SELECT ADVISOR('{escaped_query}') AS advisor_result" + + logger.info(f"Running Index Advisor for query in {bucket_name}.{scope_name}") + + # Execute the ADVISOR function at cluster level + result = cluster.query(advisor_query) + + # Extract the advisor result from the query response + advisor_results = [] + for row in result: + advisor_results.append(row) + + if not advisor_results: + return { + "message": "No recommendations available", + "current_used_indexes": [], + "recommended_indexes": [], + "recommended_covering_indexes": [], + } + + # The result is wrapped in advisor_result key + advisor_data = advisor_results[0].get("advisor_result", {}) + + # Extract the relevant fields with defaults + response = { + "current_used_indexes": advisor_data.get("current_used_indexes", []), + "recommended_indexes": advisor_data.get("recommended_indexes", []), + "recommended_covering_indexes": advisor_data.get( + "recommended_covering_indexes", [] + ), + } + + # Add summary information for better user experience + response["summary"] = { + "current_indexes_count": len(response["current_used_indexes"]), + "recommended_indexes_count": len(response["recommended_indexes"]), + "recommended_covering_indexes_count": len( + response["recommended_covering_indexes"] + ), + "has_recommendations": bool( + response["recommended_indexes"] + or response["recommended_covering_indexes"] + ), + } + + logger.info( + f"Index Advisor completed. Found {response['summary']['recommended_indexes_count']} recommended indexes" + ) + + return response + + except Exception as e: + logger.error(f"Error running Index Advisor: {e!s}", exc_info=True) + raise + + +def create_index_from_recommendation( + ctx: Context, index_definition: str +) -> dict[str, Any]: + """Create an index using the provided CREATE INDEX statement. + + This tool executes a CREATE INDEX statement, typically from Index Advisor recommendations. + Note: This operation requires write permissions and will fail if: + - The read-only query mode is enabled + - The user lacks CREATE INDEX permissions + - An index with the same name already exists + + The index_definition should be a complete CREATE INDEX statement, for example: + CREATE INDEX adv_city_activity ON `travel-sample`.inventory.landmark(city, activity) + + Returns a dictionary with: + - status: 'success' or 'error' + - message: Description of the result + - index_definition: The index statement that was executed (on success) + """ + cluster = get_cluster_connection(ctx) + + app_context = ctx.request_context.lifespan_context + read_only_query_mode = app_context.read_only_query_mode + + # Check if read-only mode is enabled + if read_only_query_mode: + logger.error("Cannot create index in read-only query mode") + return { + "status": "error", + "message": "Index creation is not allowed in read-only query mode. Please disable read-only mode (CB_MCP_READ_ONLY_QUERY_MODE=false) to create indexes.", + "index_definition": index_definition, + } + + try: + # Validate that the statement is a CREATE INDEX statement + if not index_definition.strip().upper().startswith("CREATE INDEX"): + logger.error("Invalid index definition: must start with CREATE INDEX") + return { + "status": "error", + "message": "Invalid index definition. The statement must be a CREATE INDEX command.", + "index_definition": index_definition, + } + + logger.info(f"Creating index with definition: {index_definition}") + + # Execute the CREATE INDEX statement at cluster level + result = cluster.query(index_definition) + + # Consume the result to ensure the query completes + for _ in result: + pass + + logger.info("Index created successfully") + + return { + "status": "success", + "message": "Index created successfully", + "index_definition": index_definition, + } + + except Exception as e: + error_message = str(e) + logger.error(f"Error creating index: {error_message}", exc_info=True) + + # Provide helpful error messages for common issues + if "already exists" in error_message.lower(): + message = "An index with this name already exists. Consider using a different name or dropping the existing index first." + elif ( + "permission" in error_message.lower() + or "authorized" in error_message.lower() + ): + message = "Insufficient permissions to create index. Please ensure your user has the required permissions." + else: + message = f"Failed to create index: {error_message}" + + return { + "status": "error", + "message": message, + "index_definition": index_definition, + "error_details": error_message, + } From 805f79221729b5f7df7a7139e88800bab62cf698 Mon Sep 17 00:00:00 2001 From: AayushTyagi1 Date: Thu, 16 Oct 2025 10:58:14 +0530 Subject: [PATCH 2/5] test: DA-1170 Added Unit Tests for Index Advisor --- tests/test_index_advisor.py | 200 ++++++++++++++++++++++++++++++++++++ 1 file changed, 200 insertions(+) create mode 100755 tests/test_index_advisor.py diff --git a/tests/test_index_advisor.py b/tests/test_index_advisor.py new file mode 100755 index 0000000..02cd303 --- /dev/null +++ b/tests/test_index_advisor.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +""" +Test script for Index Advisor functionality. + +This script demonstrates how to use the new Index Advisor tools: +1. get_index_advisor_recommendations - Get index recommendations for a query +2. create_index_from_recommendation - Create an index from recommendations + +Usage: + uv run test_index_advisor.py + +Prerequisites: + - Set environment variables: CB_CONNECTION_STRING, CB_USERNAME, CB_PASSWORD + - Have a Couchbase cluster accessible with the travel-sample bucket (or modify the script) +""" + +import json +import os +import sys +import traceback +from datetime import timedelta +from pathlib import Path + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from couchbase.auth import PasswordAuthenticator +from couchbase.cluster import Cluster +from couchbase.options import ClusterOptions + +from tools.index import ( + create_index_from_recommendation, + get_index_advisor_recommendations, +) + + +def print_section(title): + """Print a formatted section header.""" + print("\n" + "=" * 80) + print(f" {title}") + print("=" * 80 + "\n") + + +def create_mock_context(cluster, read_only_mode=True): + """Create a mock context for testing.""" + + class MockLifespanContext: + def __init__(self, cluster, read_only): + self.cluster = cluster + self.read_only_query_mode = read_only + + class MockRequestContext: + def __init__(self, cluster, read_only): + self.lifespan_context = MockLifespanContext(cluster, read_only) + + class MockContext: + def __init__(self, cluster, read_only): + self.request_context = MockRequestContext(cluster, read_only) + self._cluster = cluster + + return MockContext(cluster, read_only_mode) + + +def main(): # noqa: PLR0915 + print_section("Couchbase Index Advisor Test Script") + + # Get connection details from environment + connection_string = os.getenv("CB_CONNECTION_STRING") + username = os.getenv("CB_USERNAME") + password = os.getenv("CB_PASSWORD") + + if not all([connection_string, username, password]): + print("โŒ Error: Missing required environment variables") + print("\nPlease set the following environment variables:") + print(" - CB_CONNECTION_STRING") + print(" - CB_USERNAME") + print(" - CB_PASSWORD") + print("\nExample:") + print( + " export CB_CONNECTION_STRING='couchbases://your-cluster.cloud.couchbase.com'" + ) + print(" export CB_USERNAME='your-username'") + print(" export CB_PASSWORD='your-password'") + return 1 + + print(f"๐Ÿ“ก Connecting to: {connection_string}") + print(f"๐Ÿ‘ค Username: {username}") + + try: + # Connect to cluster + auth = PasswordAuthenticator(username, password) + cluster = Cluster(connection_string, ClusterOptions(auth)) + + # Wait for cluster to be ready + cluster.wait_until_ready(timedelta(seconds=10)) + print("โœ… Connected successfully!\n") + + except Exception as e: + print(f"โŒ Failed to connect to Couchbase: {e}") + return 1 + + # Test 1: Get Index Advisor Recommendations + print_section("Test 1: Get Index Advisor Recommendations") + + # You can modify these values for your specific use case + bucket_name = "travel-sample" + scope_name = "inventory" + test_query = "SELECT * FROM landmark WHERE activity = 'eat' AND city = 'Paris'" + + print(f"Bucket: {bucket_name}") + print(f"Scope: {scope_name}") + print(f"Query: {test_query}\n") + + try: + ctx = create_mock_context(cluster, read_only_mode=True) + + print("๐Ÿ” Running Index Advisor...\n") + recommendations = get_index_advisor_recommendations( + ctx, bucket_name, scope_name, test_query + ) + + print("๐Ÿ“Š Index Advisor Results:") + print(json.dumps(recommendations, indent=2)) + + # Display summary + summary = recommendations.get("summary", {}) + print("\n๐Ÿ“ˆ Summary:") + print(f" Current indexes used: {summary.get('current_indexes_count', 0)}") + print(f" Recommended indexes: {summary.get('recommended_indexes_count', 0)}") + print( + f" Recommended covering indexes: {summary.get('recommended_covering_indexes_count', 0)}" + ) + print( + f" Has recommendations: {'โœ… Yes' if summary.get('has_recommendations') else 'โŒ No'}" + ) + + # Test 2: Create Index from Recommendation (if recommendations exist) + if recommendations.get("recommended_indexes"): + print_section("Test 2: Create Index from Recommendation") + + first_recommendation = recommendations["recommended_indexes"][0] + index_definition = first_recommendation["index"] + + print("๐Ÿ“ Index to create:") + print(f" {index_definition}\n") + + # Ask user if they want to create the index + print("โš ๏ธ Note: Creating an index requires:") + print( + " - Read-only mode to be disabled (CB_MCP_READ_ONLY_QUERY_MODE=false)" + ) + print(" - Appropriate permissions on the cluster") + print(" - The index name must not already exist") + + response = input( + "\nโ“ Do you want to attempt to create this index? (y/N): " + ) + + if response.lower() == "y": + # Test with read-only mode disabled + ctx_write = create_mock_context(cluster, read_only_mode=False) + + print("\n๐Ÿ”จ Attempting to create index...\n") + result = create_index_from_recommendation(ctx_write, index_definition) + + print("๐Ÿ“Š Create Index Result:") + print(json.dumps(result, indent=2)) + + if result.get("status") == "success": + print("\nโœ… Index created successfully!") + else: + print(f"\nโŒ Failed to create index: {result.get('message')}") + else: + print("\nโญ๏ธ Skipped index creation") + + # Still test the function with read-only mode enabled + print( + "\n๐Ÿงช Testing create_index_from_recommendation with read-only mode..." + ) + ctx_readonly = create_mock_context(cluster, read_only_mode=True) + result = create_index_from_recommendation( + ctx_readonly, index_definition + ) + print(json.dumps(result, indent=2)) + else: + print("\nโ„น๏ธ No index recommendations available to test creation.") # noqa: RUF001 + + print_section("Test Completed Successfully") + cluster.close() + return 0 + + except Exception as e: + print(f"\nโŒ Error during test: {e}") + traceback.print_exc() + cluster.close() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From d4087871f85212646681da70214f6cfed885f12f Mon Sep 17 00:00:00 2001 From: AayushTyagi1 Date: Thu, 16 Oct 2025 11:30:57 +0530 Subject: [PATCH 3/5] DA-1170: Fixed: Review comments --- src/tools/index.py | 13 +++++-------- tests/test_index_advisor.py | 21 ++++++++++----------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/src/tools/index.py b/src/tools/index.py index 75c63f7..eca6535 100644 --- a/src/tools/index.py +++ b/src/tools/index.py @@ -16,13 +16,12 @@ logger = logging.getLogger(f"{MCP_SERVER_NAME}.tools.index") -def get_index_advisor_recommendations( - ctx: Context, bucket_name: str, scope_name: str, query: str -) -> dict[str, Any]: +def get_index_advisor_recommendations(ctx: Context, query: str) -> dict[str, Any]: """Get index recommendations from Couchbase Index Advisor for a given SQL++ query. The Index Advisor analyzes the query and provides recommendations for optimal indexes. This tool works with SELECT, UPDATE, DELETE, or MERGE queries. + The query should contain fully qualified keyspace (e.g., bucket.scope.collection). Returns a dictionary with: - current_used_indexes: Array of currently used indexes (if any) @@ -39,18 +38,16 @@ def get_index_advisor_recommendations( # Escape single quotes in the query by doubling them for SQL++ string literal escaped_query = query.replace("'", "''") - # Build the ADVISOR query with fully qualified keyspace + # Build the ADVISOR query advisor_query = f"SELECT ADVISOR('{escaped_query}') AS advisor_result" - logger.info(f"Running Index Advisor for query in {bucket_name}.{scope_name}") + logger.info("Running Index Advisor for the provided query") # Execute the ADVISOR function at cluster level result = cluster.query(advisor_query) # Extract the advisor result from the query response - advisor_results = [] - for row in result: - advisor_results.append(row) + advisor_results = list(result) if not advisor_results: return { diff --git a/tests/test_index_advisor.py b/tests/test_index_advisor.py index 02cd303..9324209 100755 --- a/tests/test_index_advisor.py +++ b/tests/test_index_advisor.py @@ -7,11 +7,12 @@ 2. create_index_from_recommendation - Create an index from recommendations Usage: - uv run test_index_advisor.py + uv run tests/test_index_advisor.py Prerequisites: - Set environment variables: CB_CONNECTION_STRING, CB_USERNAME, CB_PASSWORD - - Have a Couchbase cluster accessible with the travel-sample bucket (or modify the script) + - Have a Couchbase cluster accessible with the travel-sample bucket (or modify the query) + - Query should use fully qualified keyspace (bucket.scope.collection) """ import json @@ -102,22 +103,20 @@ def main(): # noqa: PLR0915 # Test 1: Get Index Advisor Recommendations print_section("Test 1: Get Index Advisor Recommendations") - # You can modify these values for your specific use case - bucket_name = "travel-sample" - scope_name = "inventory" - test_query = "SELECT * FROM landmark WHERE activity = 'eat' AND city = 'Paris'" + # You can modify this query for your specific use case + # Note: The query should contain fully qualified keyspace (bucket.scope.collection) + test_query = ( + "SELECT * FROM `travel-sample`.inventory.landmark " + "WHERE activity = 'eat' AND city = 'Paris'" + ) - print(f"Bucket: {bucket_name}") - print(f"Scope: {scope_name}") print(f"Query: {test_query}\n") try: ctx = create_mock_context(cluster, read_only_mode=True) print("๐Ÿ” Running Index Advisor...\n") - recommendations = get_index_advisor_recommendations( - ctx, bucket_name, scope_name, test_query - ) + recommendations = get_index_advisor_recommendations(ctx, test_query) print("๐Ÿ“Š Index Advisor Results:") print(json.dumps(recommendations, indent=2)) From 190d0f7f8f1a3be9b1455ed5c7f5a3b5dc919c18 Mon Sep 17 00:00:00 2001 From: AayushTyagi1 Date: Wed, 22 Oct 2025 04:02:13 +0530 Subject: [PATCH 4/5] DA-1170 Modified Review comment addressed --- src/tools/__init__.py | 7 +-- src/tools/index.py | 102 +++--------------------------------------- 2 files changed, 6 insertions(+), 103 deletions(-) diff --git a/src/tools/__init__.py b/src/tools/__init__.py index 83c36f5..352d458 100644 --- a/src/tools/__init__.py +++ b/src/tools/__init__.py @@ -5,10 +5,7 @@ """ # Index tools -from .index import ( - create_index_from_recommendation, - get_index_advisor_recommendations, -) +from .index import get_index_advisor_recommendations # Key-Value tools from .kv import ( @@ -47,7 +44,6 @@ get_schema_for_collection, run_sql_plus_plus_query, get_index_advisor_recommendations, - create_index_from_recommendation, ] __all__ = [ @@ -64,7 +60,6 @@ "get_schema_for_collection", "run_sql_plus_plus_query", "get_index_advisor_recommendations", - "create_index_from_recommendation", # Convenience "ALL_TOOLS", ] diff --git a/src/tools/index.py b/src/tools/index.py index eca6535..841126d 100644 --- a/src/tools/index.py +++ b/src/tools/index.py @@ -1,8 +1,7 @@ """ Tools for index operations and optimization. -This module contains tools for getting index recommendations using the Couchbase Index Advisor -and creating indexes based on those recommendations. +This module contains tools for getting index recommendations using the Couchbase Index Advisor. """ import logging @@ -10,8 +9,8 @@ from mcp.server.fastmcp import Context +from tools.query import run_cluster_query from utils.constants import MCP_SERVER_NAME -from utils.context import get_cluster_connection logger = logging.getLogger(f"{MCP_SERVER_NAME}.tools.index") @@ -32,22 +31,14 @@ def get_index_advisor_recommendations(ctx: Context, query: str) -> dict[str, Any - index: The CREATE INDEX SQL++ command - statements: Array of statement objects with the query and run count """ - cluster = get_cluster_connection(ctx) - try: - # Escape single quotes in the query by doubling them for SQL++ string literal - escaped_query = query.replace("'", "''") - # Build the ADVISOR query - advisor_query = f"SELECT ADVISOR('{escaped_query}') AS advisor_result" + advisor_query = f"SELECT ADVISOR('{query}') AS advisor_result" logger.info("Running Index Advisor for the provided query") - # Execute the ADVISOR function at cluster level - result = cluster.query(advisor_query) - - # Extract the advisor result from the query response - advisor_results = list(result) + # Execute the ADVISOR function at cluster level using run_cluster_query + advisor_results = run_cluster_query(ctx, advisor_query) if not advisor_results: return { @@ -91,86 +82,3 @@ def get_index_advisor_recommendations(ctx: Context, query: str) -> dict[str, Any except Exception as e: logger.error(f"Error running Index Advisor: {e!s}", exc_info=True) raise - - -def create_index_from_recommendation( - ctx: Context, index_definition: str -) -> dict[str, Any]: - """Create an index using the provided CREATE INDEX statement. - - This tool executes a CREATE INDEX statement, typically from Index Advisor recommendations. - Note: This operation requires write permissions and will fail if: - - The read-only query mode is enabled - - The user lacks CREATE INDEX permissions - - An index with the same name already exists - - The index_definition should be a complete CREATE INDEX statement, for example: - CREATE INDEX adv_city_activity ON `travel-sample`.inventory.landmark(city, activity) - - Returns a dictionary with: - - status: 'success' or 'error' - - message: Description of the result - - index_definition: The index statement that was executed (on success) - """ - cluster = get_cluster_connection(ctx) - - app_context = ctx.request_context.lifespan_context - read_only_query_mode = app_context.read_only_query_mode - - # Check if read-only mode is enabled - if read_only_query_mode: - logger.error("Cannot create index in read-only query mode") - return { - "status": "error", - "message": "Index creation is not allowed in read-only query mode. Please disable read-only mode (CB_MCP_READ_ONLY_QUERY_MODE=false) to create indexes.", - "index_definition": index_definition, - } - - try: - # Validate that the statement is a CREATE INDEX statement - if not index_definition.strip().upper().startswith("CREATE INDEX"): - logger.error("Invalid index definition: must start with CREATE INDEX") - return { - "status": "error", - "message": "Invalid index definition. The statement must be a CREATE INDEX command.", - "index_definition": index_definition, - } - - logger.info(f"Creating index with definition: {index_definition}") - - # Execute the CREATE INDEX statement at cluster level - result = cluster.query(index_definition) - - # Consume the result to ensure the query completes - for _ in result: - pass - - logger.info("Index created successfully") - - return { - "status": "success", - "message": "Index created successfully", - "index_definition": index_definition, - } - - except Exception as e: - error_message = str(e) - logger.error(f"Error creating index: {error_message}", exc_info=True) - - # Provide helpful error messages for common issues - if "already exists" in error_message.lower(): - message = "An index with this name already exists. Consider using a different name or dropping the existing index first." - elif ( - "permission" in error_message.lower() - or "authorized" in error_message.lower() - ): - message = "Insufficient permissions to create index. Please ensure your user has the required permissions." - else: - message = f"Failed to create index: {error_message}" - - return { - "status": "error", - "message": message, - "index_definition": index_definition, - "error_details": error_message, - } From c1bd8e43b3273d2073ab06bdb7b123c365c68b8b Mon Sep 17 00:00:00 2001 From: AayushTyagi1 Date: Wed, 22 Oct 2025 23:37:23 +0530 Subject: [PATCH 5/5] DA-1170: Deleted: Tests --- tests/test_index_advisor.py | 199 ------------------------------------ 1 file changed, 199 deletions(-) delete mode 100755 tests/test_index_advisor.py diff --git a/tests/test_index_advisor.py b/tests/test_index_advisor.py deleted file mode 100755 index 9324209..0000000 --- a/tests/test_index_advisor.py +++ /dev/null @@ -1,199 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for Index Advisor functionality. - -This script demonstrates how to use the new Index Advisor tools: -1. get_index_advisor_recommendations - Get index recommendations for a query -2. create_index_from_recommendation - Create an index from recommendations - -Usage: - uv run tests/test_index_advisor.py - -Prerequisites: - - Set environment variables: CB_CONNECTION_STRING, CB_USERNAME, CB_PASSWORD - - Have a Couchbase cluster accessible with the travel-sample bucket (or modify the query) - - Query should use fully qualified keyspace (bucket.scope.collection) -""" - -import json -import os -import sys -import traceback -from datetime import timedelta -from pathlib import Path - -# Add src to path -sys.path.insert(0, str(Path(__file__).parent.parent / "src")) - -from couchbase.auth import PasswordAuthenticator -from couchbase.cluster import Cluster -from couchbase.options import ClusterOptions - -from tools.index import ( - create_index_from_recommendation, - get_index_advisor_recommendations, -) - - -def print_section(title): - """Print a formatted section header.""" - print("\n" + "=" * 80) - print(f" {title}") - print("=" * 80 + "\n") - - -def create_mock_context(cluster, read_only_mode=True): - """Create a mock context for testing.""" - - class MockLifespanContext: - def __init__(self, cluster, read_only): - self.cluster = cluster - self.read_only_query_mode = read_only - - class MockRequestContext: - def __init__(self, cluster, read_only): - self.lifespan_context = MockLifespanContext(cluster, read_only) - - class MockContext: - def __init__(self, cluster, read_only): - self.request_context = MockRequestContext(cluster, read_only) - self._cluster = cluster - - return MockContext(cluster, read_only_mode) - - -def main(): # noqa: PLR0915 - print_section("Couchbase Index Advisor Test Script") - - # Get connection details from environment - connection_string = os.getenv("CB_CONNECTION_STRING") - username = os.getenv("CB_USERNAME") - password = os.getenv("CB_PASSWORD") - - if not all([connection_string, username, password]): - print("โŒ Error: Missing required environment variables") - print("\nPlease set the following environment variables:") - print(" - CB_CONNECTION_STRING") - print(" - CB_USERNAME") - print(" - CB_PASSWORD") - print("\nExample:") - print( - " export CB_CONNECTION_STRING='couchbases://your-cluster.cloud.couchbase.com'" - ) - print(" export CB_USERNAME='your-username'") - print(" export CB_PASSWORD='your-password'") - return 1 - - print(f"๐Ÿ“ก Connecting to: {connection_string}") - print(f"๐Ÿ‘ค Username: {username}") - - try: - # Connect to cluster - auth = PasswordAuthenticator(username, password) - cluster = Cluster(connection_string, ClusterOptions(auth)) - - # Wait for cluster to be ready - cluster.wait_until_ready(timedelta(seconds=10)) - print("โœ… Connected successfully!\n") - - except Exception as e: - print(f"โŒ Failed to connect to Couchbase: {e}") - return 1 - - # Test 1: Get Index Advisor Recommendations - print_section("Test 1: Get Index Advisor Recommendations") - - # You can modify this query for your specific use case - # Note: The query should contain fully qualified keyspace (bucket.scope.collection) - test_query = ( - "SELECT * FROM `travel-sample`.inventory.landmark " - "WHERE activity = 'eat' AND city = 'Paris'" - ) - - print(f"Query: {test_query}\n") - - try: - ctx = create_mock_context(cluster, read_only_mode=True) - - print("๐Ÿ” Running Index Advisor...\n") - recommendations = get_index_advisor_recommendations(ctx, test_query) - - print("๐Ÿ“Š Index Advisor Results:") - print(json.dumps(recommendations, indent=2)) - - # Display summary - summary = recommendations.get("summary", {}) - print("\n๐Ÿ“ˆ Summary:") - print(f" Current indexes used: {summary.get('current_indexes_count', 0)}") - print(f" Recommended indexes: {summary.get('recommended_indexes_count', 0)}") - print( - f" Recommended covering indexes: {summary.get('recommended_covering_indexes_count', 0)}" - ) - print( - f" Has recommendations: {'โœ… Yes' if summary.get('has_recommendations') else 'โŒ No'}" - ) - - # Test 2: Create Index from Recommendation (if recommendations exist) - if recommendations.get("recommended_indexes"): - print_section("Test 2: Create Index from Recommendation") - - first_recommendation = recommendations["recommended_indexes"][0] - index_definition = first_recommendation["index"] - - print("๐Ÿ“ Index to create:") - print(f" {index_definition}\n") - - # Ask user if they want to create the index - print("โš ๏ธ Note: Creating an index requires:") - print( - " - Read-only mode to be disabled (CB_MCP_READ_ONLY_QUERY_MODE=false)" - ) - print(" - Appropriate permissions on the cluster") - print(" - The index name must not already exist") - - response = input( - "\nโ“ Do you want to attempt to create this index? (y/N): " - ) - - if response.lower() == "y": - # Test with read-only mode disabled - ctx_write = create_mock_context(cluster, read_only_mode=False) - - print("\n๐Ÿ”จ Attempting to create index...\n") - result = create_index_from_recommendation(ctx_write, index_definition) - - print("๐Ÿ“Š Create Index Result:") - print(json.dumps(result, indent=2)) - - if result.get("status") == "success": - print("\nโœ… Index created successfully!") - else: - print(f"\nโŒ Failed to create index: {result.get('message')}") - else: - print("\nโญ๏ธ Skipped index creation") - - # Still test the function with read-only mode enabled - print( - "\n๐Ÿงช Testing create_index_from_recommendation with read-only mode..." - ) - ctx_readonly = create_mock_context(cluster, read_only_mode=True) - result = create_index_from_recommendation( - ctx_readonly, index_definition - ) - print(json.dumps(result, indent=2)) - else: - print("\nโ„น๏ธ No index recommendations available to test creation.") # noqa: RUF001 - - print_section("Test Completed Successfully") - cluster.close() - return 0 - - except Exception as e: - print(f"\nโŒ Error during test: {e}") - traceback.print_exc() - cluster.close() - return 1 - - -if __name__ == "__main__": - sys.exit(main())