You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@shardingsphere.apache.org by zh...@apache.org on 2026/07/18 03:39:35 UTC
(shardingsphere) branch master updated: Improve MCP protocol and SQL safety boundaries (#39174)
This is an automated email from the ASF dual-hosted git repository.
terrymanu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shardingsphere.git
The following commit(s) were added to refs/heads/master by this push:
new cfab2d7abda Improve MCP protocol and SQL safety boundaries (#39174)
cfab2d7abda is described below
commit cfab2d7abda637ba8043e437ffae31033352fe65
Author: Liang Zhang <zh...@apache.org>
AuthorDate: Sat Jul 18 11:39:28 2026 +0800
Improve MCP protocol and SQL safety boundaries (#39174)
* Improve MCP protocol and SQL safety boundaries
- enforce initialization and session lifecycle contracts
- make SQL scanning dialect-aware and preserve object attribution
- tighten descriptor validation and simplify workflow guidance
* Improve MCP protocol and SQL safety boundaries
- enforce initialization and session lifecycle contracts
- make SQL scanning dialect-aware and preserve object attribution
- tighten descriptor validation and simplify workflow guidance
---
.../mcp/api/session/MCPSessionIdentity.java | 9 ++-
.../mcp/api/session/MCPSessionIdentityTest.java | 18 ++++-
.../server/http/StreamableHttpMCPServlet.java | 14 ++--
...dingSphereServerTransportSecurityValidator.java | 7 +-
.../server/http/StreamableHttpMCPServletTest.java | 39 +++++++++-
...SphereServerTransportSecurityValidatorTest.java | 15 ++++
.../execute/ExplainSQLCandidateValidator.java | 21 +++--
.../tool/handler/execute/MCPStatementAnalyzer.java | 63 +++++++++++++--
.../execute/SQLStatementObjectExtractor.java | 38 ++++++++-
.../tool/handler/execute/SQLStatementScanner.java | 83 ++++++++++++++++++--
.../tool/handler/execute/SQLStatementToken.java | 9 ++-
.../handler/execute/MCPSQLExecutionFacadeTest.java | 5 ++
.../handler/execute/MCPStatementAnalyzerTest.java | 66 ++++++++++++++++
.../MCPToolDescriptorCatalogValidator.java | 8 +-
.../service/WorkflowGuidancePayloadBuilder.java | 12 +--
.../MCPDescriptorCatalogValidatorTest.java | 57 ++++++++++++++
.../HttpTransportProtocolContractE2ETest.java | 12 +++
.../HttpTransportSessionAttributionE2ETest.java | 89 ++++++++++++++++++++++
.../HttpTransportSessionLifecycleE2ETest.java | 2 +-
19 files changed, 501 insertions(+), 66 deletions(-)
diff --git a/mcp/api/src/main/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentity.java b/mcp/api/src/main/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentity.java
index a11cc6f3791..026c7576640 100644
--- a/mcp/api/src/main/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentity.java
+++ b/mcp/api/src/main/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentity.java
@@ -18,7 +18,6 @@
package org.apache.shardingsphere.mcp.api.session;
import lombok.Getter;
-import lombok.RequiredArgsConstructor;
import java.util.Map;
@@ -28,7 +27,6 @@ import java.util.Map;
* <p>The session identifier is always present. Subject, source and attributes are optional trusted attribution captured when the session is initialized.
* Attribution does not represent authentication or authorization.</p>
*/
-@RequiredArgsConstructor
@Getter
public final class MCPSessionIdentity {
@@ -40,4 +38,11 @@ public final class MCPSessionIdentity {
private final Map<String, String> attributes;
+ public MCPSessionIdentity(final String sessionId, final String subject, final String source, final Map<String, String> attributes) {
+ this.sessionId = sessionId;
+ this.subject = subject;
+ this.source = source;
+ this.attributes = Map.copyOf(attributes);
+ }
+
}
diff --git a/mcp/api/src/test/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentityTest.java b/mcp/api/src/test/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentityTest.java
index 53f65c79b40..30bdc8332cc 100644
--- a/mcp/api/src/test/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentityTest.java
+++ b/mcp/api/src/test/java/org/apache/shardingsphere/mcp/api/session/MCPSessionIdentityTest.java
@@ -19,16 +19,26 @@ package org.apache.shardingsphere.mcp.api.session;
import org.junit.jupiter.api.Test;
+import java.util.HashMap;
import java.util.Map;
import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.sameInstance;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertThrows;
class MCPSessionIdentityTest {
@Test
- void assertAttributes() {
- Map<String, String> attributes = Map.of("region", "ap-south");
- assertThat(new MCPSessionIdentity("session-1", "subject", "gateway", attributes).getAttributes(), sameInstance(attributes));
+ void assertAttributesSnapshot() {
+ Map<String, String> attributes = new HashMap<>(Map.of("region", "ap-south"));
+ MCPSessionIdentity identity = new MCPSessionIdentity("session-1", "subject", "gateway", attributes);
+ attributes.put("region", "eu-west");
+ assertThat(identity.getAttributes(), is(Map.of("region", "ap-south")));
+ }
+
+ @Test
+ void assertAttributesImmutable() {
+ Map<String, String> actual = new MCPSessionIdentity("session-1", "subject", "gateway", Map.of("region", "ap-south")).getAttributes();
+ assertThrows(UnsupportedOperationException.class, () -> actual.put("region", "eu-west"));
}
}
diff --git a/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServlet.java b/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServlet.java
index b763d5aa3a4..a3b8229baa5 100644
--- a/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServlet.java
+++ b/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServlet.java
@@ -112,17 +112,14 @@ final class StreamableHttpMCPServlet extends HttpServlet implements McpStreamabl
}
private McpSchema.InitializeRequest negotiateInitializeRequest(final McpSchema.InitializeRequest initializeRequest) {
- String negotiatedProtocolVersion = negotiateProtocolVersion(initializeRequest.protocolVersion());
+ String requestedProtocolVersion = Objects.toString(initializeRequest.protocolVersion(), "").trim();
+ ShardingSpherePreconditions.checkNotEmpty(requestedProtocolVersion, () -> new IllegalArgumentException("protocolVersion is required."));
+ String negotiatedProtocolVersion = protocolVersions().contains(requestedProtocolVersion) ? requestedProtocolVersion : MCPTransportConstants.PROTOCOL_VERSION;
return negotiatedProtocolVersion.equals(initializeRequest.protocolVersion())
? initializeRequest
: new McpSchema.InitializeRequest(negotiatedProtocolVersion, initializeRequest.capabilities(), initializeRequest.clientInfo(), initializeRequest.meta());
}
- private String negotiateProtocolVersion(final String requestedProtocolVersion) {
- String actualRequestedProtocolVersion = Objects.toString(requestedProtocolVersion, "").trim();
- return protocolVersions().contains(actualRequestedProtocolVersion) ? actualRequestedProtocolVersion : MCPTransportConstants.PROTOCOL_VERSION;
- }
-
@Override
public Mono<Void> notifyClients(final String method, final Object params) {
return delegate.notifyClients(method, params);
@@ -152,6 +149,11 @@ final class StreamableHttpMCPServlet extends HttpServlet implements McpStreamabl
if (!validateTransportSecurity(request, response)) {
return;
}
+ String sessionId = Objects.toString(request.getHeader(SESSION_HEADER), "");
+ if (!sessionId.isBlank() && !sessionManager.hasSession(sessionId)) {
+ response.sendError(HttpServletResponse.SC_NOT_FOUND);
+ return;
+ }
if (!isEventStreamAccepted(request)) {
response.sendError(HttpServletResponse.SC_BAD_REQUEST, "Accept must include text/event-stream.");
return;
diff --git a/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidator.java b/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidator.java
index 4af739439d8..cc4f6309d6c 100644
--- a/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidator.java
+++ b/mcp/bootstrap/src/main/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidator.java
@@ -56,11 +56,8 @@ public final class ShardingSphereServerTransportSecurityValidator implements Ser
if ((each instanceof SessionRequiredTransportHeaderConstraint) != sessionRequired) {
continue;
}
- if (sessionRequired) {
- String sessionId = getSessionId(headers);
- if (sessionId.isBlank() || !sessionManager.hasSession(sessionId)) {
- continue;
- }
+ if (sessionRequired && getSessionId(headers).isBlank()) {
+ continue;
}
each.validate(getFirstHeaderValue(headers, each.getConstraintKey()));
}
diff --git a/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServletTest.java b/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServletTest.java
index 987ea7ffde9..548b6b42173 100644
--- a/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServletTest.java
+++ b/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/StreamableHttpMCPServletTest.java
@@ -142,12 +142,30 @@ class StreamableHttpMCPServletTest {
private static Stream<Arguments> unsupportedProtocolVersions() {
return Stream.of(
- Arguments.of("null protocol version", null),
- Arguments.of("blank protocol version", " "),
Arguments.of("legacy protocol version", ProtocolVersions.MCP_2025_06_18),
Arguments.of("unsupported protocol version", "2024-11-05"));
}
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("missingProtocolVersions")
+ void assertSetSessionFactoryRejectMissingProtocolVersion(final String name, final String requestedProtocolVersion) {
+ McpStreamableServerSession.Factory sessionFactory = mock(McpStreamableServerSession.Factory.class);
+ StreamableHttpMCPServlet actual = createServlet(mock(MCPSessionManager.class));
+ actual.setSessionFactory(sessionFactory);
+ ArgumentCaptor<McpStreamableServerSession.Factory> actualFactory = ArgumentCaptor.forClass(McpStreamableServerSession.Factory.class);
+ verify(getDelegate()).setSessionFactory(actualFactory.capture());
+ McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(requestedProtocolVersion,
+ new McpSchema.ClientCapabilities(Map.of(), null, null, null), new McpSchema.Implementation("foo_client", "1.0.0"));
+ assertThrows(IllegalArgumentException.class, () -> actualFactory.getValue().startSession(initializeRequest));
+ verify(sessionFactory, never()).startSession(any());
+ }
+
+ private static Stream<Arguments> missingProtocolVersions() {
+ return Stream.of(
+ Arguments.of("null protocol version", null),
+ Arguments.of("blank protocol version", " "));
+ }
+
@Test
void assertNotifyClients() {
StreamableHttpMCPServlet actual = createServlet(mock(MCPSessionManager.class));
@@ -217,11 +235,14 @@ class StreamableHttpMCPServletTest {
@Test
void assertRejectUnsupportedEventStreamGet() throws ServletException, IOException {
+ MCPSessionManager sessionManager = new MCPSessionManager(Map.of());
+ sessionManager.createSession(new MCPSessionIdentity("session-id", "", "", Map.of()));
HttpServletRequest request = mock(HttpServletRequest.class);
when(request.getMethod()).thenReturn("GET");
when(request.getHeader(HttpHeaders.ACCEPT)).thenReturn("text/event-stream");
+ when(request.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn("session-id");
HttpServletResponse response = mock(HttpServletResponse.class);
- StreamableHttpMCPServlet actual = createServlet(mock(MCPSessionManager.class));
+ StreamableHttpMCPServlet actual = createServlet(sessionManager);
actual.service(request, response);
verify(request).setCharacterEncoding("UTF-8");
verify(response).setCharacterEncoding("UTF-8");
@@ -229,6 +250,18 @@ class StreamableHttpMCPServletTest {
verify(getDelegate(), never()).service(any(HttpServletRequest.class), any(HttpServletResponse.class));
}
+ @Test
+ void assertRejectEventStreamGetWithUnknownSession() throws ServletException, IOException {
+ HttpServletRequest request = mock(HttpServletRequest.class);
+ when(request.getMethod()).thenReturn("GET");
+ when(request.getHeader(HttpHeaders.ACCEPT)).thenReturn("text/event-stream");
+ when(request.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn("unknown-session");
+ HttpServletResponse response = mock(HttpServletResponse.class);
+ createServlet(new MCPSessionManager(Map.of())).service(request, response);
+ verify(response).sendError(HttpServletResponse.SC_NOT_FOUND);
+ verify(getDelegate(), never()).service(any(HttpServletRequest.class), any(HttpServletResponse.class));
+ }
+
@Test
void assertServicePostWithSessionHeaderSet() throws ServletException, IOException {
MCPSessionManager sessionManager = new MCPSessionManager(Map.of());
diff --git a/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidatorTest.java b/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidatorTest.java
index 35277fa85d3..10cca9d58d5 100644
--- a/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidatorTest.java
+++ b/mcp/bootstrap/src/test/java/org/apache/shardingsphere/mcp/bootstrap/transport/server/http/validator/ShardingSphereServerTransportSecurityValidatorTest.java
@@ -17,11 +17,13 @@
package org.apache.shardingsphere.mcp.bootstrap.transport.server.http.validator;
+import io.modelcontextprotocol.server.transport.ServerTransportSecurityException;
import io.modelcontextprotocol.spec.HttpHeaders;
import org.apache.shardingsphere.mcp.api.session.MCPSessionIdentity;
import org.apache.shardingsphere.mcp.bootstrap.config.SessionAttributionSourceConfiguration;
import org.apache.shardingsphere.mcp.bootstrap.transport.server.http.SessionAttributionResolver;
import org.apache.shardingsphere.mcp.bootstrap.transport.server.http.validator.constraint.OriginHeaderConstraint;
+import org.apache.shardingsphere.mcp.bootstrap.transport.server.http.validator.constraint.SessionRequiredTransportHeaderConstraint;
import org.apache.shardingsphere.mcp.core.session.MCPSessionManager;
import org.junit.jupiter.api.Test;
@@ -32,6 +34,9 @@ import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.is;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
class ShardingSphereServerTransportSecurityValidatorTest {
@@ -109,6 +114,16 @@ class ShardingSphereServerTransportSecurityValidatorTest {
Map.of(HttpHeaders.MCP_SESSION_ID, List.of("session-1"), "X-Test-Subject", List.of("subject"), "X-Test-Source", List.of("gateway"))));
}
+ @Test
+ void assertValidateSessionRequiredConstraintWithMissingSession() throws ServerTransportSecurityException {
+ SessionRequiredTransportHeaderConstraint constraint = mock(SessionRequiredTransportHeaderConstraint.class);
+ when(constraint.getConstraintKey()).thenReturn(HttpHeaders.PROTOCOL_VERSION);
+ ShardingSphereServerTransportSecurityValidator actual = new ShardingSphereServerTransportSecurityValidator(new MCPSessionManager(Map.of()),
+ List.of(constraint), new SessionAttributionResolver(null));
+ assertDoesNotThrow(() -> actual.validateHeaders(Map.of(HttpHeaders.MCP_SESSION_ID, List.of("session-1"), HttpHeaders.PROTOCOL_VERSION, List.of("version"))));
+ verify(constraint).validate("version");
+ }
+
@Test
void assertValidateOriginBeforeSessionIdentity() {
MCPSessionManager sessionManager = new MCPSessionManager(Map.of());
diff --git a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/ExplainSQLCandidateValidator.java b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/ExplainSQLCandidateValidator.java
index 55a4644c081..09984f5d936 100644
--- a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/ExplainSQLCandidateValidator.java
+++ b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/ExplainSQLCandidateValidator.java
@@ -39,30 +39,29 @@ final class ExplainSQLCandidateValidator {
private final MCPStatementAnalyzer statementAnalyzer;
- private final SQLStatementScanner scanner = new SQLStatementScanner();
-
ClassificationResult validate(final String sql, final String explainSql, final MCPDatabaseCapability databaseCapability) {
ClassificationResult explainedStatement = statementAnalyzer.analyze(sql, databaseCapability);
ShardingSpherePreconditions.checkState(SupportedMCPStatement.QUERY == explainedStatement.getStatementClass(),
() -> new MCPInvalidRequestException("database_gateway_execute_explain_query only supports QUERY statements as the explained SQL."));
+ SQLStatementScanner scanner = new SQLStatementScanner(databaseCapability.getDatabaseType());
String actualExplainSql = scanner.normalizeSingleStatement(explainSql);
ShardingSpherePreconditions.checkState(!scanner.containsExecutableComment(actualExplainSql),
() -> new MCPInvalidRequestException("Executable comments are not supported by the MCP explain query tool."));
List<SQLStatementToken> tokens = scanner.tokenize(actualExplainSql);
- checkExplainCandidate(tokens, explainedStatement.getNormalizedSql(), actualExplainSql);
+ checkExplainCandidate(scanner, tokens, explainedStatement.getNormalizedSql(), actualExplainSql);
return new ClassificationResult(SupportedMCPStatement.EXPLAIN, "EXPLAIN", actualExplainSql, "", explainedStatement.getReferencedObjects(), false);
}
- private void checkExplainCandidate(final List<SQLStatementToken> tokens, final String sql, final String explainSql) {
+ private void checkExplainCandidate(final SQLStatementScanner scanner, final List<SQLStatementToken> tokens, final String sql, final String explainSql) {
ShardingSpherePreconditions.checkState(!tokens.isEmpty() && scanner.isKeyword(tokens.get(0), "EXPLAIN"),
() -> new MCPInvalidRequestException("explain_sql must start with EXPLAIN."));
String explainPrefix = extractExplainPrefix(explainSql, sql);
List<SQLStatementToken> explainPrefixTokens = scanner.tokenize(explainPrefix);
- ShardingSpherePreconditions.checkState(!containsKeywordSequence(explainPrefixTokens, "EXPLAIN", "PLAN", "FOR"),
+ ShardingSpherePreconditions.checkState(!containsKeywordSequence(scanner, explainPrefixTokens, "EXPLAIN", "PLAN", "FOR"),
() -> new MCPInvalidRequestException("EXPLAIN PLAN FOR workflows are not supported by the MCP explain query tool."));
- ShardingSpherePreconditions.checkState(!containsKeyword(explainPrefixTokens, "ANALYZE", "ANALYSE"),
+ ShardingSpherePreconditions.checkState(!containsKeyword(scanner, explainPrefixTokens, "ANALYZE", "ANALYSE"),
() -> new MCPInvalidRequestException("EXPLAIN ANALYZE is not supported by the MCP explain query tool."));
- ShardingSpherePreconditions.checkState(!containsKeyword(explainPrefixTokens, "INTO"),
+ ShardingSpherePreconditions.checkState(!containsKeyword(scanner, explainPrefixTokens, "INTO"),
() -> new MCPInvalidRequestException("EXPLAIN output redirection is not supported by the MCP explain query tool."));
ShardingSpherePreconditions.checkState(!containsStatementStartKeyword(explainPrefixTokens),
() -> new MCPInvalidRequestException("explain_sql must not wrap the original sql argument in another statement."));
@@ -83,7 +82,7 @@ final class ExplainSQLCandidateValidator {
return false;
}
- private boolean containsKeyword(final List<SQLStatementToken> tokens, final String... keywords) {
+ private boolean containsKeyword(final SQLStatementScanner scanner, final List<SQLStatementToken> tokens, final String... keywords) {
for (SQLStatementToken each : tokens) {
if (scanner.isKeyword(each, keywords)) {
return true;
@@ -92,16 +91,16 @@ final class ExplainSQLCandidateValidator {
return false;
}
- private boolean containsKeywordSequence(final List<SQLStatementToken> tokens, final String... keywords) {
+ private boolean containsKeywordSequence(final SQLStatementScanner scanner, final List<SQLStatementToken> tokens, final String... keywords) {
for (int index = 0; index + keywords.length <= tokens.size(); index++) {
- if (containsKeywordSequence(tokens, index, keywords)) {
+ if (containsKeywordSequence(scanner, tokens, index, keywords)) {
return true;
}
}
return false;
}
- private boolean containsKeywordSequence(final List<SQLStatementToken> tokens, final int startIndex, final String... keywords) {
+ private boolean containsKeywordSequence(final SQLStatementScanner scanner, final List<SQLStatementToken> tokens, final int startIndex, final String... keywords) {
for (int index = 0; index < keywords.length; index++) {
if (!scanner.isKeyword(tokens.get(startIndex + index), keywords[index])) {
return false;
diff --git a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzer.java b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzer.java
index 41a821414bc..93f4bbc05a6 100644
--- a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzer.java
+++ b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzer.java
@@ -38,6 +38,7 @@ import org.apache.shardingsphere.sql.parser.statement.core.statement.SQLStatemen
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dcl.GrantStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dcl.RevokeStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.ddl.DDLStatement;
+import org.apache.shardingsphere.sql.parser.statement.core.statement.type.ddl.table.CreateTableStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.DeleteStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.InsertStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.MergeStatement;
@@ -45,8 +46,10 @@ import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.Se
import org.apache.shardingsphere.sql.parser.statement.core.statement.type.dml.UpdateStatement;
import java.util.Collection;
+import java.util.List;
import java.util.Locale;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
@@ -56,21 +59,18 @@ final class MCPStatementAnalyzer {
private static final Set<String> DDL_STATEMENT_TYPES = Set.of("CREATE", "ALTER", "DROP", "TRUNCATE");
- private final SQLStatementScanner scanner = new SQLStatementScanner();
-
- private final SQLStatementSafetyValidator safetyValidator = new SQLStatementSafetyValidator(scanner);
-
- private final SQLStatementObjectExtractor objectExtractor = new SQLStatementObjectExtractor(scanner);
-
private final Collection<MCPRuleDistSQLStatementClassifier> ruleDistSQLStatementClassifiers =
ShardingSphereServiceLoader.getServiceInstances(MCPRuleDistSQLStatementClassifier.class);
private final Map<String, SQLParserEngine> parserEngines = new ConcurrentHashMap<>();
ClassificationResult analyze(final String sql, final MCPDatabaseCapability databaseCapability) {
+ String databaseType = databaseCapability.getDatabaseType();
+ SQLStatementScanner scanner = new SQLStatementScanner(databaseType);
String actualSql = scanner.normalizeSingleStatement(sql);
String leadingSql = actualSql.substring(scanner.skipInsignificant(actualSql, 0)).trim();
String upperLeadingSql = leadingSql.toUpperCase(Locale.ENGLISH);
+ SQLStatementSafetyValidator safetyValidator = new SQLStatementSafetyValidator(scanner);
safetyValidator.checkLeadingStatement(upperLeadingSql, actualSql);
if (isSavepointStatement(upperLeadingSql)) {
return analyzeSavepointStatement(actualSql, leadingSql, upperLeadingSql);
@@ -78,13 +78,60 @@ final class MCPStatementAnalyzer {
if (isTransactionControlStatement(upperLeadingSql)) {
return new ClassificationResult(SupportedMCPStatement.TRANSACTION_CONTROL, extractTransactionStatementType(upperLeadingSql), actualSql, "", Set.of(), false);
}
- SQLStatement sqlStatement = parse(actualSql, databaseCapability.getDatabaseType());
+ SQLStatement sqlStatement = parse(actualSql, databaseType);
safetyValidator.checkParsedStatement(sqlStatement, actualSql);
boolean ruleDistSQL = isRuleDistSQL(sqlStatement);
String leadingKeyword = scanner.extractLeadingKeyword(actualSql);
SupportedMCPStatement statementClass = resolveStatementClass(sqlStatement, leadingKeyword, ruleDistSQL);
String statementType = resolveStatementType(sqlStatement, statementClass, leadingKeyword);
- return new ClassificationResult(statementClass, statementType, actualSql, "", objectExtractor.extract(sqlStatement, actualSql), ruleDistSQL);
+ return new ClassificationResult(statementClass, statementType, actualSql, "", extractReferencedObjects(sqlStatement, actualSql, databaseType, scanner), ruleDistSQL);
+ }
+
+ private Collection<SQLStatementObjectName> extractReferencedObjects(final SQLStatement sqlStatement, final String sql, final String databaseType,
+ final SQLStatementScanner scanner) {
+ SQLStatementObjectExtractor extractor = new SQLStatementObjectExtractor(scanner);
+ Collection<SQLStatementObjectName> result = extractor.extract(sqlStatement, sql);
+ findCreateTableQuery(sqlStatement, sql, scanner).ifPresent(query -> result.addAll(extractor.extract(parse(query, databaseType), query)));
+ return result;
+ }
+
+ private Optional<String> findCreateTableQuery(final SQLStatement sqlStatement, final String sql, final SQLStatementScanner scanner) {
+ if (!(sqlStatement instanceof CreateTableStatement)) {
+ return Optional.empty();
+ }
+ CreateTableStatement createTableStatement = (CreateTableStatement) sqlStatement;
+ if (createTableStatement.getSelectStatement().isPresent()) {
+ return Optional.empty();
+ }
+ List<SQLStatementToken> tokens = scanner.tokenize(sql);
+ int tableStopIndex = createTableStatement.getTable().getStopIndex();
+ int parenthesisDepth = 0;
+ for (int index = 0; index < tokens.size(); index++) {
+ SQLStatementToken token = tokens.get(index);
+ if (token.startIndex() <= tableStopIndex) {
+ continue;
+ }
+ if (0 == parenthesisDepth) {
+ int queryStartIndex = scanner.isKeyword(token, "AS") ? index + 1 : index;
+ if (isCreateTableQueryStart(tokens, queryStartIndex, scanner)) {
+ return Optional.of(sql.substring(tokens.get(queryStartIndex).startIndex()));
+ }
+ }
+ if ("(".equals(token.text())) {
+ parenthesisDepth++;
+ } else if (")".equals(token.text())) {
+ parenthesisDepth--;
+ }
+ }
+ return Optional.empty();
+ }
+
+ private boolean isCreateTableQueryStart(final List<SQLStatementToken> tokens, final int startIndex, final SQLStatementScanner scanner) {
+ int queryKeywordIndex = startIndex;
+ while (queryKeywordIndex < tokens.size() && "(".equals(tokens.get(queryKeywordIndex).text())) {
+ queryKeywordIndex++;
+ }
+ return queryKeywordIndex < tokens.size() && scanner.isKeyword(tokens.get(queryKeywordIndex), "SELECT", "WITH", "TABLE", "VALUES");
}
private ClassificationResult analyzeSavepointStatement(final String actualSql, final String leadingSql, final String upperLeadingSql) {
diff --git a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementObjectExtractor.java b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementObjectExtractor.java
index 3b86896019f..1ab17e9f264 100644
--- a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementObjectExtractor.java
+++ b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementObjectExtractor.java
@@ -19,6 +19,8 @@ package org.apache.shardingsphere.mcp.core.tool.handler.execute;
import lombok.AccessLevel;
import lombok.RequiredArgsConstructor;
+import org.apache.shardingsphere.database.connector.core.metadata.database.enums.QuoteCharacter;
+import org.apache.shardingsphere.mcp.core.protocol.exception.MCPUnsupportedSQLStatementException;
import org.apache.shardingsphere.sql.parser.statement.core.extractor.TableExtractor;
import org.apache.shardingsphere.sql.parser.statement.core.segment.ddl.index.IndexSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.complex.CommonTableExpressionSegment;
@@ -222,8 +224,14 @@ final class SQLStatementObjectExtractor {
} else if (sqlStatement instanceof AlterSequenceStatement) {
addName(((AlterSequenceStatement) sqlStatement).getSequenceName(), result);
} else if (sqlStatement instanceof DropSequenceStatement) {
- for (String each : ((DropSequenceStatement) sqlStatement).getSequenceNames()) {
- addName(each, result);
+ for (Object each : ((DropSequenceStatement) sqlStatement).getSequenceNames()) {
+ if (each instanceof String) {
+ addName((String) each, result);
+ } else if (each instanceof SimpleTableSegment) {
+ addTable((SimpleTableSegment) each, result);
+ } else {
+ throw new MCPUnsupportedSQLStatementException();
+ }
}
}
}
@@ -310,9 +318,31 @@ final class SQLStatementObjectExtractor {
}
private void addTable(final SimpleTableSegment table, final Collection<SQLStatementObjectName> result) {
- if (null != table) {
- result.add(SQLStatementObjectName.from(table.getOwner(), table.getTableName().getIdentifier()));
+ if (null == table) {
+ return;
+ }
+ IdentifierValue identifier = table.getTableName().getIdentifier();
+ if (table.getOwner().isEmpty() && addFlattenedQualifiedTable(identifier, result)) {
+ return;
+ }
+ result.add(SQLStatementObjectName.from(table.getOwner(), identifier));
+ }
+
+ private boolean addFlattenedQualifiedTable(final IdentifierValue identifier, final Collection<SQLStatementObjectName> result) {
+ if (!identifier.getValue().contains(".")) {
+ return false;
+ }
+ QuoteCharacter quoteCharacter = identifier.getQuoteCharacter();
+ String identifierText = QuoteCharacter.NONE == quoteCharacter || quoteCharacter.isWrapped(identifier.getValue())
+ ? identifier.getValue()
+ : identifier.getValueWithQuoteCharacters();
+ List<SQLStatementToken> tokens = scanner.tokenize(identifierText);
+ int objectNameEnd = findObjectNameEnd(tokens, 0);
+ if (3 > objectNameEnd || objectNameEnd != tokens.size()) {
+ return false;
}
+ addObjectName(tokens, 0, objectNameEnd, false, result);
+ return true;
}
private void addIndex(final IndexSegment index, final Collection<SQLStatementObjectName> result) {
diff --git a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementScanner.java b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementScanner.java
index 6d0677e6fa4..7339035ca2e 100644
--- a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementScanner.java
+++ b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementScanner.java
@@ -28,6 +28,15 @@ import java.util.Locale;
final class SQLStatementScanner {
+ private final boolean mysqlCommentSyntax;
+
+ private final boolean dollarQuotedStrings;
+
+ SQLStatementScanner(final String databaseType) {
+ mysqlCommentSyntax = "MySQL".equalsIgnoreCase(databaseType) || "MariaDB".equalsIgnoreCase(databaseType);
+ dollarQuotedStrings = "PostgreSQL".equalsIgnoreCase(databaseType) || "openGauss".equalsIgnoreCase(databaseType);
+ }
+
String normalizeSingleStatement(final String sql) {
String result = sql.trim();
ShardingSpherePreconditions.checkState(!result.isEmpty(), () -> new MCPInvalidRequestException("sql cannot be empty."));
@@ -84,13 +93,17 @@ final class SQLStatementScanner {
break;
}
char currentChar = sql.charAt(currentIndex);
+ if (isDollarQuotedTextStart(sql, currentIndex)) {
+ currentIndex = skipDollarQuotedText(sql, currentIndex) + 1;
+ continue;
+ }
if ('\'' == currentChar) {
currentIndex = skipQuotedText(sql, currentIndex) + 1;
continue;
}
if (isQuotedIdentifierStart(currentChar)) {
int stopIndex = skipQuotedText(sql, currentIndex);
- result.add(new SQLStatementToken(sql.substring(currentIndex, stopIndex + 1), true));
+ result.add(new SQLStatementToken(sql.substring(currentIndex, stopIndex + 1), true, currentIndex));
currentIndex = stopIndex + 1;
continue;
}
@@ -99,12 +112,12 @@ final class SQLStatementScanner {
while (stopIndex < sql.length() && isIdentifierCharacter(sql.charAt(stopIndex))) {
stopIndex++;
}
- result.add(new SQLStatementToken(sql.substring(currentIndex, stopIndex), false));
+ result.add(new SQLStatementToken(sql.substring(currentIndex, stopIndex), false, currentIndex));
currentIndex = stopIndex;
continue;
}
if ('.' == currentChar || '(' == currentChar || ')' == currentChar || ',' == currentChar || '*' == currentChar) {
- result.add(new SQLStatementToken(String.valueOf(currentChar), false));
+ result.add(new SQLStatementToken(String.valueOf(currentChar), false, currentIndex));
}
currentIndex++;
}
@@ -138,6 +151,10 @@ final class SQLStatementScanner {
currentIndex = skipBlockComment(sql, currentIndex) + 1;
continue;
}
+ if (isDollarQuotedTextStart(sql, currentIndex)) {
+ currentIndex = skipDollarQuotedText(sql, currentIndex) + 1;
+ continue;
+ }
if (isQuotedTextStart(sql.charAt(currentIndex))) {
currentIndex = skipQuotedText(sql, currentIndex) + 1;
continue;
@@ -165,6 +182,10 @@ final class SQLStatementScanner {
currentIndex = skipBlockComment(sql, currentIndex) + 1;
continue;
}
+ if (isDollarQuotedTextStart(sql, currentIndex)) {
+ currentIndex = skipDollarQuotedText(sql, currentIndex) + 1;
+ continue;
+ }
if (isQuotedTextStart(sql.charAt(currentIndex))) {
currentIndex = skipQuotedText(sql, currentIndex) + 1;
continue;
@@ -204,7 +225,6 @@ final class SQLStatementScanner {
private int findStatementDelimiter(final String sql) {
int result = 0;
while (result < sql.length()) {
- char currentChar = sql.charAt(result);
if (isLineCommentStart(sql, result)) {
result = skipLineComment(sql, result) + 1;
continue;
@@ -213,6 +233,11 @@ final class SQLStatementScanner {
result = skipBlockComment(sql, result) + 1;
continue;
}
+ if (isDollarQuotedTextStart(sql, result)) {
+ result = skipDollarQuotedText(sql, result) + 1;
+ continue;
+ }
+ char currentChar = sql.charAt(result);
if (isQuotedTextStart(currentChar)) {
result = skipQuotedText(sql, result) + 1;
continue;
@@ -254,7 +279,14 @@ final class SQLStatementScanner {
}
private boolean isLineCommentStart(final String sql, final int startIndex) {
- return startIndex + 1 < sql.length() && '-' == sql.charAt(startIndex) && '-' == sql.charAt(startIndex + 1);
+ if (mysqlCommentSyntax && '#' == sql.charAt(startIndex)) {
+ return true;
+ }
+ if (startIndex + 1 >= sql.length() || '-' != sql.charAt(startIndex) || '-' != sql.charAt(startIndex + 1)) {
+ return false;
+ }
+ return !mysqlCommentSyntax || startIndex + 2 >= sql.length()
+ || Character.isWhitespace(sql.charAt(startIndex + 2)) || Character.isISOControl(sql.charAt(startIndex + 2));
}
private boolean isBlockCommentStart(final String sql, final int startIndex) {
@@ -262,7 +294,7 @@ final class SQLStatementScanner {
}
private int skipLineComment(final String sql, final int startIndex) {
- int result = startIndex + 2;
+ int result = startIndex + ('#' == sql.charAt(startIndex) ? 1 : 2);
while (result < sql.length() && '\n' != sql.charAt(result)) {
result++;
}
@@ -279,4 +311,43 @@ final class SQLStatementScanner {
}
throw new MCPUnsupportedSQLStatementException();
}
+
+ private boolean isDollarQuotedTextStart(final String sql, final int startIndex) {
+ return dollarQuotedStrings && (0 == startIndex || !isIdentifierCharacter(sql.charAt(startIndex - 1))) && !getDollarQuoteDelimiter(sql, startIndex).isEmpty();
+ }
+
+ private int skipDollarQuotedText(final String sql, final int startIndex) {
+ String delimiter = getDollarQuoteDelimiter(sql, startIndex);
+ int closingDelimiterIndex = sql.indexOf(delimiter, startIndex + delimiter.length());
+ if (-1 == closingDelimiterIndex) {
+ throw new MCPUnsupportedSQLStatementException();
+ }
+ return closingDelimiterIndex + delimiter.length() - 1;
+ }
+
+ private String getDollarQuoteDelimiter(final String sql, final int startIndex) {
+ if (startIndex >= sql.length() || '$' != sql.charAt(startIndex) || startIndex + 1 >= sql.length()) {
+ return "";
+ }
+ int result = startIndex + 1;
+ if ('$' == sql.charAt(result)) {
+ return "$$";
+ }
+ if (!isDollarQuoteTagStart(sql.charAt(result))) {
+ return "";
+ }
+ result++;
+ while (result < sql.length() && isDollarQuoteTagCharacter(sql.charAt(result))) {
+ result++;
+ }
+ return result < sql.length() && '$' == sql.charAt(result) ? sql.substring(startIndex, result + 1) : "";
+ }
+
+ private boolean isDollarQuoteTagStart(final char value) {
+ return Character.isLetter(value) || '_' == value;
+ }
+
+ private boolean isDollarQuoteTagCharacter(final char value) {
+ return Character.isLetterOrDigit(value) || '_' == value;
+ }
}
diff --git a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementToken.java b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementToken.java
index 918b38ef0a5..f5b2256faa5 100644
--- a/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementToken.java
+++ b/mcp/core/src/main/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/SQLStatementToken.java
@@ -27,10 +27,13 @@ final class SQLStatementToken {
private final boolean quotedIdentifier;
- SQLStatementToken(final String text, final boolean quotedIdentifier) {
+ private final int startIndex;
+
+ SQLStatementToken(final String text, final boolean quotedIdentifier, final int startIndex) {
this.text = text;
upperText = text.toUpperCase(Locale.ENGLISH);
this.quotedIdentifier = quotedIdentifier;
+ this.startIndex = startIndex;
}
String text() {
@@ -45,6 +48,10 @@ final class SQLStatementToken {
return quotedIdentifier;
}
+ int startIndex() {
+ return startIndex;
+ }
+
boolean identifier() {
return quotedIdentifier || !text.isEmpty() && isIdentifierStart(text.charAt(0));
}
diff --git a/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPSQLExecutionFacadeTest.java b/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPSQLExecutionFacadeTest.java
index ce48088b4f2..b45988a96ba 100644
--- a/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPSQLExecutionFacadeTest.java
+++ b/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPSQLExecutionFacadeTest.java
@@ -509,6 +509,11 @@ class MCPSQLExecutionFacadeTest {
Arguments.of("delete target before from", "DELETE other_db.orders FROM logic_db.orders JOIN other_db.items ON 1 = 1", "other_db.orders", "DML"),
Arguments.of("create view object list", "CREATE VIEW logic_db.active_orders AS SELECT * FROM logic_db.orders, other_db.items", "other_db.items", "DDL"),
Arguments.of("create table like", "CREATE TABLE logic_db.orders_archive LIKE other_db.orders", "other_db.orders", "DDL"),
+ Arguments.of("create table select without as", "CREATE TABLE logic_db.orders_archive SELECT * FROM other_db.orders", "other_db.orders", "DDL"),
+ Arguments.of("create table parenthesized select", "CREATE TABLE logic_db.orders_archive AS (SELECT * FROM other_db.orders)", "other_db.orders", "DDL"),
+ Arguments.of("create table table query", "CREATE TABLE logic_db.orders_archive AS TABLE other_db.orders", "other_db.orders", "DDL"),
+ Arguments.of("create table quoted table query", "CREATE TABLE logic_db.orders_archive AS TABLE `other_db`.`orders`", "other_db.orders", "DDL"),
+ Arguments.of("create table ANSI quoted table query", "CREATE TABLE logic_db.orders_archive AS TABLE \"other_db\".\"orders\"", "other_db.orders", "DDL"),
Arguments.of("create table foreign key reference", "CREATE TABLE logic_db.order_items (order_id INT REFERENCES other_db.orders(id))", "other_db.orders", "DDL"),
Arguments.of("alter table", "ALTER TABLE other_db.orders ADD COLUMN status VARCHAR(10)", "other_db.orders", "DDL"),
Arguments.of("alter table foreign key reference",
diff --git a/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzerTest.java b/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzerTest.java
index f8cc7979e59..76842bb9220 100644
--- a/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzerTest.java
+++ b/mcp/core/src/test/java/org/apache/shardingsphere/mcp/core/tool/handler/execute/MCPStatementAnalyzerTest.java
@@ -30,6 +30,8 @@ import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
+import java.util.ArrayList;
+import java.util.List;
import java.util.stream.Stream;
import static org.hamcrest.MatcherAssert.assertThat;
@@ -88,6 +90,12 @@ class MCPStatementAnalyzerTest {
assertThat(actual.getReferencedObjectNames(), contains("logic_db.orders", "other_db.order_items"));
}
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("referencedObjectCases")
+ void assertAnalyzeExactReferencedObjects(final String name, final String databaseType, final String sql, final List<String> expected) {
+ assertThat(new ArrayList<>(analyzer.analyze(sql, createCapability(databaseType)).getReferencedObjectNames()), is(expected));
+ }
+
@Test
void assertAnalyzeCommonTableExpressionReferences() {
ClassificationResult actual = analyzer.analyze(
@@ -117,6 +125,13 @@ class MCPStatementAnalyzerTest {
assertThat(actual.getTargetObjectName().orElse(""), is("orders"));
}
+ @Test
+ void assertAnalyzeQuotedTableQueryIdentifier() {
+ ClassificationResult actual = analyzer.analyze("CREATE TABLE order_archive AS TABLE `other.db.orders`", createCapability("MySQL"));
+ assertThat(actual.getReferencedObjectNames(), contains("order_archive", "other.db.orders"));
+ assertFalse(actual.getReferencedObjects().stream().filter(each -> "other.db.orders".equals(each.getObjectName())).findFirst().orElseThrow().isQualified());
+ }
+
@Test
void assertAnalyzeDCLReference() {
ClassificationResult actual = analyzer.analyze("GRANT SELECT ON other_db.orders TO PUBLIC", createCapability("MySQL"));
@@ -184,10 +199,49 @@ class MCPStatementAnalyzerTest {
Arguments.of("schema", "GRANT USAGE ON SCHEMA other_db TO PUBLIC"));
}
+ private static Stream<Arguments> referencedObjectCases() {
+ return Stream.of(
+ Arguments.of("merge target and source", "Oracle", "MERGE INTO order_archive USING orders ON (order_archive.order_id = orders.order_id)",
+ List.of("order_archive", "orders")),
+ Arguments.of("create table like target and source", "MySQL", "CREATE TABLE order_archive LIKE orders", List.of("order_archive", "orders")),
+ Arguments.of("create table select target and source", "MySQL", "CREATE TABLE order_archive (order_id BIGINT) AS SELECT * FROM other_db.orders",
+ List.of("order_archive", "other_db.orders")),
+ Arguments.of("create table select without as", "MySQL", "CREATE TABLE order_archive SELECT * FROM other_db.orders",
+ List.of("order_archive", "other_db.orders")),
+ Arguments.of("create table with query without as", "MySQL",
+ "CREATE TABLE order_archive WITH order_source AS (SELECT * FROM other_db.orders) SELECT * FROM order_source",
+ List.of("order_archive", "other_db.orders")),
+ Arguments.of("create table parenthesized select", "MySQL", "CREATE TABLE order_archive AS (SELECT * FROM other_db.orders)",
+ List.of("order_archive", "other_db.orders")),
+ Arguments.of("create table table query", "MySQL", "CREATE TABLE order_archive AS TABLE other_db.orders",
+ List.of("order_archive", "other_db.orders")),
+ Arguments.of("create table partition values", "MySQL",
+ "CREATE TABLE order_archive (order_id INT) PARTITION BY RANGE (order_id) (PARTITION p0 VALUES LESS THAN (10))", List.of("order_archive")),
+ Arguments.of("create index and table", "PostgreSQL", "CREATE INDEX order_idx ON orders (status)", List.of("order_idx", "orders")),
+ Arguments.of("create database", "MySQL", "CREATE DATABASE order_archive", List.of("order_archive")),
+ Arguments.of("drop database", "MySQL", "DROP DATABASE order_archive", List.of("order_archive")),
+ Arguments.of("create schema", "PostgreSQL", "CREATE SCHEMA order_archive", List.of("order_archive")),
+ Arguments.of("alter schema rename", "PostgreSQL", "ALTER SCHEMA order_archive RENAME TO order_history", List.of("order_archive", "order_history")),
+ Arguments.of("drop schema", "PostgreSQL", "DROP SCHEMA order_archive", List.of("order_archive")),
+ Arguments.of("create sequence", "PostgreSQL", "CREATE SEQUENCE order_seq", List.of("order_seq")),
+ Arguments.of("alter sequence", "PostgreSQL", "ALTER SEQUENCE order_seq INCREMENT BY 2", List.of("order_seq")),
+ Arguments.of("drop sequence", "PostgreSQL", "DROP SEQUENCE order_seq", List.of("order_seq")));
+ }
+
private static Stream<Arguments> statementCases() {
return Stream.of(
Arguments.of("trim trailing semicolon", "MySQL", " SELECT * FROM orders ; ", SupportedMCPStatement.QUERY, "SELECT", "SELECT * FROM orders", "orders", ""),
Arguments.of("semicolon literal", "MySQL", "SELECT ';' AS literal_value", SupportedMCPStatement.QUERY, "SELECT", "SELECT ';' AS literal_value", "", ""),
+ Arguments.of("mysql hash comment", "MySQL", "SELECT 1 # @order_status := 1", SupportedMCPStatement.QUERY, "SELECT",
+ "SELECT 1 # @order_status := 1", "", ""),
+ Arguments.of("postgresql adjacent dash comment", "PostgreSQL", "SELECT 1--@order_status:=1", SupportedMCPStatement.QUERY, "SELECT",
+ "SELECT 1--@order_status:=1", "", ""),
+ Arguments.of("postgresql dollar quoted literal", "PostgreSQL", "SELECT $$@order_status:=1; /*!80018 ANALYZE */$$", SupportedMCPStatement.QUERY, "SELECT",
+ "SELECT $$@order_status:=1; /*!80018 ANALYZE */$$", "", ""),
+ Arguments.of("opengauss tagged dollar quoted literal", "openGauss", "SELECT $tag$a;b$tag$", SupportedMCPStatement.QUERY, "SELECT",
+ "SELECT $tag$a;b$tag$", "", ""),
+ Arguments.of("user variable assignment string", "MySQL", "SELECT '@order_status := 1'", SupportedMCPStatement.QUERY, "SELECT",
+ "SELECT '@order_status := 1'", "", ""),
Arguments.of("with query", "PostgreSQL", "WITH order_result AS (SELECT * FROM orders) SELECT * FROM order_result",
SupportedMCPStatement.QUERY, "SELECT", "WITH order_result AS (SELECT * FROM orders) SELECT * FROM order_result", "orders", ""),
Arguments.of("with update", "PostgreSQL", "WITH order_result AS (SELECT * FROM orders) UPDATE order_archive SET status = 'DONE' FROM order_result",
@@ -226,6 +280,12 @@ class MCPStatementAnalyzerTest {
Arguments.of("blank", "MySQL", " ", MCPInvalidRequestException.class, "sql cannot be empty."),
Arguments.of("delimiter only", "MySQL", ";", MCPInvalidRequestException.class, "sql cannot be empty."),
Arguments.of("multiple statements", "MySQL", "SELECT 1; SELECT 2", MCPMultipleSQLStatementsException.class, "Only one SQL statement is allowed."),
+ Arguments.of("multiple statements after dollar quoted literal", "PostgreSQL", "SELECT $$a;b$$; SELECT 2", MCPMultipleSQLStatementsException.class,
+ "Only one SQL statement is allowed."),
+ Arguments.of("unterminated quote", "MySQL", "SELECT 'value", MCPUnsupportedSQLStatementException.class,
+ "Statement is not supported by the MCP contract."),
+ Arguments.of("unterminated dollar quote", "PostgreSQL", "SELECT $tag$value", MCPUnsupportedSQLStatementException.class,
+ "Statement is not supported by the MCP contract."),
Arguments.of("parser unavailable", "H2", "SELECT 1", MCPUnsupportedException.class, "SQL parser is not available for database type `H2`."),
Arguments.of("savepoint without name", "MySQL", "SAVEPOINT", MCPInvalidRequestException.class, "Savepoint name is required."),
Arguments.of("savepoint with extra token", "MySQL", "SAVEPOINT order_sp extra", MCPInvalidRequestException.class, "Savepoint name is required."),
@@ -272,6 +332,12 @@ class MCPStatementAnalyzerTest {
"Statement is banned by the MCP contract."),
Arguments.of("banned user variable assignment", "MySQL", "SELECT @order_status := status FROM orders", MCPBannedSQLStatementException.class,
"Statement is banned by the MCP contract."),
+ Arguments.of("banned quoted user variable assignment", "MySQL", "SELECT @'order-status' := status FROM orders", MCPBannedSQLStatementException.class,
+ "Statement is banned by the MCP contract."),
+ Arguments.of("banned mysql adjacent dash assignment", "MySQL", "SELECT 1--@order_status:=1", MCPBannedSQLStatementException.class,
+ "Statement is banned by the MCP contract."),
+ Arguments.of("banned mariadb adjacent dash assignment", "MariaDB", "SELECT 1--@order_status:=1", MCPBannedSQLStatementException.class,
+ "Statement is banned by the MCP contract."),
Arguments.of("banned last insert id mutation", "MySQL", "SELECT LAST_INSERT_ID(1)", MCPBannedSQLStatementException.class,
"Statement is banned by the MCP contract."),
Arguments.of("banned metadata lookup", "PostgreSQL", "SELECT to_regclass('orders')", MCPBannedSQLStatementException.class,
diff --git a/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/descriptor/MCPToolDescriptorCatalogValidator.java b/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/descriptor/MCPToolDescriptorCatalogValidator.java
index eca5789b56c..e8e9135cf97 100644
--- a/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/descriptor/MCPToolDescriptorCatalogValidator.java
+++ b/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/descriptor/MCPToolDescriptorCatalogValidator.java
@@ -80,7 +80,7 @@ public final class MCPToolDescriptorCatalogValidator {
MCPToolOutputSchemaValidator.validate(each);
validateToolDescriptorContract(each, runtimes.get(each.getName()));
validateDestructiveToolDescriptor(each, runtimes.get(each.getName()));
- validatePlanningExecutionMode(each);
+ validateNonDestructiveExecutionMode(each);
validateRelatedResourceUris(each, resourceIdentifiers, shardingSphereResourceIdentifiers);
}
validatePlanningToolRuntimeDescriptors(registered, runtimeDescriptors);
@@ -346,7 +346,7 @@ public final class MCPToolDescriptorCatalogValidator {
() -> new IllegalStateException(String.format("Destructive tool `%s` must declare sideEffectScope in internal runtime.", descriptor.getName())));
}
- private static void validatePlanningExecutionMode(final MCPToolDescriptor descriptor) {
+ private static void validateNonDestructiveExecutionMode(final MCPToolDescriptor descriptor) {
if (descriptor.getAnnotations().isDestructiveHint()) {
return;
}
@@ -356,9 +356,9 @@ public final class MCPToolDescriptorCatalogValidator {
}
Collection<?> executionModes = executionMode.get().get("enum") instanceof Collection ? (Collection<?>) executionMode.get().get("enum") : List.of();
ShardingSpherePreconditions.checkState(!executionModes.contains("preview"),
- () -> new IllegalStateException(String.format("Planning tool `%s` execution_mode must not expose preview.", descriptor.getName())));
+ () -> new IllegalStateException(String.format("Non-destructive tool `%s` execution_mode must not expose preview.", descriptor.getName())));
ShardingSpherePreconditions.checkState(!executionModes.contains("auto-execute"),
- () -> new IllegalStateException(String.format("Planning tool `%s` execution_mode must not expose auto-execute.", descriptor.getName())));
+ () -> new IllegalStateException(String.format("Non-destructive tool `%s` execution_mode must not expose auto-execute.", descriptor.getName())));
}
}
diff --git a/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/workflow/service/WorkflowGuidancePayloadBuilder.java b/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/workflow/service/WorkflowGuidancePayloadBuilder.java
index c797f71ff0a..d79375b0627 100644
--- a/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/workflow/service/WorkflowGuidancePayloadBuilder.java
+++ b/mcp/support/src/main/java/org/apache/shardingsphere/mcp/support/workflow/service/WorkflowGuidancePayloadBuilder.java
@@ -93,7 +93,7 @@ public final class WorkflowGuidancePayloadBuilder {
} else if (WorkflowLifecycle.STATUS_FAILED.equals(status)) {
nextActions.add(createUserAction("Inspect issues and retry database_gateway_apply_workflow only after the failed artifact is corrected.", List.of("issues")));
}
- payload.put(MCPPayloadFieldNames.NEXT_ACTIONS, addSequencing(nextActions));
+ payload.put(MCPPayloadFieldNames.NEXT_ACTIONS, nextActions);
}
private static boolean isSecretReferenceRecovery(final Map<String, Object> payload) {
@@ -312,16 +312,6 @@ public final class WorkflowGuidancePayloadBuilder {
return MCPNextActionUtils.stop("Validation passed. Report the confirmed workflow result to the user.");
}
- private static List<Map<String, Object>> addSequencing(final List<Map<String, Object>> nextActions) {
- for (int index = 0; index < nextActions.size(); index++) {
- nextActions.get(index).put("order", index + 1);
- if (0 < index && "ask_user".equals(nextActions.get(index - 1).get("type"))) {
- nextActions.get(index).put("depends_on", List.of(index));
- }
- }
- return nextActions;
- }
-
private static String resolvePlanningTool(final WorkflowContextSnapshot snapshot) {
return MCPDescriptorCatalogIndex.findPlanningToolNameByWorkflowKind(resolveWorkflowKind(snapshot)).orElse("");
}
diff --git a/mcp/support/src/test/java/org/apache/shardingsphere/mcp/support/descriptor/MCPDescriptorCatalogValidatorTest.java b/mcp/support/src/test/java/org/apache/shardingsphere/mcp/support/descriptor/MCPDescriptorCatalogValidatorTest.java
index c5e099e4186..9d824043b72 100644
--- a/mcp/support/src/test/java/org/apache/shardingsphere/mcp/support/descriptor/MCPDescriptorCatalogValidatorTest.java
+++ b/mcp/support/src/test/java/org/apache/shardingsphere/mcp/support/descriptor/MCPDescriptorCatalogValidatorTest.java
@@ -157,6 +157,44 @@ class MCPDescriptorCatalogValidatorTest {
"Planning tool `database_gateway_test_plan_rule` annotations.idempotentHint must be false.");
}
+ @Test
+ void assertValidateRejectsDestructiveToolWithoutExecutionMode() {
+ assertDestructiveToolValidationError(false, List.of(), false, List.of("database"),
+ "Destructive tool `database_gateway_test_tool` must declare execution_mode.");
+ }
+
+ @Test
+ void assertValidateRejectsDestructiveToolWithOptionalExecutionMode() {
+ assertDestructiveToolValidationError(true, List.of("preview"), false, List.of("database"),
+ "Destructive tool `database_gateway_test_tool` execution_mode must be required.");
+ }
+
+ @Test
+ void assertValidateRejectsDestructiveToolWithoutPreview() {
+ assertDestructiveToolValidationError(true, List.of("review-then-execute"), true, List.of("database"),
+ "Destructive tool `database_gateway_test_tool` execution_mode must allow preview.");
+ }
+
+ @Test
+ void assertValidateRejectsDestructiveToolWithAutoExecute() {
+ assertDestructiveToolValidationError(true, List.of("preview", "auto-execute"), true, List.of("database"),
+ "Destructive tool `database_gateway_test_tool` execution_mode must not expose auto-execute.");
+ }
+
+ @Test
+ void assertValidateRejectsDestructiveToolWithoutSideEffectScope() {
+ assertDestructiveToolValidationError(true, List.of("preview"), true, List.of(),
+ "Destructive tool `database_gateway_test_tool` must declare sideEffectScope in internal runtime.");
+ }
+
+ @ParameterizedTest(name = "{0}")
+ @ValueSource(strings = {"preview", "auto-execute"})
+ void assertValidateRejectsNonDestructiveExecutionMode(final String executionMode) {
+ MCPToolDescriptor descriptor = createExecutionModeTool(false, true, List.of(executionMode), true);
+ assertValidationError(createToolRuntimeCatalog(List.of(), List.of(descriptor), List.of()),
+ String.format("Non-destructive tool `database_gateway_test_tool` execution_mode must not expose %s.", executionMode));
+ }
+
@Test
void assertValidateRejectsUnsupportedNextActionSchemaField() {
assertValidationError(createCatalog(List.of(), List.of(createToolDescriptor(
@@ -268,6 +306,25 @@ class MCPDescriptorCatalogValidatorTest {
assertThat(actual.getMessage(), is(expectedMessage));
}
+ private void assertDestructiveToolValidationError(final boolean executionModePresent, final List<String> executionModes, final boolean executionModeRequired,
+ final List<String> sideEffectScope, final String expectedMessage) {
+ MCPToolDescriptor descriptor = createExecutionModeTool(true, executionModePresent, executionModes, executionModeRequired);
+ assertValidationError(createToolRuntimeCatalog(List.of(), List.of(descriptor),
+ List.of(new MCPToolRuntimeDescriptor(descriptor.getName(), "execute", sideEffectScope))), expectedMessage);
+ }
+
+ private MCPToolDescriptor createExecutionModeTool(final boolean destructive, final boolean executionModePresent, final List<String> executionModes,
+ final boolean executionModeRequired) {
+ Map<String, Object> properties = executionModePresent
+ ? Map.of(MCPPayloadFieldNames.EXECUTION_MODE, Map.of("type", "string", "description", "Execution mode.", "enum", executionModes))
+ : Map.of();
+ List<String> required = executionModeRequired ? List.of(MCPPayloadFieldNames.EXECUTION_MODE) : List.of();
+ MCPToolAnnotations annotations = MCPToolAnnotations.builder().title("Test Tool").readOnlyHint(false).destructiveHint(destructive)
+ .idempotentHint(false).openWorldHint(true).build();
+ return new MCPToolDescriptor("database_gateway_test_tool", "Test Tool", "Run a test tool.", createInputSchema(properties, required),
+ createOutputSchema(), annotations, Map.of());
+ }
+
private MCPDescriptorCatalog createCatalog(final List<MCPResourceDescriptor> resourceDescriptors, final List<MCPToolDescriptor> toolDescriptors) {
return createCatalog(resourceDescriptors, List.of(), toolDescriptors);
}
diff --git a/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportProtocolContractE2ETest.java b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportProtocolContractE2ETest.java
index 8586ddbc42e..e798912fdac 100644
--- a/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportProtocolContractE2ETest.java
+++ b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportProtocolContractE2ETest.java
@@ -52,6 +52,7 @@ class HttpTransportProtocolContractE2ETest extends AbstractHttpProtocolOnlyE2ETe
assertThat(actual.headers().firstValue("MCP-Protocol-Version").orElse(""), is(getProtocolVersion()));
String actualSessionId = actual.headers().firstValue("MCP-Session-Id").orElse("");
assertFalse(actualSessionId.isEmpty());
+ assertTrue(actualSessionId.chars().allMatch(each -> 0x21 <= each && each <= 0x7E));
Map<String, Object> actualPayload = parseJsonBody(actual.body());
Map<String, Object> actualResult = MCPInteractionPayloads.getRequiredJsonRpcResult(actualPayload);
assertThat(String.valueOf(actualPayload.get("jsonrpc")), is("2.0"));
@@ -230,6 +231,17 @@ class HttpTransportProtocolContractE2ETest extends AbstractHttpProtocolOnlyE2ETe
assertThat(String.valueOf(MCPInteractionPayloads.getRequiredJsonRpcResult(actualPayload).get("protocolVersion")), is(getProtocolVersion()));
}
+ @Test
+ void assertRejectInitializeWithoutProtocolVersion() throws IOException, InterruptedException {
+ launchHttpTransport();
+ HttpClient httpClient = HttpClient.newHttpClient();
+ Map<String, Object> initializeRequestParams = new LinkedHashMap<>(MCPInteractionProtocolSupport.createInitializeRequestParams("mcp-e2e-programmatic"));
+ initializeRequestParams.remove("protocolVersion");
+ HttpResponse<String> actual = sendInitializeRequest(httpClient, initializeRequestParams);
+ assertThat(actual.statusCode(), is(400));
+ assertFalse(actual.headers().firstValue("MCP-Session-Id").isPresent());
+ }
+
@Test
void assertRejectFollowUpRequestWithProtocolMismatch() throws IOException, InterruptedException {
launchHttpTransport();
diff --git a/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionAttributionE2ETest.java b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionAttributionE2ETest.java
new file mode 100644
index 00000000000..b2bdd6b83ab
--- /dev/null
+++ b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionAttributionE2ETest.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.shardingsphere.test.e2e.mcp.runtime.programmatic;
+
+import org.apache.shardingsphere.mcp.bootstrap.config.HttpTransportConfiguration;
+import org.apache.shardingsphere.mcp.bootstrap.config.SessionAttributionSourceConfiguration;
+import org.apache.shardingsphere.test.e2e.mcp.support.transport.MCPInteractionPayloads;
+import org.apache.shardingsphere.test.e2e.mcp.support.transport.MCPInteractionProtocolSupport;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.net.http.HttpClient;
+import java.net.http.HttpResponse;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+
+class HttpTransportSessionAttributionE2ETest extends AbstractHttpProtocolOnlyE2ETest {
+
+ private static final String SUBJECT_HEADER = "X-Test-Subject";
+
+ private static final String SOURCE_HEADER = "X-Test-Source";
+
+ private static final String ATTRIBUTE_HEADER = "X-Test-Attr-Region";
+
+ @Test
+ void assertAcceptMatchedSessionAttribution() throws IOException, InterruptedException {
+ launchHttpTransport();
+ HttpClient httpClient = HttpClient.newHttpClient();
+ String sessionId = initializeAttributedSession(httpClient);
+ assertThat(sendCapabilitiesRequest(httpClient, createAttributedSessionHeaders(sessionId, "subject")).statusCode(), is(200));
+ }
+
+ @Test
+ void assertRejectChangedSessionAttribution() throws IOException, InterruptedException {
+ launchHttpTransport();
+ HttpClient httpClient = HttpClient.newHttpClient();
+ String sessionId = initializeAttributedSession(httpClient);
+ HttpResponse<String> actual = sendCapabilitiesRequest(httpClient, createAttributedSessionHeaders(sessionId, "other"));
+ assertThat(actual.statusCode(), is(400));
+ Map<String, Object> error = MCPInteractionPayloads.getRequiredObject(parseJsonBody(actual.body()), "error");
+ Map<String, Object> data = MCPInteractionPayloads.getRequiredObject(error, "data");
+ assertThat(MCPInteractionPayloads.getRequiredObject(data, "recovery").get("category"), is("session_attribution_mismatch"));
+ }
+
+ private String initializeAttributedSession(final HttpClient httpClient) throws IOException, InterruptedException {
+ Map<String, String> headers = createAttributionHeaders("subject");
+ HttpResponse<String> initializeResponse = sendInitializeRequest(
+ httpClient, headers, MCPInteractionProtocolSupport.createInitializeRequestParams("mcp-e2e-programmatic"));
+ assertThat(initializeResponse.statusCode(), is(200));
+ String result = initializeResponse.headers().firstValue("MCP-Session-Id").orElseThrow();
+ assertThat(sendRawPostRequest(httpClient, createAttributedSessionHeaders(result, "subject"),
+ MCPInteractionProtocolSupport.createJsonRpcNotificationBody("notifications/initialized", Map.of())).statusCode(), is(202));
+ return result;
+ }
+
+ private Map<String, String> createAttributedSessionHeaders(final String sessionId, final String subject) {
+ Map<String, String> result = new LinkedHashMap<>(createSessionHeaders(sessionId));
+ result.putAll(createAttributionHeaders(subject));
+ return result;
+ }
+
+ private Map<String, String> createAttributionHeaders(final String subject) {
+ return Map.of(SUBJECT_HEADER, subject, SOURCE_HEADER, "gateway", ATTRIBUTE_HEADER, "ap-south");
+ }
+
+ @Override
+ protected HttpTransportConfiguration createHttpTransportConfiguration() {
+ return new HttpTransportConfiguration("127.0.0.1", 0, getHttpEndpointPath(),
+ new SessionAttributionSourceConfiguration(SUBJECT_HEADER, SOURCE_HEADER, "X-Test-Attr-"));
+ }
+}
diff --git a/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionLifecycleE2ETest.java b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionLifecycleE2ETest.java
index 23b951d248d..86671d5cf54 100644
--- a/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionLifecycleE2ETest.java
+++ b/test/e2e/mcp/src/test/java/org/apache/shardingsphere/test/e2e/mcp/runtime/programmatic/HttpTransportSessionLifecycleE2ETest.java
@@ -39,7 +39,7 @@ class HttpTransportSessionLifecycleE2ETest extends AbstractHttpProtocolOnlyE2ETe
"Accept", "text/event-stream",
"MCP-Session-Id", sessionId,
"MCP-Protocol-Version", getProtocolVersion()));
- assertThat(actual.statusCode(), is(405));
+ assertThat(actual.statusCode(), is(404));
}
@Test