feat: add mcp server support

.github/workflows/build.yaml

@@ -10,7 +10,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
       - uses: actions/[email protected]
       - name: Unit Test
         run: |
@@ -29,7 +29,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
       - uses: actions/[email protected]
       - name: Run GoReleaser
         uses: goreleaser/goreleaser-action@v6

.github/workflows/release.yaml

@@ -20,7 +20,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
       - uses: actions/[email protected]
       - name: Unit Test
         run: |
@@ -48,7 +48,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v3
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
       - name: Image Registry Login
         run: |
           docker login --username linuxsuren --password ${{secrets.DOCKER_HUB_PUBLISH_SECRETS}}

Dockerfile

@@ -1,4 +1,4 @@
-ARG GO_BUILDER=docker.io/library/golang:1.22
+ARG GO_BUILDER=docker.io/library/golang:1.23
 ARG BASE_IMAGE=docker.io/library/alpine:3.12
 
 FROM ${GO_BUILDER} AS builder

README.md

@@ -9,13 +9,34 @@ This project provides an ORM-based database store extension for API testing, sim
 - Simplified database operations using ORM.
 - Integration with API testing frameworks.
 - Support for multiple databases (SQLite, MySQL, PostgreSQL, TDengine, etc.).
+- Database query as a MCP server
 
 ## Usage
 To use this extension in your API testing project, follow these steps:
 1. Install the necessary dependencies.
 2. Configure the database connection settings.
 3. Integrate the extension into your API tests.
 
+## MCP Server
+
+```json
+{
+  "mcpServers": {
+    "database-mcp": {
+      "name": "database-mcp",
+      "type": "stdio",
+      "description": "Database query MCP Server",
+      "isActive": true,
+      "command": "atest-store-orm",
+      "args": [
+        "mcp",
+        "--mode=stdio"
+      ]
+    }
+  }
+}
+```
+
 ## Quick MySQL Setup with TiUP Playground
 
 You can quickly set up a MySQL-compatible database using [TiUP Playground](https://docs.pingcap.com/tidb/stable/tiup-playground):

cmd/mcp.go

@@ -0,0 +1,124 @@
+/*
+Copyright 2025 API Testing Authors.
+
+Licensed 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 cmd
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+
+	"github.com/linuxsuren/api-testing/pkg/testing"
+	"github.com/linuxsuren/atest-ext-store-orm/pkg"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/spf13/cobra"
+)
+
+func newMCPCommand() (c *cobra.Command) {
+	opt := &mcpOption{}
+	c = &cobra.Command{
+		Use:     "mcp",
+		Short:   "Multi-Cluster-Platform related commands",
+		PreRunE: opt.preRunE,
+		RunE:    opt.runE,
+	}
+	flags := c.Flags()
+	flags.StringVarP(&opt.mode, "mode", "", "http", "Server mode, one of http/stdio/sse")
+	flags.IntVarP(&opt.port, "port", "", 7072, "Server port for http or sse mode")
+	flags.StringVarP(&opt.url, "url", "", "", "Database URL")
+	flags.StringVarP(&opt.username, "username", "", "", "Database username")
+	flags.StringVarP(&opt.password, "password", "", "", "Database password")
+	flags.StringVarP(&opt.database, "database", "", "", "Database name")
+	flags.StringVarP(&opt.driver, "driver", "", "mysql", "Database driver, one of mysql/postgres/sqlite")
+	return
+}
+
+type mcpOption struct {
+	mode     string
+	port     int
+	url      string
+	username string
+	password string
+	database string
+	driver   string
+}
+
+func (o *mcpOption) preRunE(c *cobra.Command, args []string) (err error) {
+	if o.url = getValueOrEnv(o.url, "DB_URL"); o.url == "" {
+		err = fmt.Errorf("database url is required")
+		return
+	}
+	o.username = getValueOrEnv(o.username, "DB_USERNAME")
+	o.password = getValueOrEnv(o.password, "DB_PASSWORD")
+	o.database = getValueOrEnv(o.database, "DB_DATABASE")
+	o.driver = getValueOrEnv(o.driver, "DB_DRIVER")
+	return
+}
+
+func getValueOrEnv(value, envKey string) (result string) {
+	if value != "" {
+		result = value
+	} else {
+		result = os.Getenv(envKey)
+	}
+	return
+}
+
+func (o *mcpOption) runE(c *cobra.Command, args []string) (err error) {
+	opts := &mcp.ServerOptions{
+		Instructions: "Database query mcp server",
+	}
+
+	server := mcp.NewServer(&mcp.Implementation{
+		Name:  "database-mcp-server",
+		Title: "ORM Database MCP Server",
+	}, opts)
+
+	store := &testing.Store{
+		URL:      o.url,
+		Username: o.username,
+		Password: o.password,
+		Properties: map[string]string{
+			"database": o.database,
+			"driver":   o.driver,
+		},
+	}
+
+	dbServer := pkg.NewMcpServer(store)
+	mcp.AddTool(server, &mcp.Tool{
+		Name:        "database-query",
+		Description: "Query the database by SQL",
+	}, dbServer.Query)
+
+	switch o.mode {
+	case "sse":
+		handler := mcp.NewSSEHandler(func(request *http.Request) *mcp.Server {
+			return server
+		})
+		c.Println("Starting SSE server on port:", o.port)
+		err = http.ListenAndServe(fmt.Sprintf(":%d", o.port), handler)
+	case "stdio":
+		err = server.Run(c.Context(), &mcp.StdioTransport{})
+	case "http":
+		fallthrough
+	default:
+		handler := mcp.NewStreamableHTTPHandler(func(request *http.Request) *mcp.Server {
+			return server
+		}, nil)
+		c.Println("Starting HTTP server on port:", o.port)
+		err = http.ListenAndServe(fmt.Sprintf(":%d", o.port), handler)
+	}
+	return
+}

cmd/root.go

@@ -34,6 +34,8 @@ func NewRootCommand() (c *cobra.Command) {
 	opt.AddFlags(c.Flags())
 	c.Flags().IntVarP(&opt.historyLimit, "history-limit", "", 1000, "History record items count limit")
 	c.Flags().BoolVarP(&opt.version, "version", "", false, "Print the version then exit")
+
+	c.AddCommand(newMCPCommand())
 	return
 }
 

e2e/compose.yaml

@@ -29,7 +29,7 @@ services:
       context: ..
       dockerfile: Dockerfile
       args:
-        - "GO_BUILDER=ghcr.io/linuxsuren/library/golang:1.22"
+        - "GO_BUILDER=golang:1.23"
         - "BASE_IMAGE=ghcr.io/linuxsuren/library/alpine:3.12"
         - GOPROXY=${GOPROXY}
 #    ports:

go.mod

@@ -1,11 +1,12 @@
 module github.com/linuxsuren/atest-ext-store-orm
 
-go 1.22.4
+go 1.23.0
 
-toolchain go1.22.6
+toolchain go1.24.3
 
 require (
 	github.com/linuxsuren/api-testing v0.0.20-0.20250319020913-f5f9383e2948
+	github.com/modelcontextprotocol/go-sdk v0.3.1
 	github.com/spf13/cobra v1.8.1
 	github.com/stretchr/testify v1.9.0
 	github.com/taosdata/driver-go/v3 v3.6.0
@@ -43,6 +44,7 @@ require (
 	github.com/go-openapi/swag v0.23.0 // indirect
 	github.com/go-sql-driver/mysql v1.7.0 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
+	github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/gorilla/mux v1.8.1 // indirect
 	github.com/gorilla/websocket v1.5.0 // indirect
@@ -89,6 +91,7 @@ require (
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
 	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
 	github.com/xeipuuv/gojsonschema v1.2.0 // indirect
+	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/crypto v0.31.0 // indirect

go.sum

@@ -51,9 +51,11 @@ github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ
 github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
-github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
-github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
+github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 h1:mBlBwtDebdDYr+zdop8N62a44g+Nbv7o2KjWyS1deR4=
+github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
@@ -120,6 +122,8 @@ github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HK
 github.com/mitchellh/reflectwalk v1.0.0/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
 github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
 github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
+github.com/modelcontextprotocol/go-sdk v0.3.1 h1:0z04yIPlSwTluuelCBaL+wUag4YeflIU2Fr4Icb7M+o=
+github.com/modelcontextprotocol/go-sdk v0.3.1/go.mod h1:whv0wHnsTphwq7CTiKYHkLtwLC06WMoY2KpO+RB9yXQ=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -197,6 +201,8 @@ github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHo
 github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ=
 github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74=
 github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
+github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
+github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 github.com/yudai/gojsondiff v1.0.0 h1:27cbfqXLVEJ1o8I6v3y9lg8Ydm53EKqHXAOMxEGlCOA=
 github.com/yudai/gojsondiff v1.0.0/go.mod h1:AY32+k2cwILAkW1fbgxQ5mUmMiZFgLIV+FBNExI05xg=
 github.com/yudai/golcs v0.0.0-20170316035057-ecda9a501e82 h1:BHyfKlQyqbsFN5p3IfnEUduWvb9is428/nNb5L3U01M=
@@ -248,6 +254,8 @@ golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
+golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo=
+golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/genproto v0.0.0-20240123012728-ef4313101c80 h1:KAeGQVN3M9nD0/bQXnr/ClcEMJ968gUXJQ9pwfSynuQ=
 google.golang.org/genproto v0.0.0-20240123012728-ef4313101c80/go.mod h1:cc8bqMqtv9gMOr0zHg2Vzff5ULhhL2IXP4sbcn32Dro=

pkg/data_query.go

@@ -120,6 +120,7 @@ func runMultilineSQL(ctx context.Context, multilineSQL string, db *gorm.DB) (res
 }
 
 func sqlQuery(ctx context.Context, sqlText string, db *gorm.DB) (result *server.DataQueryResult, err error) {
+	fmt.Println("execute sql:", sqlText)
 	var rows *sql.Rows
 	if rows, err = db.Raw(sqlText).Rows(); err != nil {
 		return
@@ -301,14 +302,16 @@ func (q *commonDataQuery) GetCurrentDatabase() (current string, err error) {
 
 func (q *commonDataQuery) GetLabels(ctx context.Context, sql string) (metadata []*server.Pair) {
 	metadata = make([]*server.Pair, 0)
-	if databaseResult, err := sqlQuery(ctx, fmt.Sprintf("explain %s", sql), q.db); err == nil && len(databaseResult.Items) != 1 {
-		for _, data := range databaseResult.Items[0].Data {
-			switch data.Key {
-			case "type":
-				metadata = append(metadata, &server.Pair{
-					Key:   "sql_type",
-					Value: data.Value,
-				})
+	if !strings.Contains(sql, ";") {
+		if databaseResult, err := sqlQuery(ctx, fmt.Sprintf("explain %s", sql), q.db); err == nil && len(databaseResult.Items) != 1 {
+			for _, data := range databaseResult.Items[0].Data {
+				switch data.Key {
+				case "type":
+					metadata = append(metadata, &server.Pair{
+						Key:   "sql_type",
+						Value: data.Value,
+					})
+				}
 			}
 		}
 	}

pkg/mcp.go

@@ -0,0 +1,60 @@
+/*
+Copyright 2025 API Testing Authors.
+
+Licensed 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 pkg
+
+import (
+	"context"
+
+	"github.com/linuxsuren/api-testing/pkg/server"
+	"github.com/linuxsuren/api-testing/pkg/testing"
+	"github.com/linuxsuren/api-testing/pkg/testing/remote"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+type mcpServer struct {
+	store *testing.Store
+}
+
+type DBQuery struct {
+	SQL string `json:"sql" jsonschema:"the sql to be executed"`
+}
+
+type DatabaseQuery interface {
+	Query(ctx context.Context, request *mcp.CallToolRequest, query DBQuery) (
+		result *mcp.CallToolResult, a any, err error)
+}
+
+func NewMcpServer(store *testing.Store) DatabaseQuery {
+	return &mcpServer{store: store}
+}
+
+func (s *mcpServer) Query(ctx context.Context, request *mcp.CallToolRequest, query DBQuery) (
+	result *mcp.CallToolResult, a any, err error) {
+	db := &dbserver{}
+	ctx = remote.WithIncomingStoreContext(ctx, s.store)
+	result = &mcp.CallToolResult{}
+
+	var queryResult *server.DataQueryResult
+	if queryResult, err = db.Query(ctx, &server.DataQuery{
+		Sql: query.SQL,
+	}); err == nil {
+		result.StructuredContent = queryResult
+		result.Content = []mcp.Content{
+			&mcp.TextContent{Text: queryResult.String()},
+		}
+	}
+	return
+}