Skip to content

Commit c14a6d9

Browse files
oleksandrkaplunviroleksandrkaplunvir
authored
chore: migrate to official mcp sdk (#56)
- switch to use official mcp sdk from https://github.com/modelcontextprotocol/go-sdk --------- Signed-off-by: Alex Kaplun <o.kaplun@firebolt.io>
1 parent 2a9c23e commit c14a6d9

21 files changed

Lines changed: 824 additions & 831 deletions

‎go.mod‎

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -3,46 +3,45 @@ module github.com/firebolt-db/mcp-server
33
go 1.24.1
44

55
require (
6-
github.com/JohannesKaufmann/html-to-markdown/v2 v2.3.3
7-
github.com/firebolt-db/firebolt-go-sdk v1.10.0
8-
github.com/gocolly/colly/v2 v2.2.0
9-
github.com/mark3labs/mcp-go v0.34.0
6+
github.com/JohannesKaufmann/html-to-markdown/v2 v2.5.0
7+
github.com/firebolt-db/firebolt-go-sdk v1.13.0
8+
github.com/gocolly/colly/v2 v2.3.0
9+
github.com/modelcontextprotocol/go-sdk v1.2.0
1010
github.com/neilotoole/slogt v1.1.0
11-
github.com/stretchr/testify v1.10.0
12-
github.com/urfave/cli/v3 v3.3.8
13-
golang.org/x/oauth2 v0.30.0
14-
golang.org/x/sync v0.16.0
11+
github.com/stretchr/testify v1.11.1
12+
github.com/urfave/cli/v3 v3.6.1
13+
golang.org/x/oauth2 v0.34.0
14+
golang.org/x/sync v0.19.0
1515
)
1616

1717
require (
1818
github.com/JohannesKaufmann/dom v0.2.0 // indirect
19-
github.com/PuerkitoBio/goquery v1.10.2 // indirect
19+
github.com/PuerkitoBio/goquery v1.11.0 // indirect
2020
github.com/andybalholm/cascadia v1.3.3 // indirect
21-
github.com/antchfx/htmlquery v1.3.4 // indirect
22-
github.com/antchfx/xmlquery v1.4.4 // indirect
23-
github.com/antchfx/xpath v1.3.3 // indirect
21+
github.com/antchfx/htmlquery v1.3.5 // indirect
22+
github.com/antchfx/xmlquery v1.5.0 // indirect
23+
github.com/antchfx/xpath v1.3.5 // indirect
2424
github.com/astaxie/beego v1.12.3 // indirect
25-
github.com/bits-and-blooms/bitset v1.22.0 // indirect
25+
github.com/bits-and-blooms/bitset v1.24.4 // indirect
2626
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
2727
github.com/gobwas/glob v0.2.3 // indirect
2828
github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect
2929
github.com/golang/protobuf v1.5.4 // indirect
30-
github.com/google/go-cmp v0.7.0 // indirect
31-
github.com/google/uuid v1.6.0 // indirect
30+
github.com/google/jsonschema-go v0.4.2 // indirect
3231
github.com/kennygrant/sanitize v1.2.4 // indirect
33-
github.com/matishsiao/goInfo v0.0.0-20210923090445-da2e3fa8d45f // indirect
34-
github.com/nlnwa/whatwg-url v0.6.1 // indirect
32+
github.com/kr/text v0.2.0 // indirect
33+
github.com/matishsiao/goInfo v0.0.0-20241216093258-66a9250504d6 // indirect
34+
github.com/nlnwa/whatwg-url v0.6.2 // indirect
3535
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
3636
github.com/saintfish/chardet v0.0.0-20230101081208-5e3ef4b5456d // indirect
3737
github.com/shopspring/decimal v1.4.0 // indirect
38-
github.com/spf13/cast v1.7.1 // indirect
39-
github.com/stretchr/objx v0.5.2 // indirect
38+
github.com/stretchr/objx v0.5.3 // indirect
4039
github.com/temoto/robotstxt v1.1.2 // indirect
4140
github.com/xwb1989/sqlparser v0.0.0-20180606152119-120387863bf2 // indirect
4241
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
43-
golang.org/x/net v0.39.0 // indirect
44-
golang.org/x/text v0.24.0 // indirect
42+
golang.org/x/net v0.48.0 // indirect
43+
golang.org/x/text v0.33.0 // indirect
4544
google.golang.org/appengine v1.6.8 // indirect
46-
google.golang.org/protobuf v1.36.6 // indirect
45+
google.golang.org/protobuf v1.36.11 // indirect
4746
gopkg.in/yaml.v3 v3.0.1 // indirect
4847
)

‎go.sum‎

Lines changed: 51 additions & 12 deletions
Large diffs are not rendered by default.

‎pkg/prompts/firebolt.go‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import (
44
"context"
55
_ "embed"
66

7-
"github.com/mark3labs/mcp-go/mcp"
7+
"github.com/modelcontextprotocol/go-sdk/mcp"
88
)
99

1010
//go:embed firebolt.md
@@ -19,20 +19,25 @@ func NewFireboltExpert() *FireboltExpert {
1919
return &FireboltExpert{}
2020
}
2121

22-
func (p *FireboltExpert) Prompt() mcp.Prompt {
23-
return mcp.NewPrompt(
24-
"Firebolt Expert",
25-
mcp.WithPromptDescription(fireboltPromptDescription),
26-
)
22+
func (p *FireboltExpert) Prompt() *mcp.Prompt {
23+
return &mcp.Prompt{
24+
Name: "Firebolt Expert",
25+
Description: fireboltPromptDescription,
26+
}
2727
}
2828

29-
func (p *FireboltExpert) Handler(_ context.Context, _ mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
29+
func (p *FireboltExpert) Handler(_ context.Context, _ *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
3030

31-
var messages []mcp.PromptMessage
32-
messages = append(messages, mcp.NewPromptMessage(
33-
mcp.RoleAssistant,
34-
mcp.NewTextContent(fireboltMD),
35-
))
31+
var messages []*mcp.PromptMessage
32+
messages = append(messages, &mcp.PromptMessage{
33+
Role: "assistant",
34+
Content: &mcp.TextContent{
35+
Text: fireboltMD,
36+
},
37+
})
3638

37-
return mcp.NewGetPromptResult(fireboltPromptDescription, messages), nil
39+
return &mcp.GetPromptResult{
40+
Description: fireboltPromptDescription,
41+
Messages: messages,
42+
}, nil
3843
}

‎pkg/resources/accounts.go‎

Lines changed: 25 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"encoding/json"
66
"fmt"
77

8-
"github.com/mark3labs/mcp-go/mcp"
8+
"github.com/modelcontextprotocol/go-sdk/mcp"
99

1010
"github.com/firebolt-db/mcp-server/pkg/clients/discovery"
1111
"github.com/firebolt-db/mcp-server/pkg/helpers/args"
@@ -34,21 +34,24 @@ func NewAccounts(discoveryClient discovery.Client) *Accounts {
3434

3535
// ResourceTemplate defines the template for account resources.
3636
// It specifies the URI format, content type, description, and suggested usage.
37-
func (r *Accounts) ResourceTemplate() mcp.ResourceTemplate {
38-
return mcp.NewResourceTemplate(
39-
AccountURI("{account}"),
40-
"Account",
41-
mcp.WithTemplateMIMEType(mimetype.JSON),
42-
mcp.WithTemplateAnnotations([]mcp.Role{mcp.RoleUser, mcp.RoleAssistant}, 0.9),
43-
mcp.WithTemplateDescription("Brief information about the account in the Firebolt organization."),
44-
)
37+
func (r *Accounts) ResourceTemplate() *mcp.ResourceTemplate {
38+
return &mcp.ResourceTemplate{
39+
URITemplate: AccountURI("{account}"),
40+
Name: "Account",
41+
MIMEType: mimetype.JSON,
42+
Description: "Brief information about the account in the Firebolt organization.",
43+
Annotations: &mcp.Annotations{
44+
Audience: []mcp.Role{"user", "assistant"},
45+
Priority: 0.9,
46+
},
47+
}
4548
}
4649

4750
// Handler processes resource requests for account information.
4851
// It extracts the account parameter and fetches the appropriate account data.
49-
func (r *Accounts) Handler(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
52+
func (r *Accounts) Handler(ctx context.Context, request *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) {
5053

51-
accountName, err := args.String(request.Params.Arguments, "account")
54+
accountName, err := args.String(request.GetParams().GetMeta(), "account")
5255
if err != nil {
5356
return nil, fmt.Errorf("bad request: %w", err)
5457
}
@@ -58,7 +61,7 @@ func (r *Accounts) Handler(ctx context.Context, request mcp.ReadResourceRequest)
5861

5962
// FetchAccountResources retrieves account information from the Firebolt service.
6063
// If a specific account name is specified, it filters for that account; otherwise, it returns all accounts.
61-
func (r *Accounts) FetchAccountResources(ctx context.Context, accountName string) ([]mcp.ResourceContents, error) {
64+
func (r *Accounts) FetchAccountResources(ctx context.Context, accountName string) (*mcp.ReadResourceResult, error) {
6265

6366
// Fetch the list of accounts from the discovery client
6467
accounts, err := r.discoveryClient.ListAccounts(ctx)
@@ -74,17 +77,23 @@ func (r *Accounts) FetchAccountResources(ctx context.Context, accountName string
7477
}
7578
}
7679

77-
return itertools.MapWithFailure(filteredAccounts, func(i discovery.Account) (mcp.ResourceContents, error) {
78-
80+
out, err := itertools.MapWithFailure(filteredAccounts, func(i discovery.Account) (*mcp.ResourceContents, error) {
7981
data, err := json.Marshal(i)
8082
if err != nil {
81-
return mcp.TextResourceContents{}, fmt.Errorf("failed to marshal row data to JSON: %w", err)
83+
return nil, fmt.Errorf("failed to marshal row data to JSON: %w", err)
8284
}
8385

84-
return mcp.TextResourceContents{
86+
return &mcp.ResourceContents{
8587
URI: AccountURI(i.Name),
8688
MIMEType: mimetype.JSON,
8789
Text: string(data),
8890
}, nil
8991
})
92+
if err != nil {
93+
return nil, err
94+
}
95+
96+
return &mcp.ReadResourceResult{
97+
Contents: out,
98+
}, nil
9099
}

‎pkg/resources/accounts_test.go‎

Lines changed: 17 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import (
66
"errors"
77
"testing"
88

9-
"github.com/mark3labs/mcp-go/mcp"
9+
"github.com/modelcontextprotocol/go-sdk/mcp"
1010
"github.com/stretchr/testify/assert"
1111
"github.com/stretchr/testify/require"
1212

@@ -96,7 +96,7 @@ func TestAccounts_Handler(t *testing.T) {
9696
},
9797
},
9898
{
99-
name: "missing account parameter",
99+
name: "missing account parameter",
100100
params: map[string]any{
101101
// No account parameter
102102
},
@@ -126,8 +126,11 @@ func TestAccounts_Handler(t *testing.T) {
126126

127127
accounts := resources.NewAccounts(client)
128128

129-
request := mcp.ReadResourceRequest{}
130-
request.Params.Arguments = tt.params
129+
request := &mcp.ReadResourceRequest{
130+
Params: &mcp.ReadResourceParams{
131+
Meta: tt.params,
132+
},
133+
}
131134

132135
result, err := accounts.Handler(t.Context(), request)
133136

@@ -138,19 +141,17 @@ func TestAccounts_Handler(t *testing.T) {
138141
}
139142

140143
require.NoError(t, err)
141-
require.Len(t, result, len(tt.expected))
144+
require.Len(t, result.Contents, len(tt.expected))
142145

143-
for i, res := range result {
144-
textRes, ok := res.(mcp.TextResourceContents)
145-
require.True(t, ok, "Expected TextResourceContents")
146+
for i, res := range result.Contents {
146147

147148
var data discovery.Account
148-
err := json.Unmarshal([]byte(textRes.Text), &data)
149+
err := json.Unmarshal([]byte(res.Text), &data)
149150
require.NoError(t, err)
150151

151152
assert.Equal(t, tt.expected[i], data)
152-
assert.Equal(t, resources.AccountURI(data.Name), textRes.URI)
153-
assert.Equal(t, mimetype.JSON, textRes.MIMEType)
153+
assert.Equal(t, resources.AccountURI(data.Name), res.URI)
154+
assert.Equal(t, mimetype.JSON, res.MIMEType)
154155
}
155156

156157
assert.Equal(t, 1, client.ListAccountsCount, "ListAccounts should be called exactly once")
@@ -277,19 +278,16 @@ func TestAccounts_FetchAccountResources(t *testing.T) {
277278
}
278279

279280
require.NoError(t, err)
280-
require.Len(t, result, len(tt.expected))
281-
282-
for i, res := range result {
283-
textRes, ok := res.(mcp.TextResourceContents)
284-
require.True(t, ok, "Expected TextResourceContents")
281+
require.Len(t, result.Contents, len(tt.expected))
285282

283+
for i, res := range result.Contents {
286284
var data discovery.Account
287-
err := json.Unmarshal([]byte(textRes.Text), &data)
285+
err := json.Unmarshal([]byte(res.Text), &data)
288286
require.NoError(t, err)
289287

290288
assert.Equal(t, tt.expected[i], data)
291-
assert.Equal(t, resources.AccountURI(data.Name), textRes.URI)
292-
assert.Equal(t, mimetype.JSON, textRes.MIMEType)
289+
assert.Equal(t, resources.AccountURI(data.Name), res.URI)
290+
assert.Equal(t, mimetype.JSON, res.MIMEType)
293291
}
294292

295293
assert.Equal(t, 1, client.ListAccountsCount, "ListAccounts should be called exactly once")

‎pkg/resources/databases.go‎

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import (
55
"encoding/json"
66
"fmt"
77

8-
"github.com/mark3labs/mcp-go/mcp"
8+
"github.com/modelcontextprotocol/go-sdk/mcp"
99

1010
"github.com/firebolt-db/mcp-server/pkg/clients/database"
1111
"github.com/firebolt-db/mcp-server/pkg/helpers/args"
@@ -32,21 +32,24 @@ func NewDatabases(dbPool database.Pool) *Databases {
3232

3333
// ResourceTemplate defines the template for database resources.
3434
// It specifies the URI format, content type, description, and suggested usage.
35-
func (r *Databases) ResourceTemplate() mcp.ResourceTemplate {
36-
return mcp.NewResourceTemplate(
37-
DatabaseURI("{account}", "{database}"),
38-
"Database",
39-
mcp.WithTemplateMIMEType(mimetype.JSON),
40-
mcp.WithTemplateAnnotations([]mcp.Role{mcp.RoleUser, mcp.RoleAssistant}, 0.8),
41-
mcp.WithTemplateDescription("Brief information about the database in the Firebolt account."),
42-
)
35+
func (r *Databases) ResourceTemplate() *mcp.ResourceTemplate {
36+
return &mcp.ResourceTemplate{
37+
URITemplate: DatabaseURI("{account}", "{database}"),
38+
Name: "Database",
39+
MIMEType: mimetype.JSON,
40+
Description: "Brief information about the database in the Firebolt account.",
41+
Annotations: &mcp.Annotations{
42+
Audience: []mcp.Role{"user", "assistant"},
43+
Priority: 0.8,
44+
},
45+
}
4346
}
4447

4548
// Handler processes resource requests for database information.
4649
// It extracts account and database parameters and fetches the appropriate database data.
47-
func (r *Databases) Handler(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
50+
func (r *Databases) Handler(ctx context.Context, request *mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) {
4851

49-
values, err := args.Strings(request.Params.Arguments, "account", "database")
52+
values, err := args.Strings(request.GetParams().GetMeta(), "account", "database")
5053
if err != nil {
5154
return nil, fmt.Errorf("bad request: %w", err)
5255
}
@@ -56,7 +59,7 @@ func (r *Databases) Handler(ctx context.Context, request mcp.ReadResourceRequest
5659

5760
// FetchDatabaseResources retrieves database information from the Firebolt service.
5861
// If a specific database is specified, it filters for that database; otherwise, it returns all databases.
59-
func (r *Databases) FetchDatabaseResources(ctx context.Context, account, dbName string) ([]mcp.ResourceContents, error) {
62+
func (r *Databases) FetchDatabaseResources(ctx context.Context, account, dbName string) (*mcp.ReadResourceResult, error) {
6063

6164
// Acquire a connection to the database
6265
conn, err := r.dbPool.GetConnection(database.PoolParams{
@@ -83,18 +86,25 @@ func (r *Databases) FetchDatabaseResources(ctx context.Context, account, dbName
8386
}
8487

8588
// Convert rows to resources
86-
return itertools.MapWithFailure(rows, func(i map[string]any) (mcp.ResourceContents, error) {
89+
out, err := itertools.MapWithFailure(rows, func(i map[string]any) (*mcp.ResourceContents, error) {
8790

8891
i["account_name"] = account
8992
data, err := json.Marshal(i)
9093
if err != nil {
91-
return mcp.TextResourceContents{}, fmt.Errorf("failed to marshal row data to JSON: %w", err)
94+
return nil, fmt.Errorf("failed to marshal row data to JSON: %w", err)
9295
}
9396

94-
return mcp.TextResourceContents{
97+
return &mcp.ResourceContents{
9598
URI: DatabaseURI(account, i["database_name"].(string)),
9699
MIMEType: mimetype.JSON,
97100
Text: string(data),
98101
}, nil
99102
})
103+
if err != nil {
104+
return nil, err
105+
}
106+
107+
return &mcp.ReadResourceResult{
108+
Contents: out,
109+
}, nil
100110
}

0 commit comments

Comments
 (0)