# Directory Structure
```
├── .github
│ └── workflows
│ └── test.yml
├── common
│ └── client.go
├── go.mod
├── go.sum
├── main.go
├── mcp_kusto_test.png
├── mcp.json
├── README.md
└── tools
├── common.go
├── databases_test.go
├── databases.go
├── query_test.go
├── query.go
├── tables_test.go
└── tables.go
```
# Files
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
```markdown
1 | # Vibe querying with MCP server for Azure Data Explorer (Kusto)
2 |
3 | This is an implementation of an MCP server for Azure Data Explorer (Kusto) built using its [Go SDK](https://github.com/Azure/azure-kusto-go). You can use this with [VS Code](https://code.visualstudio.com/blogs/2025/02/24/introducing-copilot-agent-mode) (or other MCP tools) for making data analysis and exploration easier.
4 |
5 | It exposes tools for interacting with Azure Data Explorer:
6 |
7 | 1. **list_databases** - Lists all databases in a specific Azure Data Explorer cluster.
8 | 2. **list_tables** - Lists all tables in a specific Azure Data Explorer database.
9 | 3. **get_table_schema** - Gets the schema of a specific table in an Azure Data Explorer database.
10 | 4. **execute_query** - Executes a read-only KQL query against a database.
11 |
12 | > Word(s) of caution: As much as I want folks to benefit from this, I have to call out that Large Language Models (LLMs) are non-deterministic by nature and can make mistakes. I would recommend you to **always validate** the results and queries before making any decisions based on them.
13 |
14 | Here is a sneak peek:
15 |
16 | 
17 |
18 | ## How to run
19 |
20 | ```bash
21 | git clone https://github.com/abhirockzz/mcp_kusto
22 | cd mcp_kusto
23 |
24 | go build -o mcp_kusto main.go
25 | ```
26 |
27 | ### Configure the MCP server
28 |
29 | This will differ based on the MCP client/tool you use. For VS Code you can [follow these instructions](https://code.visualstudio.com/docs/copilot/chat/mcp-servers#_add-an-mcp-server) on how to configure this server using a `mcp.json` file.
30 |
31 | Here is an example of the [mcp.json file](mcp.json):
32 |
33 | ```json
34 | {
35 | "servers": {
36 | "Kusto MCP server": {
37 | "type": "stdio",
38 | "command": "enter path to binary e.g. /Users/demo/Desktop/mcp_kusto",
39 | "args": []
40 | },
41 | //other MCP servers...
42 | }
43 | }
44 | ```
45 |
46 | Here is an example of Claude Desktop configuration:
47 |
48 | ```json
49 | {
50 | "mcpServers": {
51 | "Kusto MCP server": {
52 | "command": "enter path to binary e.g. /Users/demo/Desktop/mcp_kusto",
53 | "args": []
54 | },
55 | //other MCP servers...
56 | }
57 | }
58 | ```
59 |
60 | ### Authentication
61 |
62 | - The user principal you use should have permissions required for `.show databases`, `.show table`, `.show tables`, and execute queries on the database. Refer to the documentation for [Azure Data Explorer](https://learn.microsoft.com/en-us/kusto/management/security-roles?view=azure-data-explorer) for more details.
63 |
64 | - Authentication (Local credentials) - To keep things secure and simple, the MCP server uses [DefaultAzureCredential](https://learn.microsoft.com/en-us/azure/developer/go/sdk/authentication/credential-chains#defaultazurecredential-overview). This approach looks in the environment variables for an application service principal or at locally installed developer tools, such as the Azure CLI, for a set of developer credentials. Either approach can be used to authenticate the MCP server to Azure Data Explorer. For example, just login locally using Azure CLI ([az login](https://learn.microsoft.com/en-us/cli/azure/authenticate-azure-cli)).
65 |
66 | You are good to go! Now spin up VS Code, Claude Desktop, or any other MCP tool and start vibe querying your Azure Data Explorer (Kusto) cluster!
67 |
68 | ## Local dev/testing
69 |
70 | Start with [MCP inspector](https://modelcontextprotocol.io/docs/tools/inspector) - `npx @modelcontextprotocol/inspector ./mcp_kusto`
71 |
```
--------------------------------------------------------------------------------
/mcp.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "servers": {
3 | "Kusto MCP server": {
4 | "type": "stdio",
5 | "command": "enter path to binary e.g. /Users/demo/Desktop/mcp_kusto",
6 | "args": []
7 | }
8 | }
9 | }
10 |
```
--------------------------------------------------------------------------------
/tools/common.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | const CLUSTER_PARAMETER_DESCRIPTION = "Name of the cluster. If not available, ask the user to provide the cluster name. Do not use a random cluster name of your choice."
4 |
5 | const clusterNameFormat = "https://%s.kusto.windows.net/"
6 |
```
--------------------------------------------------------------------------------
/main.go:
--------------------------------------------------------------------------------
```go
1 | package main
2 |
3 | import (
4 | "fmt"
5 |
6 | "github.com/abhirockzz/mcp_kusto/tools"
7 |
8 | "github.com/mark3labs/mcp-go/server"
9 | )
10 |
11 | func main() {
12 |
13 | s := server.NewMCPServer(
14 | "Kusto MCP server",
15 | "0.0.5",
16 | server.WithLogging(),
17 | )
18 |
19 | s.AddTool(tools.ListDatabases())
20 | s.AddTool(tools.ListTables())
21 | s.AddTool(tools.GetTableSchema())
22 | s.AddTool(tools.ExecuteQuery())
23 |
24 | // Start the stdio server
25 | if err := server.ServeStdio(s); err != nil {
26 | fmt.Printf("Server error: %v\n", err)
27 | }
28 | }
29 |
```
--------------------------------------------------------------------------------
/.github/workflows/test.yml:
--------------------------------------------------------------------------------
```yaml
1 | name: Run Tests
2 |
3 | on:
4 | push:
5 | branches:
6 | - main
7 |
8 | jobs:
9 | test:
10 | runs-on: ubuntu-latest
11 |
12 | env:
13 | CLUSTER_NAME: ${{ secrets.CLUSTER_NAME }}
14 | DB_NAME: testdb
15 | TABLE_NAME: test_table
16 | COLUMN_NAMES: column1,column2,column3
17 | AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }}
18 | AZURE_CLIENT_SECRET: ${{ secrets.AZURE_CLIENT_SECRET }}
19 | AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }}
20 |
21 | steps:
22 | # - name: Azure login
23 | # uses: azure/login@v2
24 | # with:
25 | # creds: ${{ secrets.AZURE_CREDENTIALS }}
26 |
27 | - name: Checkout code
28 | uses: actions/checkout@v3
29 |
30 | - name: Set up Go
31 | uses: actions/setup-go@v4
32 | with:
33 | go-version: 1.23
34 |
35 | - name: Install dependencies
36 | run: go mod tidy
37 |
38 | - name: Run tests
39 | run: go test ./...
40 |
```
--------------------------------------------------------------------------------
/tools/query_test.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "os"
6 | "testing"
7 |
8 | "github.com/mark3labs/mcp-go/mcp"
9 | )
10 |
11 | func TestExecuteQueryHandler(t *testing.T) {
12 | ctx := context.Background()
13 |
14 | // Fetch values from environment variables
15 | clusterName := os.Getenv("CLUSTER_NAME")
16 | if clusterName == "" {
17 | t.Fatal("Environment variable CLUSTER_NAME is not set")
18 | }
19 |
20 | dbName := os.Getenv("DB_NAME")
21 | if dbName == "" {
22 | t.Fatal("Environment variable DB_NAME is not set")
23 | }
24 |
25 | tableName := os.Getenv("TABLE_NAME")
26 | if tableName == "" {
27 | t.Fatal("Environment variable TABLE_NAME is not set")
28 | }
29 |
30 | // query := os.Getenv("QUERY")
31 | // if query == "" {
32 | // t.Fatal("Environment variable QUERY is not set")
33 | // }
34 |
35 | query := tableName + " | count"
36 |
37 | request := mcp.CallToolRequest{
38 | Params: struct {
39 | Name string `json:"name"`
40 | Arguments map[string]any `json:"arguments,omitempty"`
41 | Meta *struct {
42 | ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
43 | } `json:"_meta,omitempty"`
44 | }{
45 | Name: "execute_query",
46 | Arguments: map[string]any{
47 | "cluster": clusterName,
48 | "database": dbName,
49 | "query": query,
50 | },
51 | },
52 | }
53 |
54 | result, err := executeQueryHandler(ctx, request)
55 | if err != nil {
56 | t.Fatalf("executeQueryHandler failed: %v", err)
57 | }
58 |
59 | if result == nil {
60 | t.Fatal("Expected result, got nil")
61 | }
62 |
63 | content := result.Content[0].(mcp.TextContent)
64 |
65 | // t.Logf("Content: %s", content.Text)
66 |
67 | if content.Text == "" {
68 | t.Fatal("Expected non-empty content")
69 | }
70 |
71 | //validate content is not empty
72 | if content.Text == "" {
73 | t.Fatal("Expected non-empty content")
74 | }
75 | }
76 |
```
--------------------------------------------------------------------------------
/tools/databases.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "errors"
7 | "fmt"
8 |
9 | "github.com/Azure/azure-kusto-go/azkustodata/kql"
10 | "github.com/abhirockzz/mcp_kusto/common"
11 | "github.com/mark3labs/mcp-go/mcp"
12 | "github.com/mark3labs/mcp-go/server"
13 | )
14 |
15 | func ListDatabases() (mcp.Tool, server.ToolHandlerFunc) {
16 |
17 | return listDatabases(), listDatabasesHandler
18 | }
19 |
20 | func listDatabases() mcp.Tool {
21 |
22 | return mcp.NewTool("list_databases",
23 |
24 | mcp.WithString("cluster",
25 | mcp.Required(),
26 | mcp.Description(CLUSTER_PARAMETER_DESCRIPTION),
27 | ),
28 | mcp.WithDescription("List all databases in a specific Azure Data Explorer cluster"),
29 | )
30 | }
31 |
32 | func listDatabasesHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
33 |
34 | clusterName, ok := request.Params.Arguments["cluster"].(string)
35 | if !ok {
36 | return nil, errors.New("cluster name missing")
37 | }
38 |
39 | client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
40 | if err != nil {
41 | return nil, err
42 | }
43 | defer client.Close()
44 |
45 | // Use .show databases command
46 | dataset, err := client.Mgmt(ctx, "", kql.New(".show databases"))
47 | if err != nil {
48 | return nil, err
49 | }
50 |
51 | databaseNames := []string{}
52 |
53 | // Process the results
54 | for _, row := range dataset.Tables()[0].Rows() {
55 | // Access database name by column name
56 | databaseName, err := row.StringByName("DatabaseName")
57 | if err != nil {
58 | return nil, err
59 | }
60 | //fmt.Println("Database:", databaseName)
61 | databaseNames = append(databaseNames, databaseName)
62 | }
63 |
64 | var result ListDatabasesResponse
65 |
66 | result.Databases = databaseNames
67 |
68 | jsonResult, err := json.Marshal(result)
69 | if err != nil {
70 | return nil, err
71 | }
72 |
73 | return mcp.NewToolResultText(string(jsonResult)), nil
74 | }
75 |
76 | type ListDatabasesResponse struct {
77 | Databases []string `json:"databases"`
78 | }
79 |
```
--------------------------------------------------------------------------------
/common/client.go:
--------------------------------------------------------------------------------
```go
1 | package common
2 |
3 | import (
4 | "github.com/Azure/azure-kusto-go/azkustodata"
5 | )
6 |
7 | func GetClient(endpoint string) (*azkustodata.Client, error) {
8 | // Create a connection string builder with authentication
9 | kustoConnectionString := azkustodata.NewConnectionStringBuilder(endpoint).WithDefaultAzureCredential()
10 |
11 | // Initialize the client
12 | client, err := azkustodata.New(kustoConnectionString)
13 | if err != nil {
14 | return nil, err
15 | }
16 | return client, nil
17 | }
18 |
19 | // const clusterNameFormat = "https://%s.kusto.windows.net/"
20 |
21 | // func CreateTable(clusterName, dbName, createTableCommand string) error {
22 |
23 | // endpoint := fmt.Sprintf(clusterNameFormat, clusterName)
24 | // fmt.Println("Cluster URL:", endpoint)
25 |
26 | // // Initialize the client
27 | // client, err := GetClient(endpoint)
28 | // if err != nil {
29 | // return fmt.Errorf("error creating Kusto client: %w", err)
30 | // }
31 | // defer client.Close()
32 |
33 | // ctx := context.Background()
34 |
35 | // _, err = client.Mgmt(ctx, dbName, kql.New("").AddUnsafe(createTableCommand))
36 | // if err != nil {
37 | // return fmt.Errorf("error executing create table: %w", err)
38 | // }
39 |
40 | // log.Println("table created successfully with command", createTableCommand)
41 | // return nil
42 | // }
43 |
44 | // func DropTable(clusterName, dbName, tableName string) error {
45 | // endpoint := fmt.Sprintf(clusterNameFormat, clusterName)
46 | // fmt.Println("Cluster URL:", endpoint)
47 |
48 | // // Initialize the client
49 | // client, err := GetClient(endpoint)
50 | // if err != nil {
51 | // return fmt.Errorf("error creating Kusto client: %w", err)
52 | // }
53 | // defer client.Close()
54 |
55 | // ctx := context.Background()
56 |
57 | // _, err = client.Mgmt(ctx, dbName, kql.New(". drop table ").AddUnsafe(tableName))
58 | // if err != nil {
59 | // return fmt.Errorf("error executing drop table command: %w", err)
60 | // }
61 |
62 | // log.Println("table dropped successfully", tableName)
63 | // return nil
64 | // }
65 |
```
--------------------------------------------------------------------------------
/tools/databases_test.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "os"
7 | "testing"
8 |
9 | "github.com/mark3labs/mcp-go/mcp"
10 | )
11 |
12 | func TestListDatabasesHandler(t *testing.T) {
13 | ctx := context.Background()
14 |
15 | clusterName := os.Getenv("CLUSTER_NAME")
16 | if clusterName == "" {
17 | t.Fatal("Environment variable CLUSTER_NAME is not set")
18 | }
19 |
20 | dbName := os.Getenv("DB_NAME")
21 | if dbName == "" {
22 | t.Fatal("Environment variable DB_NAME is not set")
23 | }
24 |
25 | request := mcp.CallToolRequest{
26 | Params: struct {
27 | Name string `json:"name"`
28 | Arguments map[string]any `json:"arguments,omitempty"`
29 | Meta *struct {
30 | ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
31 | } `json:"_meta,omitempty"`
32 | }{
33 | Name: "list_databases",
34 | Arguments: map[string]any{
35 | "cluster": clusterName,
36 | },
37 | },
38 | }
39 |
40 | result, err := listDatabasesHandler(ctx, request)
41 | if err != nil {
42 | t.Fatalf("listDatabasesHandler failed: %v", err)
43 | }
44 |
45 | if result == nil {
46 | t.Fatal("Expected result, got nil")
47 | }
48 |
49 | content := result.Content[0].(mcp.TextContent)
50 |
51 | t.Logf("Content: %s", content.Text)
52 |
53 | if content.Text == "" {
54 | t.Fatal("Expected non-empty content")
55 | }
56 |
57 | // Unmarshal the content
58 | var output ListDatabasesResponse
59 | if err := json.Unmarshal([]byte(content.Text), &output); err != nil {
60 | t.Fatalf("Failed to unmarshal content: %v", err)
61 | }
62 |
63 | if len(output.Databases) == 0 {
64 | t.Fatal("Expected 'databases' key in unmarshaled output with non-empty slice")
65 | }
66 |
67 | if output.Databases[0] != dbName {
68 | t.Fatalf("Expected database name %s, got %s", dbName, output.Databases[0])
69 | }
70 | }
71 |
72 | // func createDatabase(clusterName, dbName string) error {
73 |
74 | // if dbName == "" {
75 | // return errors.New("database name cannot be empty")
76 | // }
77 |
78 | // client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
79 | // if err != nil {
80 | // return fmt.Errorf("failed to get client: %v", err)
81 | // }
82 |
83 | // _, err = client.Mgmt(context.Background(), "", kql.New("").AddUnsafe(".create database "+dbName))
84 | // if err != nil {
85 | // return fmt.Errorf("failed to create database: %v", err)
86 | // }
87 | // return nil
88 | // }
89 |
90 | // func deleteDatabase(clusterName, dbName string) error {
91 |
92 | // if dbName == "" {
93 | // return errors.New("database name cannot be empty")
94 | // }
95 |
96 | // client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
97 | // if err != nil {
98 | // return fmt.Errorf("failed to get client: %v", err)
99 | // }
100 |
101 | // _, err = client.Mgmt(context.Background(), "", kql.New(".drop database ").AddUnsafe(dbName))
102 | // if err != nil {
103 | // return fmt.Errorf("failed to drop database: %v", err)
104 | // }
105 |
106 | // return nil
107 | // }
108 |
```
--------------------------------------------------------------------------------
/tools/query.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "errors"
6 | "fmt"
7 |
8 | "github.com/Azure/azure-kusto-go/azkustodata/kql"
9 | "github.com/abhirockzz/mcp_kusto/common"
10 | "github.com/mark3labs/mcp-go/mcp"
11 | "github.com/mark3labs/mcp-go/server"
12 | )
13 |
14 | func ExecuteQuery() (mcp.Tool, server.ToolHandlerFunc) {
15 |
16 | return executeQuery(), executeQueryHandler
17 | }
18 |
19 | func executeQuery() mcp.Tool {
20 |
21 | return mcp.NewTool("execute_query",
22 |
23 | mcp.WithString("cluster",
24 | mcp.Required(),
25 | mcp.Description(CLUSTER_PARAMETER_DESCRIPTION),
26 | ),
27 | mcp.WithString("database",
28 | mcp.Required(),
29 | mcp.Description("Name of the database."),
30 | ),
31 |
32 | // mcp.WithString("table",
33 | // mcp.Required(),
34 | // mcp.Description("Name of the table."),
35 | // ),
36 | mcp.WithString("query",
37 | mcp.Required(),
38 | mcp.Description("The query to execute."),
39 | ),
40 | mcp.WithDescription("Execute a read-only query. Ask the user for permission before executing the query. It has to be a valid KQL query. Write queries are not allowed. Result truncation is a limit set by default on the result set returned by the query. Kusto limits the number of records returned to the client to 500,000, and the overall data size for those records to 64 MB. When either of these limits is exceeded, the query fails with a partial query failure. Exceeding these limits will generate an exception. Reduce the result set size by modifying the query to only return interesting data. There are several strategies to avoid this. 1/ Use the summarize operator group and aggregate over similar records in the query output. 2/ Potentially sample some columns by using the take_any aggregation function. 3/ Use a take operator to sample the query output. 4/Use the substring function to trim wide free-text columns. 5/ Use the project operator to drop any uninteresting column from the result set."),
41 | )
42 | }
43 |
44 | func executeQueryHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
45 |
46 | clusterName, ok := request.Params.Arguments["cluster"].(string)
47 | if !ok {
48 | return nil, errors.New("cluster name missing")
49 | }
50 |
51 | dbName, ok := request.Params.Arguments["database"].(string)
52 | if !ok {
53 | return nil, errors.New("database name missing")
54 | }
55 |
56 | // table, ok := request.Params.Arguments["table"].(string)
57 | // if !ok {
58 | // return nil, errors.New("table name missing")
59 | // }
60 |
61 | query, ok := request.Params.Arguments["query"].(string)
62 | if !ok {
63 | return nil, errors.New("query missing")
64 | }
65 |
66 | client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
67 | if err != nil {
68 | return nil, err
69 | }
70 | defer client.Close()
71 |
72 | stmt := kql.New("").AddUnsafe(query)
73 |
74 | queryResponse, err := client.QueryToJson(context.Background(), dbName, stmt)
75 | if err != nil {
76 | return nil, err
77 | }
78 |
79 | return mcp.NewToolResultText(queryResponse), nil
80 | }
81 |
```
--------------------------------------------------------------------------------
/tools/tables_test.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "os"
7 | "strings"
8 | "testing"
9 |
10 | "slices"
11 |
12 | "github.com/mark3labs/mcp-go/mcp"
13 | )
14 |
15 | func TestListTablesHandler(t *testing.T) {
16 | ctx := context.Background()
17 |
18 | // Fetch values from environment variables
19 | clusterName := os.Getenv("CLUSTER_NAME")
20 | if clusterName == "" {
21 | t.Fatal("Environment variable CLUSTER_NAME is not set")
22 | }
23 |
24 | dbName := os.Getenv("DB_NAME")
25 | if dbName == "" {
26 | t.Fatal("Environment variable DB_NAME is not set")
27 | }
28 |
29 | tableName := os.Getenv("TABLE_NAME")
30 | if tableName == "" {
31 | t.Fatal("Environment variable TABLE_NAME is not set")
32 | }
33 |
34 | request := mcp.CallToolRequest{
35 | Params: struct {
36 | Name string `json:"name"`
37 | Arguments map[string]any `json:"arguments,omitempty"`
38 | Meta *struct {
39 | ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
40 | } `json:"_meta,omitempty"`
41 | }{
42 | Name: "list_tables",
43 | Arguments: map[string]any{
44 | "cluster": clusterName,
45 | "database": dbName,
46 | },
47 | },
48 | }
49 |
50 | result, err := listTablesHandler(ctx, request)
51 | if err != nil {
52 | t.Fatalf("listTablesHandler failed: %v", err)
53 | }
54 |
55 | if result == nil {
56 | t.Fatal("Expected result, got nil")
57 | }
58 |
59 | content := result.Content[0].(mcp.TextContent)
60 |
61 | t.Logf("Content: %s", content.Text)
62 |
63 | if content.Text == "" {
64 | t.Fatal("Expected non-empty content")
65 | }
66 |
67 | // Unmarshal the content to check for the response struct
68 | var output ListTablesResponse
69 | if err := json.Unmarshal([]byte(content.Text), &output); err != nil {
70 | t.Fatalf("Failed to unmarshal content: %v", err)
71 | }
72 |
73 | // Validate the result contains the expected cluster name
74 | expectedClusterName := clusterName
75 | if output.Cluster != expectedClusterName {
76 | t.Fatalf("Expected cluster name '%s', but got '%v'", expectedClusterName, output.Cluster)
77 | }
78 |
79 | // Validate the result contains the expected database name
80 | expectedDatabaseName := dbName
81 | if output.Database != expectedDatabaseName {
82 | t.Fatalf("Expected database name '%s', but got '%v'", expectedDatabaseName, output.Database)
83 | }
84 |
85 | // Validate the result contains the expected table name
86 | expectedTableName := tableName
87 | found := slices.Contains(output.Tables, expectedTableName)
88 | if !found {
89 | t.Fatalf("Expected table name '%s' not found in tables", expectedTableName)
90 | }
91 | }
92 |
93 | func TestGetSchemaHandler(t *testing.T) {
94 | ctx := context.Background()
95 |
96 | clusterName := os.Getenv("CLUSTER_NAME")
97 | if clusterName == "" {
98 | t.Fatal("Environment variable CLUSTER_NAME is not set")
99 | }
100 |
101 | dbName := os.Getenv("DB_NAME")
102 | if dbName == "" {
103 | t.Fatal("Environment variable DB_NAME is not set")
104 | }
105 |
106 | tableName := os.Getenv("TABLE_NAME")
107 | if tableName == "" {
108 | t.Fatal("Environment variable TABLE_NAME is not set")
109 | }
110 |
111 | // columnName := os.Getenv("COLUMN_NAME")
112 | // if columnName == "" {
113 | // t.Fatal("Environment variable COLUMN_NAME is not set")
114 | // }
115 |
116 | request := mcp.CallToolRequest{
117 | Params: struct {
118 | Name string `json:"name"`
119 | Arguments map[string]any `json:"arguments,omitempty"`
120 | Meta *struct {
121 | ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
122 | } `json:"_meta,omitempty"`
123 | }{
124 | Name: "get_table_schema",
125 | Arguments: map[string]any{
126 | "cluster": clusterName,
127 | "database": dbName,
128 | "table": tableName,
129 | },
130 | },
131 | }
132 |
133 | // Call the handler
134 | result, err := getSchemaHandler(ctx, request)
135 | if err != nil {
136 | t.Fatalf("Handler returned an error: %v", err)
137 | }
138 |
139 | // Validate the result
140 | if result == nil {
141 | t.Fatal("Result is nil")
142 | }
143 |
144 | schema := result.Content[0].(mcp.TextContent)
145 | if schema.Text == "" {
146 | t.Fatal("Schema text is empty")
147 | }
148 | t.Logf("Schema text: %s", schema.Text)
149 |
150 | // Verify the schema response
151 | var schemaResponse TableSchemaResponse
152 | if err := json.Unmarshal([]byte(schema.Text), &schemaResponse); err != nil {
153 | t.Fatalf("Failed to unmarshal schema response: %v", err)
154 | }
155 |
156 | // Check table name
157 | if schemaResponse.Name != tableName {
158 | t.Fatalf("Expected table name '%s', but got '%v'", tableName, schemaResponse.Name)
159 | }
160 |
161 | // Check ordered columns
162 | expectedColumnNames := os.Getenv("COLUMN_NAMES")
163 | if expectedColumnNames == "" {
164 | t.Fatal("Environment variable COLUMN_NAMES is not set")
165 | }
166 |
167 | expectedColumns := make(map[string]bool)
168 | for _, col := range strings.Split(expectedColumnNames, ",") {
169 | expectedColumns[col] = true
170 | }
171 |
172 | for _, col := range schemaResponse.OrderedColumns {
173 | delete(expectedColumns, col.Name)
174 | }
175 |
176 | if len(expectedColumns) > 0 {
177 | t.Fatalf("Expected column names not found: %v", expectedColumns)
178 | }
179 | }
180 |
```
--------------------------------------------------------------------------------
/tools/tables.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "errors"
7 | "fmt"
8 |
9 | "github.com/Azure/azure-kusto-go/azkustodata/kql"
10 | "github.com/abhirockzz/mcp_kusto/common"
11 | "github.com/mark3labs/mcp-go/mcp"
12 | "github.com/mark3labs/mcp-go/server"
13 | )
14 |
15 | func ListTables() (mcp.Tool, server.ToolHandlerFunc) {
16 |
17 | return listTables(), listTablesHandler
18 | }
19 |
20 | // listTables returns a tool that lists all tables in a specific Azure Data Explorer database.
21 | func listTables() mcp.Tool {
22 |
23 | return mcp.NewTool("list_tables",
24 |
25 | mcp.WithString("cluster",
26 | mcp.Required(),
27 | mcp.Description(CLUSTER_PARAMETER_DESCRIPTION),
28 | ),
29 | mcp.WithString("database",
30 | mcp.Required(),
31 | mcp.Description("Name of the database to list tables from."),
32 | ),
33 | mcp.WithDescription("List all tables in a specific Azure Data Explorer database"),
34 | )
35 | }
36 |
37 | // Define a struct to represent the response for listTablesHandler
38 | // This struct will replace the map currently used
39 |
40 | type ListTablesResponse struct {
41 | Cluster string `json:"cluster"`
42 | Database string `json:"database"`
43 | Tables []string `json:"tables"`
44 | }
45 |
46 | // listTablesHandler handles the request to list all tables in a specific Azure Data Explorer database.
47 | func listTablesHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
48 |
49 | clusterName, ok := request.Params.Arguments["cluster"].(string)
50 | if !ok {
51 | return nil, errors.New("cluster name missing")
52 | }
53 |
54 | dbName, ok := request.Params.Arguments["database"].(string)
55 | if !ok {
56 | return nil, errors.New("database name missing")
57 | }
58 |
59 | client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
60 | if err != nil {
61 | return nil, err
62 | }
63 | defer client.Close()
64 |
65 | dataset, err := client.Mgmt(ctx, dbName, kql.New(".show tables"))
66 | if err != nil {
67 | return nil, err
68 | }
69 |
70 | tableNames := []string{}
71 |
72 | // Process the results
73 | for _, row := range dataset.Tables()[0].Rows() {
74 | // Access table name by column name
75 | tableName, err := row.StringByName("TableName")
76 | if err != nil {
77 | return nil, err
78 | }
79 | tableNames = append(tableNames, tableName)
80 | }
81 |
82 | response := ListTablesResponse{
83 | Cluster: clusterName,
84 | Database: dbName,
85 | Tables: tableNames,
86 | }
87 |
88 | jsonResult, err := json.Marshal(response)
89 | if err != nil {
90 | return nil, err
91 | }
92 |
93 | return mcp.NewToolResultText(string(jsonResult)), nil
94 | }
95 |
96 | // GetTableSchema returns a tool that retrieves the schema of a specific table in an Azure Data Explorer database.
97 | func GetTableSchema() (mcp.Tool, server.ToolHandlerFunc) {
98 |
99 | return getSchema(), getSchemaHandler
100 | }
101 |
102 | // getSchema returns a tool that retrieves the schema of a specific table in an Azure Data Explorer database.
103 | func getSchema() mcp.Tool {
104 |
105 | return mcp.NewTool("get_table_schema",
106 |
107 | mcp.WithString("cluster",
108 | mcp.Required(),
109 | mcp.Description(CLUSTER_PARAMETER_DESCRIPTION),
110 | ),
111 | mcp.WithString("database",
112 | mcp.Required(),
113 | mcp.Description("Name of the database."),
114 | ),
115 |
116 | mcp.WithString("table",
117 | mcp.Required(),
118 | mcp.Description("Name of the table to get the schema for."),
119 | ),
120 | mcp.WithDescription("Get the schema of a specific table in an Azure Data Explorer database"),
121 | )
122 | }
123 |
124 | // Define a struct to represent the schema response
125 | // This is to aid testing
126 |
127 | type TableSchemaResponse struct {
128 | Name string `json:"Name"`
129 | OrderedColumns []struct {
130 | Name string `json:"Name"`
131 | Type string `json:"Type"`
132 | CslType string `json:"CslType"`
133 | } `json:"OrderedColumns"`
134 | }
135 |
136 | func getSchemaHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
137 |
138 | clusterName, ok := request.Params.Arguments["cluster"].(string)
139 | if !ok {
140 | return nil, errors.New("cluster name missing")
141 | }
142 |
143 | dbName, ok := request.Params.Arguments["database"].(string)
144 | if !ok {
145 | return nil, errors.New("database name missing")
146 | }
147 |
148 | table, ok := request.Params.Arguments["table"].(string)
149 | if !ok {
150 | return nil, errors.New("table name missing")
151 | }
152 |
153 | client, err := common.GetClient(fmt.Sprintf(clusterNameFormat, clusterName))
154 | if err != nil {
155 | return nil, err
156 | }
157 | defer client.Close()
158 |
159 | command := kql.New(".show table ").AddTable(table).AddLiteral(" schema as json")
160 |
161 | //fmt.Println("Command:", command.String())
162 |
163 | dataset, err := client.Mgmt(ctx, dbName, command)
164 | if err != nil {
165 | return nil, err
166 | }
167 |
168 | // Process the schema information
169 | //fmt.Println("Schema for table", table)
170 | jsonSchema, err := dataset.Tables()[0].Rows()[0].StringByName("Schema")
171 |
172 | if err != nil {
173 | return nil, err
174 | }
175 |
176 | // var schemaResponse TableSchemaResponse
177 | // err = json.Unmarshal([]byte(jsonSchema), &schemaResponse)
178 | // if err != nil {
179 | // return nil, err
180 | // }
181 |
182 | // responseJSON, err := json.Marshal(schemaResponse)
183 | // if err != nil {
184 | // return nil, err
185 | // }
186 |
187 | //return mcp.NewToolResultText(string(responseJSON)), nil
188 |
189 | return mcp.NewToolResultText(jsonSchema), nil
190 | }
191 |
```