Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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();
}
Expand All @@ -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();

Expand All @@ -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();
Expand All @@ -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;
Expand All @@ -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;
}
Expand All @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) {
Expand All @@ -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) {
Expand Down Expand Up @@ -87,6 +98,8 @@ public McpServerInitializer(final McpServerContext context) {
connectionSource = EmptyFactory.createEmptyDatabaseConnectionSource();
}
this.connectionSource = connectionSource;

excludeTools = context.excludeTools();
}

@Override
Expand All @@ -102,5 +115,6 @@ public void initialize(@NonNull final GenericApplicationContext context) {
"functionDefinitionRegistry",
FunctionDefinitionRegistry.class,
() -> FunctionDefinitionRegistry.getFunctionDefinitionRegistry());
context.registerBean("excludeTools", Collection.class, () -> excludeTools);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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;
Expand All @@ -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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
Loading
Loading