diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 760af7c..90494d1 100644 --- a/.github/workflows/build.yaml +++ b/.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/checkout@v3.0.0 - 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/checkout@v3.0.0 - name: Run GoReleaser uses: goreleaser/goreleaser-action@v6 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 62f9965..de21285 100644 --- a/.github/workflows/release.yaml +++ b/.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/checkout@v3.0.0 - 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}} diff --git a/Dockerfile b/Dockerfile index 41e0a3f..298275d 100644 --- a/Dockerfile +++ b/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 diff --git a/README.md b/README.md index 2c11308..6e6aaa6 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,7 @@ 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: @@ -16,6 +17,26 @@ To use this extension in your API testing project, follow these steps: 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): diff --git a/cmd/mcp.go b/cmd/mcp.go new file mode 100644 index 0000000..64e1d27 --- /dev/null +++ b/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 +} diff --git a/cmd/root.go b/cmd/root.go index 916f5b7..0b3c3c4 100644 --- a/cmd/root.go +++ b/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 } diff --git a/e2e/compose.yaml b/e2e/compose.yaml index ea4ff4e..c075333 100644 --- a/e2e/compose.yaml +++ b/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: diff --git a/go.mod b/go.mod index 43d0c50..43f8f4b 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 0f5bc8e..3b07961 100644 --- a/go.sum +++ b/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= diff --git a/pkg/data_query.go b/pkg/data_query.go index 3ca8de0..6c48c3d 100644 --- a/pkg/data_query.go +++ b/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, + }) + } } } } diff --git a/pkg/mcp.go b/pkg/mcp.go new file mode 100644 index 0000000..0bc97b1 --- /dev/null +++ b/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 +}