Merge pull request #125 from schemacrawler/exclude

Allow excluding of certain tools

Author: sualeh

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/McpServerContext.java

@@ -13,6 +13,7 @@
 import static us.fatehi.utility.Utility.trimToEmpty;
 
 import java.nio.file.Path;
+import java.util.Collection;
 import java.util.logging.Level;
 import schemacrawler.schema.Catalog;
 import schemacrawler.schemacrawler.InfoLevel;
@@ -22,6 +23,7 @@
 import schemacrawler.schemacrawler.LoadOptionsBuilder;
 import schemacrawler.schemacrawler.SchemaCrawlerOptions;
 import schemacrawler.schemacrawler.SchemaCrawlerOptionsBuilder;
+import schemacrawler.tools.ai.mcpserver.utility.CollectionsUtility;
 import schemacrawler.tools.databaseconnector.EnvironmentalDatabaseConnectionSourceBuilder;
 import schemacrawler.tools.offline.jdbc.OfflineConnectionUtility;
 import schemacrawler.tools.utility.SchemaCrawlerUtility;
@@ -32,6 +34,12 @@
 /** Inner class that handles the MCP server setup. */
 final class McpServerContext {
 
+  private static final String EXCLUDE_TOOLS = "SCHCRWLR_EXCLUDE_TOOLS";
+  private static final String INFO_LEVEL = "SCHCRWLR_INFO_LEVEL";
+  private static final String LOG_LEVEL = "SCHCRWLR_LOG_LEVEL";
+  private static final String MCP_SERVER_TRANSPORT = "SCHCRWLR_MCP_SERVER_TRANSPORT";
+  private static final String OFFLINE_DATABASE = "SCHCRWLR_OFFLINE_DATABASE";
+
   private final EnvironmentVariableAccessor envAccessor;
   private final McpServerTransportType transport;
   private final SchemaCrawlerOptions schemaCrawlerOptions;
@@ -57,8 +65,7 @@ protected McpServerContext(final EnvironmentVariableAccessor envAccessor) {
   }
 
   protected DatabaseConnectionSource buildCatalogDatabaseConnectionSource() {
-    final String offlineDatabasePathString =
-        trimToEmpty(envAccessor.getenv("SCHCRWLR_OFFLINE_DATABASE"));
+    final String offlineDatabasePathString = trimToEmpty(envAccessor.getenv(OFFLINE_DATABASE));
     if (isBlank(offlineDatabasePathString)) {
       return buildOperationsDatabaseConnectionSource();
     }
@@ -80,11 +87,6 @@ protected DatabaseConnectionSource buildOperationsDatabaseConnectionSource() {
     return databaseConnectionSource;
   }
 
-  /**
-   * Adds SchemaCrawler specific arguments to the arguments list.
-   *
-   * @param arguments The list of arguments to add to
-   */
   protected SchemaCrawlerOptions buildSchemaCrawlerOptions() {
     final InfoLevel infoLevel = readInfoLevel();
 
@@ -104,6 +106,10 @@ protected SchemaCrawlerOptions buildSchemaCrawlerOptions() {
     return schemaCrawlerOptions;
   }
 
+  protected Collection<String> excludeTools() {
+    return CollectionsUtility.setOfStrings(envAccessor.getenv(EXCLUDE_TOOLS));
+  }
+
   protected Catalog getCatalog() {
     final DatabaseConnectionSource connectionSource = buildCatalogDatabaseConnectionSource();
     final SchemaCrawlerOptions schemaCrawlerOptions = getSchemaCrawlerOptions();
@@ -129,7 +135,7 @@ protected InfoLevel readInfoLevel() {
 
     final InfoLevel defaultValue = InfoLevel.standard;
     try {
-      final String value = envAccessor.getenv("SCHCRWLR_INFO_LEVEL");
+      final String value = envAccessor.getenv(INFO_LEVEL);
       InfoLevel infoLevel = InfoLevel.valueOfFromString(value);
       if (infoLevel == InfoLevel.unknown) {
         infoLevel = defaultValue;
@@ -150,7 +156,7 @@ protected Level readLogLevel() {
 
     final Level defaultValue = Level.INFO;
 
-    final String value = envAccessor.getenv("SCHCRWLR_LOG_LEVEL");
+    final String value = envAccessor.getenv(LOG_LEVEL);
     if (isBlank(value)) {
       return defaultValue;
     }
@@ -172,7 +178,7 @@ private McpServerTransportType readTransport() {
 
     final McpServerTransportType defaultValue = McpServerTransportType.stdio;
 
-    final String value = envAccessor.getenv("SCHCRWLR_MCP_SERVER_TRANSPORT");
+    final String value = envAccessor.getenv(MCP_SERVER_TRANSPORT);
     if (isBlank(value)) {
       return defaultValue;
     }

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/McpServerInitializer.java

@@ -8,9 +8,12 @@
 
 package schemacrawler.tools.ai.mcpserver;
 
+import static java.util.Collections.emptySet;
+import static java.util.Collections.unmodifiableCollection;
 import static java.util.Objects.requireNonNull;
 
 import java.sql.Connection;
+import java.util.Collection;
 import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.lang.NonNull;
@@ -28,11 +31,13 @@ public class McpServerInitializer
   private final Catalog catalog;
   private final DatabaseConnectionSource connectionSource;
   private final McpServerTransportType mcpTransport;
+  private final Collection<String> excludeTools;
 
   public McpServerInitializer(
       final Catalog catalog,
       final Connection connection,
-      final McpServerTransportType mcpTransport) {
+      final McpServerTransportType mcpTransport,
+      final Collection<String> excludeTools) {
 
     this.mcpTransport = requireNonNull(mcpTransport, "No MCP Server transport provided");
     if (mcpTransport == McpServerTransportType.unknown) {
@@ -54,6 +59,12 @@ public McpServerInitializer(
     this.isInErrorState = isInErrorState;
 
     this.catalog = requireNonNull(catalog, "No catalog provided");
+
+    if (excludeTools == null) {
+      this.excludeTools = emptySet();
+    } else {
+      this.excludeTools = unmodifiableCollection(excludeTools);
+    }
   }
 
   public McpServerInitializer(final McpServerContext context) {
@@ -87,6 +98,8 @@ public McpServerInitializer(final McpServerContext context) {
       connectionSource = EmptyFactory.createEmptyDatabaseConnectionSource();
     }
     this.connectionSource = connectionSource;
+
+    excludeTools = context.excludeTools();
   }
 
   @Override
@@ -102,5 +115,6 @@ public void initialize(@NonNull final GenericApplicationContext context) {
         "functionDefinitionRegistry",
         FunctionDefinitionRegistry.class,
         () -> FunctionDefinitionRegistry.getFunctionDefinitionRegistry());
+    context.registerBean("excludeTools", Collection.class, () -> excludeTools);
   }
 }

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/McpServerMain.java

@@ -9,6 +9,7 @@
 package schemacrawler.tools.ai.mcpserver;
 
 import java.sql.Connection;
+import java.util.Collection;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
@@ -51,9 +52,10 @@ public static void startMcpServer() {
   public static void startMcpServer(
       final Catalog catalog,
       final Connection connection,
-      final McpServerTransportType mcpTransport) {
+      final McpServerTransportType mcpTransport,
+      final Collection<String> excludeTools) {
     new SpringApplicationBuilder(McpServer.class)
-        .initializers(new McpServerInitializer(catalog, connection, mcpTransport))
+        .initializers(new McpServerInitializer(catalog, connection, mcpTransport, excludeTools))
         .profiles(mcpTransport.name())
         .run();
     LOGGER.log(

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/server/ToolProvider.java

@@ -15,6 +15,7 @@
 import io.modelcontextprotocol.server.McpSyncServerExchange;
 import io.modelcontextprotocol.spec.McpSchema.Implementation;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -42,6 +43,7 @@ public class ToolProvider {
   @Autowired private ServerHealth serverHealth;
   @Autowired private FunctionDefinitionRegistry functionDefinitionRegistry;
   @Autowired private ToolHelper toolHelper;
+  @Autowired private Collection<String> excludeTools;
 
   @McpTool(
       name = "mcp-server-health",
@@ -81,9 +83,14 @@ public List<McpServerFeatures.SyncToolSpecification> schemaCrawlerTools() {
     final List<McpServerFeatures.SyncToolSpecification> tools = new ArrayList<>();
     for (final FunctionDefinition<?> functionDefinition :
         functionDefinitionRegistry.getFunctionDefinitions()) {
-      LOGGER.log(
-          Level.INFO,
-          new StringFormat("Adding callback for:%n%s", functionDefinition.getFunctionName()));
+      final String functionName = functionDefinition.getFunctionName().getName();
+      if (excludeTools == null || excludeTools.contains(functionName)) {
+        LOGGER.log(Level.WARNING, new StringFormat("Excluding tool <%s>", functionName));
+        continue;
+      }
+      LOGGER.log(Level.INFO, new StringFormat("Adding callback for <%s>", functionName));
+      LOGGER.log(Level.FINE, new StringFormat("%s", functionDefinition.getFunctionName()));
+
       final McpServerFeatures.SyncToolSpecification toolSpecification =
           toolHelper.toSyncToolSpecification(functionDefinition);
       tools.add(toolSpecification);

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/utility/CollectionsUtility.java

@@ -0,0 +1,29 @@
+package schemacrawler.tools.ai.mcpserver.utility;
+
+import static us.fatehi.utility.Utility.trimToEmpty;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+import us.fatehi.utility.UtilityMarker;
+
+@UtilityMarker
+public class CollectionsUtility {
+
+  public static Collection<String> setOfStrings(final String input) {
+    final Set<String> setOfStrings =
+        Arrays.stream(trimToEmpty(input).split(","))
+            .filter(Objects::nonNull)
+            .map(String::strip)
+            .filter(s -> !s.isEmpty())
+            .collect(Collectors.toSet());
+
+    return setOfStrings;
+  }
+
+  private CollectionsUtility() {
+    // Prevent instantiation
+  }
+}

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/ai/mcpserver/utility/EmptyFactory.java

@@ -12,8 +12,10 @@
 import java.lang.reflect.Proxy;
 import java.sql.DriverManager;
 import schemacrawler.schema.Catalog;
+import us.fatehi.utility.UtilityMarker;
 import us.fatehi.utility.datasource.DatabaseConnectionSource;
 
+@UtilityMarker
 public class EmptyFactory {
 
   public static Catalog createEmptyCatalog(final Exception e) {

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/command/mcpserver/McpServerCommand.java

@@ -8,6 +8,7 @@
 
 package schemacrawler.tools.command.mcpserver;
 
+import java.util.Collection;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import schemacrawler.tools.ai.mcpserver.McpServerMain;
@@ -35,7 +36,8 @@ public void checkAvailability() throws RuntimeException {
   @Override
   public void execute() {
     final McpServerTransportType mcpTransport = commandOptions.mcpTransport();
-    McpServerMain.startMcpServer(catalog, connection, mcpTransport);
+    final Collection<String> excludeTools = commandOptions.excludeTools();
+    McpServerMain.startMcpServer(catalog, connection, mcpTransport, excludeTools);
   }
 
   @Override

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/command/mcpserver/McpServerCommandOptions.java

@@ -8,15 +8,20 @@
 
 package schemacrawler.tools.command.mcpserver;
 
+import static java.util.Objects.requireNonNull;
+
+import java.util.Collection;
 import schemacrawler.tools.ai.mcpserver.McpServerTransportType;
 import schemacrawler.tools.executable.CommandOptions;
 
-public record McpServerCommandOptions(McpServerTransportType mcpTransport)
+public record McpServerCommandOptions(
+    McpServerTransportType mcpTransport, Collection<String> excludeTools)
     implements CommandOptions {
 
   public McpServerCommandOptions {
     if (mcpTransport == null || mcpTransport == McpServerTransportType.unknown) {
       throw new IllegalArgumentException("No MCP Server transport specified");
     }
+    requireNonNull(excludeTools, "No exclude tools list provided");
   }
 }

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/command/mcpserver/McpServerCommandOptionsBuilder.java

@@ -9,8 +9,12 @@
 package schemacrawler.tools.command.mcpserver;
 
 import static schemacrawler.tools.ai.mcpserver.McpServerTransportType.unknown;
+
+import java.util.Collection;
+import java.util.Collections;
 import schemacrawler.schemacrawler.OptionsBuilder;
 import schemacrawler.tools.ai.mcpserver.McpServerTransportType;
+import schemacrawler.tools.ai.mcpserver.utility.CollectionsUtility;
 import schemacrawler.tools.options.Config;
 import schemacrawler.tools.options.ConfigOptionsBuilder;
 
@@ -23,17 +27,20 @@ public static McpServerCommandOptionsBuilder builder() {
   }
 
   private McpServerTransportType mcpTransport;
+  private Collection<String> excludeTools;
 
   private McpServerCommandOptionsBuilder() {
     // MCP Server transport needs to be explicitly specified,
     // so default to known
     mcpTransport = unknown;
+    excludeTools = Collections.emptySet();
   }
 
   @Override
   public McpServerCommandOptionsBuilder fromConfig(final Config config) {
     if (config != null) {
       mcpTransport = config.getEnumValue("transport", unknown);
+      excludeTools = CollectionsUtility.setOfStrings(config.getStringValue("exclude-tools", ""));
     }
 
     return this;
@@ -43,18 +50,19 @@ public McpServerCommandOptionsBuilder fromConfig(final Config config) {
   public McpServerCommandOptionsBuilder fromOptions(final McpServerCommandOptions options) {
     if (options != null) {
       mcpTransport = options.mcpTransport();
+      excludeTools = options.excludeTools();
     }
     return this;
   }
 
   @Override
   public Config toConfig() {
-    throw new UnsupportedOperationException("Cannot load transport from config file");
+    throw new UnsupportedOperationException("Cannot load MCP Server options from config file");
   }
 
   @Override
   public McpServerCommandOptions toOptions() {
-    return new McpServerCommandOptions(mcpTransport);
+    return new McpServerCommandOptions(mcpTransport, excludeTools);
   }
 
   public McpServerCommandOptionsBuilder withMcpTransport(

schemacrawler-ai-mcpserver/src/main/java/schemacrawler/tools/command/mcpserver/McpServerCommandProvider.java

@@ -30,6 +30,8 @@ public PluginCommand getCommandLineCommand() {
     final PluginCommand pluginCommand = newPluginCommand(COMMAND);
     pluginCommand.addOption(
         "transport", McpServerTransportType.class, "Type of MCP server to start");
+    pluginCommand.addOption(
+        "exclude-tools", String.class, "Comma-separated list of tools to exclude");
     return pluginCommand;
   }