From f16d322501af94ec9c9a48ed484e78edabff0f71 Mon Sep 17 00:00:00 2001 From: Shalev Shalit Date: Tue, 19 Aug 2025 18:21:31 +0300 Subject: [PATCH] feat(tools): Add ClickHouse support with querying and metadata retrieval - Updated README to include ClickHouse as a supported datasource. - Implemented ClickHouse querying capabilities, including SQL execution, database and table listing, and table description. - Added integration and unit tests for ClickHouse functionalities. - Enhanced the disabled tools configuration to include ClickHouse options. --- README.md | 12 +- cmd/mcp-grafana/main.go | 8 +- tools/clickhouse.go | 517 ++++++++++++++++++++++++++++++++++ tools/clickhouse_test.go | 225 +++++++++++++++ tools/clickhouse_unit_test.go | 77 +++++ 5 files changed, 835 insertions(+), 4 deletions(-) create mode 100644 tools/clickhouse.go create mode 100644 tools/clickhouse_test.go create mode 100644 tools/clickhouse_unit_test.go diff --git a/README.md b/README.md index 97853645..8c2cd62f 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ The dashboard tools now include several strategies to manage context window usag ### Datasources - **List and fetch datasource information:** View all configured datasources and retrieve detailed information about each. - - _Supported datasource types: Prometheus, Loki._ + - _Supported datasource types: Prometheus, Loki, ClickHouse._ ### Prometheus Querying @@ -51,6 +51,12 @@ The dashboard tools now include several strategies to manage context window usag - **Query Loki logs and metrics:** Run both log queries and metric queries using LogQL against Loki datasources. - **Query Loki metadata:** Retrieve label names, label values, and stream statistics from Loki datasources. +### ClickHouse Querying + +- **Query ClickHouse:** Execute SQL queries against ClickHouse datasources (SELECT, SHOW, and DESCRIBE queries only for safety). +- **Query ClickHouse metadata:** Retrieve database lists, table information with detailed metadata (engine, row count, size), and table schemas. +- **Database exploration:** List databases, tables with filtering support, and describe table structures including column types and comments. + ### Incidents - **Search, create, and update incidents:** Manage incidents in Grafana Incident, including searching, creating, and adding activities to incidents. @@ -177,6 +183,10 @@ Scopes define the specific resources that permissions apply to. Each action requ | `list_loki_label_names` | Loki | List all available label names in logs | `datasources:query` | `datasources:uid:loki-uid` | | `list_loki_label_values` | Loki | List values for a specific log label | `datasources:query` | `datasources:uid:loki-uid` | | `query_loki_stats` | Loki | Get statistics about log streams | `datasources:query` | `datasources:uid:loki-uid` | +| `query_clickhouse` | ClickHouse | Execute SQL queries against a ClickHouse datasource | `datasources:query` | `datasources:uid:clickhouse-uid` | +| `list_clickhouse_databases` | ClickHouse | List all available databases in a ClickHouse datasource | `datasources:query` | `datasources:uid:clickhouse-uid` | +| `list_clickhouse_tables` | ClickHouse | List tables in a database with detailed metadata | `datasources:query` | `datasources:uid:clickhouse-uid` | +| `describe_clickhouse_table` | ClickHouse | Describe table structure with column types and comments | `datasources:query` | `datasources:uid:clickhouse-uid` | | `list_alert_rules` | Alerting | List alert rules | `alert.rules:read` | `folders:*` or `folders:uid:alerts-folder` | | `get_alert_rule_by_uid` | Alerting | Get alert rule by UID | `alert.rules:read` | `folders:uid:alerts-folder` | | `list_contact_points` | Alerting | List notification contact points | `alert.notifications:read` | Global scope | diff --git a/cmd/mcp-grafana/main.go b/cmd/mcp-grafana/main.go index befea298..eab230b8 100644 --- a/cmd/mcp-grafana/main.go +++ b/cmd/mcp-grafana/main.go @@ -33,7 +33,7 @@ type disabledTools struct { enabledTools string search, datasource, incident, - prometheus, loki, alerting, + prometheus, loki, clickhouse, alerting, dashboard, oncall, asserts, sift, admin, pyroscope, navigation bool } @@ -51,13 +51,14 @@ type grafanaConfig struct { } func (dt *disabledTools) addFlags() { - flag.StringVar(&dt.enabledTools, "enabled-tools", "search,datasource,incident,prometheus,loki,alerting,dashboard,oncall,asserts,sift,admin,pyroscope,navigation", "A comma separated list of tools enabled for this server. Can be overwritten entirely or by disabling specific components, e.g. --disable-search.") + flag.StringVar(&dt.enabledTools, "enabled-tools", "search,datasource,incident,prometheus,loki,clickhouse,alerting,dashboard,oncall,asserts,sift,admin,pyroscope,navigation", "A comma separated list of tools enabled for this server. Can be overwritten entirely or by disabling specific components, e.g. --disable-search.") flag.BoolVar(&dt.search, "disable-search", false, "Disable search tools") flag.BoolVar(&dt.datasource, "disable-datasource", false, "Disable datasource tools") flag.BoolVar(&dt.incident, "disable-incident", false, "Disable incident tools") flag.BoolVar(&dt.prometheus, "disable-prometheus", false, "Disable prometheus tools") flag.BoolVar(&dt.loki, "disable-loki", false, "Disable loki tools") + flag.BoolVar(&dt.clickhouse, "disable-clickhouse", false, "Disable clickhouse tools") flag.BoolVar(&dt.alerting, "disable-alerting", false, "Disable alerting tools") flag.BoolVar(&dt.dashboard, "disable-dashboard", false, "Disable dashboard tools") flag.BoolVar(&dt.oncall, "disable-oncall", false, "Disable oncall tools") @@ -85,6 +86,7 @@ func (dt *disabledTools) addTools(s *server.MCPServer) { maybeAddTools(s, tools.AddIncidentTools, enabledTools, dt.incident, "incident") maybeAddTools(s, tools.AddPrometheusTools, enabledTools, dt.prometheus, "prometheus") maybeAddTools(s, tools.AddLokiTools, enabledTools, dt.loki, "loki") + maybeAddTools(s, tools.AddClickHouseTools, enabledTools, dt.clickhouse, "clickhouse") maybeAddTools(s, tools.AddAlertingTools, enabledTools, dt.alerting, "alerting") maybeAddTools(s, tools.AddDashboardTools, enabledTools, dt.dashboard, "dashboard") maybeAddTools(s, tools.AddOnCallTools, enabledTools, dt.oncall, "oncall") @@ -102,7 +104,7 @@ func newServer(dt disabledTools) *server.MCPServer { Available Capabilities: - Dashboards: Search, retrieve, update, and create dashboards. Extract panel queries and datasource information. - Datasources: List and fetch details for datasources. - - Prometheus & Loki: Run PromQL and LogQL queries, retrieve metric/log metadata, and explore label names/values. + - Prometheus & Loki & ClickHouse: Run PromQL, LogQL, and SQL queries, retrieve metric/log/table metadata, and explore label names/values and database schemas. - Incidents: Search, create, update, and resolve incidents in Grafana Incident. - Sift Investigations: Start and manage Sift investigations, analyze logs/traces, find error patterns, and detect slow requests. - Alerting: List and fetch alert rules and notification contact points. diff --git a/tools/clickhouse.go b/tools/clickhouse.go new file mode 100644 index 00000000..93d18e24 --- /dev/null +++ b/tools/clickhouse.go @@ -0,0 +1,517 @@ +package tools + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + + mcpgrafana "github.com/grafana/mcp-grafana" + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" +) + +const ( + // DefaultClickHouseLimit is the default number of rows to return if not specified + DefaultClickHouseLimit = 100 + + // MaxClickHouseLimit is the maximum number of rows that can be requested + MaxClickHouseLimit = 1000 +) + +// ClickHouseClient represents a client for ClickHouse HTTP interface +type ClickHouseClient struct { + httpClient *http.Client + baseURL string +} + +// ClickHouseTable represents a ClickHouse table with metadata +type ClickHouseTable struct { + Database string `json:"database"` + Name string `json:"name"` + Engine string `json:"engine"` + TotalRows uint64 `json:"total_rows"` + TotalBytes uint64 `json:"total_bytes"` + TotalBytesUncompressed uint64 `json:"total_bytes_uncompressed"` + Parts uint64 `json:"parts"` + ActiveParts uint64 `json:"active_parts"` + Comment string `json:"comment"` +} + +// ClickHouseColumn represents a ClickHouse column definition +type ClickHouseColumn struct { + Database string `json:"database"` + Table string `json:"table"` + Name string `json:"name"` + Type string `json:"type"` + DefaultKind string `json:"default_kind"` + DefaultExpression string `json:"default_expression"` + Comment string `json:"comment"` +} + +// ClickHouseQueryResult represents the result of a ClickHouse query +type ClickHouseQueryResult struct { + Columns []string `json:"columns"` + Rows [][]interface{} `json:"rows"` + Summary struct { + ReadRows uint64 `json:"read_rows"` + ReadBytes uint64 `json:"read_bytes"` + Written uint64 `json:"written_rows"` + } `json:"summary,omitempty"` +} + +// newClickHouseClient creates a new ClickHouse client +func newClickHouseClient(ctx context.Context, uid string) (*ClickHouseClient, error) { + // First check if the datasource exists + _, err := getDatasourceByUID(ctx, GetDatasourceByUIDParams{UID: uid}) + if err != nil { + return nil, err + } + + cfg := mcpgrafana.GrafanaConfigFromContext(ctx) + url := fmt.Sprintf("%s/api/datasources/proxy/uid/%s", strings.TrimRight(cfg.URL, "/"), uid) + + // Create custom transport with TLS configuration if available + var transport = http.DefaultTransport + if tlsConfig := cfg.TLSConfig; tlsConfig != nil { + var err error + transport, err = tlsConfig.HTTPTransport(transport.(*http.Transport)) + if err != nil { + return nil, fmt.Errorf("failed to create custom transport: %w", err) + } + } + + authTransport := &clickhouseAuthRoundTripper{ + accessToken: cfg.AccessToken, + idToken: cfg.IDToken, + apiKey: cfg.APIKey, + underlying: transport, + } + + client := &http.Client{ + Transport: mcpgrafana.NewUserAgentTransport( + authTransport, + ), + } + + return &ClickHouseClient{ + httpClient: client, + baseURL: url, + }, nil +} + +// clickhouseAuthRoundTripper handles authentication for ClickHouse requests +type clickhouseAuthRoundTripper struct { + accessToken string + idToken string + apiKey string + underlying http.RoundTripper +} + +func (rt *clickhouseAuthRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if rt.accessToken != "" && rt.idToken != "" { + req.Header.Set("X-Access-Token", rt.accessToken) + req.Header.Set("X-Grafana-Id", rt.idToken) + } else if rt.apiKey != "" { + req.Header.Set("Authorization", "Bearer "+rt.apiKey) + } + + resp, err := rt.underlying.RoundTrip(req) + if err != nil { + return nil, err + } + + return resp, nil +} + +// buildURL constructs a full URL for a ClickHouse HTTP endpoint +func (c *ClickHouseClient) buildURL(params url.Values) string { + fullURL := c.baseURL + if params != nil { + fullURL += "?" + params.Encode() + } + return fullURL +} + +// executeQuery executes a ClickHouse query and returns the response +func (c *ClickHouseClient) executeQuery(ctx context.Context, query string, format string, limit int) ([]byte, error) { + params := url.Values{} + params.Add("query", query) + if format != "" { + params.Add("format", format) + } + if limit > 0 { + // ClickHouse LIMIT clause should be part of the query, but we can also use settings + params.Add("max_result_rows", strconv.Itoa(limit)) + } + + fullURL := c.buildURL(params) + + req, err := http.NewRequestWithContext(ctx, "POST", fullURL, nil) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("executing request: %w", err) + } + defer func() { + _ = resp.Body.Close() //nolint:errcheck + }() + + // Check for non-200 status code + if resp.StatusCode != http.StatusOK { + bodyBytes, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("ClickHouse returned status code %d: %s", resp.StatusCode, string(bodyBytes)) + } + + // Read the response body with a limit to prevent memory issues + body := io.LimitReader(resp.Body, 1024*1024*50) // 50MB limit + bodyBytes, err := io.ReadAll(body) + if err != nil { + return nil, fmt.Errorf("reading response body: %w", err) + } + + return bytes.TrimSpace(bodyBytes), nil +} + +// queryJSON executes a query and returns JSON results +func (c *ClickHouseClient) queryJSON(ctx context.Context, query string, limit int) (*ClickHouseQueryResult, error) { + bodyBytes, err := c.executeQuery(ctx, query, "JSONCompact", limit) + if err != nil { + return nil, err + } + + if len(bodyBytes) == 0 { + return &ClickHouseQueryResult{Columns: []string{}, Rows: [][]interface{}{}}, nil + } + + var jsonResult struct { + Meta []struct { + Name string `json:"name"` + Type string `json:"type"` + } `json:"meta"` + Data [][]interface{} `json:"data"` + Statistics struct { + Elapsed float64 `json:"elapsed"` + RowsRead uint64 `json:"rows_read"` + BytesRead uint64 `json:"bytes_read"` + } `json:"statistics,omitempty"` + } + + err = json.Unmarshal(bodyBytes, &jsonResult) + if err != nil { + return nil, fmt.Errorf("unmarshalling ClickHouse response: %w", err) + } + + // Extract column names + columns := make([]string, len(jsonResult.Meta)) + for i, col := range jsonResult.Meta { + columns[i] = col.Name + } + + result := &ClickHouseQueryResult{ + Columns: columns, + Rows: jsonResult.Data, + } + + if jsonResult.Statistics.RowsRead > 0 { + result.Summary.ReadRows = jsonResult.Statistics.RowsRead + result.Summary.ReadBytes = jsonResult.Statistics.BytesRead + } + + return result, nil +} + +// queryPlainText executes a query and returns plain text results +func (c *ClickHouseClient) queryPlainText(ctx context.Context, query string) ([]string, error) { + bodyBytes, err := c.executeQuery(ctx, query, "TSV", 0) + if err != nil { + return nil, err + } + + if len(bodyBytes) == 0 { + return []string{}, nil + } + + lines := strings.Split(string(bodyBytes), "\n") + // Remove empty lines + var result []string + for _, line := range lines { + line = strings.TrimSpace(line) + if line != "" { + result = append(result, line) + } + } + + return result, nil +} + +// QueryClickHouseParams defines the parameters for querying ClickHouse +type QueryClickHouseParams struct { + DatasourceUID string `json:"datasourceUid" jsonschema:"required,description=The UID of the datasource to query"` + Query string `json:"query" jsonschema:"required,description=The SQL query to execute against ClickHouse. Must be a SELECT query for safety."` + Limit int `json:"limit,omitempty" jsonschema:"description=Optionally\\, the maximum number of rows to return (default: 100\\, max: 1000)"` +} + +// enforceLimit ensures a limit value is within acceptable bounds +func enforceClickHouseLimit(requestedLimit int) int { + if requestedLimit <= 0 { + return DefaultClickHouseLimit + } + if requestedLimit > MaxClickHouseLimit { + return MaxClickHouseLimit + } + return requestedLimit +} + +// queryClickHouse executes a SQL query against ClickHouse +func queryClickHouse(ctx context.Context, args QueryClickHouseParams) (*ClickHouseQueryResult, error) { + client, err := newClickHouseClient(ctx, args.DatasourceUID) + if err != nil { + return nil, fmt.Errorf("creating ClickHouse client: %w", err) + } + + // Basic validation - ensure it's a SELECT query + trimmedQuery := strings.TrimSpace(strings.ToUpper(args.Query)) + if !strings.HasPrefix(trimmedQuery, "SELECT") && !strings.HasPrefix(trimmedQuery, "SHOW") && !strings.HasPrefix(trimmedQuery, "DESCRIBE") { + return nil, fmt.Errorf("only SELECT, SHOW, and DESCRIBE queries are allowed for safety") + } + + // Apply limit constraints + limit := enforceClickHouseLimit(args.Limit) + + result, err := client.queryJSON(ctx, args.Query, limit) + if err != nil { + return nil, err + } + + return result, nil +} + +// QueryClickHouse is a tool for querying ClickHouse +var QueryClickHouse = mcpgrafana.MustTool( + "query_clickhouse", + "Executes a SQL query against a ClickHouse datasource. Only SELECT, SHOW, and DESCRIBE queries are allowed for safety. Returns column names, data rows, and query statistics. Supports a configurable row limit (default: 100, max: 1000).", + queryClickHouse, + mcp.WithTitleAnnotation("Query ClickHouse"), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithReadOnlyHintAnnotation(true), +) + +// ListClickHouseDatabasesParams defines the parameters for listing ClickHouse databases +type ListClickHouseDatabasesParams struct { + DatasourceUID string `json:"datasourceUid" jsonschema:"required,description=The UID of the datasource to query"` +} + +// listClickHouseDatabases lists all databases in a ClickHouse datasource +func listClickHouseDatabases(ctx context.Context, args ListClickHouseDatabasesParams) ([]string, error) { + client, err := newClickHouseClient(ctx, args.DatasourceUID) + if err != nil { + return nil, fmt.Errorf("creating ClickHouse client: %w", err) + } + + databases, err := client.queryPlainText(ctx, "SHOW DATABASES") + if err != nil { + return nil, err + } + + return databases, nil +} + +// ListClickHouseDatabases is a tool for listing ClickHouse databases +var ListClickHouseDatabases = mcpgrafana.MustTool( + "list_clickhouse_databases", + "Lists all available databases in a ClickHouse datasource. Returns a list of database names that can be used for further table exploration.", + listClickHouseDatabases, + mcp.WithTitleAnnotation("List ClickHouse databases"), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithReadOnlyHintAnnotation(true), +) + +// ListClickHouseTablesParams defines the parameters for listing ClickHouse tables +type ListClickHouseTablesParams struct { + DatasourceUID string `json:"datasourceUid" jsonschema:"required,description=The UID of the datasource to query"` + Database string `json:"database" jsonschema:"required,description=The name of the database to list tables from"` + Like string `json:"like,omitempty" jsonschema:"description=Optionally\\, filter table names using LIKE pattern (e.g. 'log_%' for tables starting with 'log_')"` + Limit int `json:"limit,omitempty" jsonschema:"description=Optionally\\, the maximum number of tables to return (default: 100)"` +} + +// listClickHouseTables lists all tables in a ClickHouse database with metadata +func listClickHouseTables(ctx context.Context, args ListClickHouseTablesParams) ([]ClickHouseTable, error) { + client, err := newClickHouseClient(ctx, args.DatasourceUID) + if err != nil { + return nil, fmt.Errorf("creating ClickHouse client: %w", err) + } + + // Build the query to get table information from system.tables + query := fmt.Sprintf(`SELECT + database, + name, + engine, + total_rows, + total_bytes, + total_bytes_uncompressed, + parts, + active_parts, + comment + FROM system.tables + WHERE database = '%s'`, strings.ReplaceAll(args.Database, "'", "''")) + + if args.Like != "" { + query += fmt.Sprintf(" AND name LIKE '%s'", strings.ReplaceAll(args.Like, "'", "''")) + } + + query += " ORDER BY name" + + if args.Limit > 0 { + query += fmt.Sprintf(" LIMIT %d", args.Limit) + } else { + query += " LIMIT 100" + } + + result, err := client.queryJSON(ctx, query, 0) + if err != nil { + return nil, err + } + + var tables []ClickHouseTable + for _, row := range result.Rows { + if len(row) >= 9 { + table := ClickHouseTable{ + Database: toString(row[0]), + Name: toString(row[1]), + Engine: toString(row[2]), + TotalRows: toUint64(row[3]), + TotalBytes: toUint64(row[4]), + TotalBytesUncompressed: toUint64(row[5]), + Parts: toUint64(row[6]), + ActiveParts: toUint64(row[7]), + Comment: toString(row[8]), + } + tables = append(tables, table) + } + } + + return tables, nil +} + +// Helper functions for type conversion +func toString(val interface{}) string { + if val == nil { + return "" + } + if str, ok := val.(string); ok { + return str + } + return fmt.Sprintf("%v", val) +} + +func toUint64(val interface{}) uint64 { + if val == nil { + return 0 + } + + switch v := val.(type) { + case float64: + return uint64(v) + case int: + return uint64(v) + case int64: + return uint64(v) + case uint64: + return v + case string: + if num, err := strconv.ParseUint(v, 10, 64); err == nil { + return num + } + } + return 0 +} + +// ListClickHouseTables is a tool for listing ClickHouse tables +var ListClickHouseTables = mcpgrafana.MustTool( + "list_clickhouse_tables", + "Lists all tables in a specified ClickHouse database with detailed metadata including table engine, row count, size information, and comments. Supports filtering with LIKE patterns and result limiting.", + listClickHouseTables, + mcp.WithTitleAnnotation("List ClickHouse tables"), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithReadOnlyHintAnnotation(true), +) + +// DescribeClickHouseTableParams defines the parameters for describing a ClickHouse table +type DescribeClickHouseTableParams struct { + DatasourceUID string `json:"datasourceUid" jsonschema:"required,description=The UID of the datasource to query"` + Database string `json:"database" jsonschema:"required,description=The name of the database"` + Table string `json:"table" jsonschema:"required,description=The name of the table to describe"` +} + +// describeClickHouseTable describes the structure of a ClickHouse table +func describeClickHouseTable(ctx context.Context, args DescribeClickHouseTableParams) ([]ClickHouseColumn, error) { + client, err := newClickHouseClient(ctx, args.DatasourceUID) + if err != nil { + return nil, fmt.Errorf("creating ClickHouse client: %w", err) + } + + // Get column information from system.columns + query := fmt.Sprintf(`SELECT + database, + table, + name, + type, + default_kind, + default_expression, + comment + FROM system.columns + WHERE database = '%s' AND table = '%s' + ORDER BY position`, + strings.ReplaceAll(args.Database, "'", "''"), + strings.ReplaceAll(args.Table, "'", "''")) + + result, err := client.queryJSON(ctx, query, 0) + if err != nil { + return nil, err + } + + var columns []ClickHouseColumn + for _, row := range result.Rows { + if len(row) >= 7 { + column := ClickHouseColumn{ + Database: toString(row[0]), + Table: toString(row[1]), + Name: toString(row[2]), + Type: toString(row[3]), + DefaultKind: toString(row[4]), + DefaultExpression: toString(row[5]), + Comment: toString(row[6]), + } + columns = append(columns, column) + } + } + + return columns, nil +} + +// DescribeClickHouseTable is a tool for describing ClickHouse table structure +var DescribeClickHouseTable = mcpgrafana.MustTool( + "describe_clickhouse_table", + "Describes the structure of a ClickHouse table, returning detailed information about each column including name, data type, default values, and comments. Useful for understanding table schema before writing queries.", + describeClickHouseTable, + mcp.WithTitleAnnotation("Describe ClickHouse table"), + mcp.WithIdempotentHintAnnotation(true), + mcp.WithReadOnlyHintAnnotation(true), +) + +// AddClickHouseTools registers all ClickHouse tools with the MCP server +func AddClickHouseTools(mcp *server.MCPServer) { + QueryClickHouse.Register(mcp) + ListClickHouseDatabases.Register(mcp) + ListClickHouseTables.Register(mcp) + DescribeClickHouseTable.Register(mcp) +} diff --git a/tools/clickhouse_test.go b/tools/clickhouse_test.go new file mode 100644 index 00000000..cb8f354d --- /dev/null +++ b/tools/clickhouse_test.go @@ -0,0 +1,225 @@ +// Requires a Grafana instance running on localhost:3000, +// with a ClickHouse datasource provisioned. +// Run with `go test -tags integration`. +//go:build integration + +package tools + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestClickHouseTools(t *testing.T) { + t.Run("list clickhouse databases", func(t *testing.T) { + ctx := newTestContext() + result, err := listClickHouseDatabases(ctx, ListClickHouseDatabasesParams{ + DatasourceUID: "clickhouse", + }) + require.NoError(t, err) + assert.Greater(t, len(result), 0, "Expected at least one database") + + // ClickHouse should at least have the 'system' database + assert.Contains(t, result, "system", "Expected 'system' database to be present") + }) + + t.Run("list clickhouse tables", func(t *testing.T) { + ctx := newTestContext() + result, err := listClickHouseTables(ctx, ListClickHouseTablesParams{ + DatasourceUID: "clickhouse", + Database: "system", + Limit: 10, + }) + require.NoError(t, err) + assert.Greater(t, len(result), 0, "Expected at least one table in system database") + + // Check that we got proper table metadata + for _, table := range result { + assert.Equal(t, "system", table.Database) + assert.NotEmpty(t, table.Name) + assert.NotEmpty(t, table.Engine) + } + }) + + t.Run("list clickhouse tables with like filter", func(t *testing.T) { + ctx := newTestContext() + result, err := listClickHouseTables(ctx, ListClickHouseTablesParams{ + DatasourceUID: "clickhouse", + Database: "system", + Like: "columns", + Limit: 5, + }) + require.NoError(t, err) + + // Should find the 'columns' table + found := false + for _, table := range result { + if table.Name == "columns" { + found = true + break + } + } + assert.True(t, found, "Expected to find 'columns' table with LIKE filter") + }) + + t.Run("describe clickhouse table", func(t *testing.T) { + ctx := newTestContext() + result, err := describeClickHouseTable(ctx, DescribeClickHouseTableParams{ + DatasourceUID: "clickhouse", + Database: "system", + Table: "databases", + }) + require.NoError(t, err) + assert.Greater(t, len(result), 0, "Expected at least one column in system.databases table") + + // Check that we got proper column metadata + nameFound := false + for _, col := range result { + assert.Equal(t, "system", col.Database) + assert.Equal(t, "databases", col.Table) + assert.NotEmpty(t, col.Name) + assert.NotEmpty(t, col.Type) + if col.Name == "name" { + nameFound = true + } + } + assert.True(t, nameFound, "Expected 'name' column in system.databases table") + }) + + t.Run("query clickhouse basic select", func(t *testing.T) { + ctx := newTestContext() + result, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "SELECT name FROM system.databases LIMIT 5", + Limit: 5, + }) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Greater(t, len(result.Columns), 0, "Expected at least one column") + assert.Equal(t, "name", result.Columns[0]) + assert.GreaterOrEqual(t, len(result.Rows), 1, "Expected at least one row") + }) + + t.Run("query clickhouse show databases", func(t *testing.T) { + ctx := newTestContext() + result, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "SHOW DATABASES", + }) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Greater(t, len(result.Columns), 0, "Expected at least one column") + assert.GreaterOrEqual(t, len(result.Rows), 1, "Expected at least one database") + }) + + t.Run("query clickhouse describe table", func(t *testing.T) { + ctx := newTestContext() + result, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "DESCRIBE system.tables", + }) + require.NoError(t, err) + assert.NotNil(t, result) + assert.Greater(t, len(result.Columns), 0, "Expected at least one column from DESCRIBE") + assert.GreaterOrEqual(t, len(result.Rows), 1, "Expected at least one row from DESCRIBE") + }) + + t.Run("query clickhouse with limit", func(t *testing.T) { + ctx := newTestContext() + result, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "SELECT name FROM system.tables", + Limit: 3, + }) + require.NoError(t, err) + assert.NotNil(t, result) + assert.LessOrEqual(t, len(result.Rows), 3, "Expected no more than 3 rows due to limit") + }) + + t.Run("query clickhouse reject unsafe queries", func(t *testing.T) { + ctx := newTestContext() + + unsafeQueries := []string{ + "INSERT INTO test VALUES (1)", + "UPDATE test SET x = 1", + "DELETE FROM test", + "CREATE TABLE test (x Int32)", + "DROP TABLE test", + "TRUNCATE TABLE test", + } + + for _, query := range unsafeQueries { + t.Run("reject "+query, func(t *testing.T) { + _, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: query, + }) + assert.Error(t, err, "Expected error for unsafe query: %s", query) + assert.Contains(t, err.Error(), "only SELECT, SHOW, and DESCRIBE queries are allowed") + }) + } + }) + + t.Run("query clickhouse with invalid datasource", func(t *testing.T) { + ctx := newTestContext() + _, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "nonexistent", + Query: "SELECT 1", + }) + assert.Error(t, err) + assert.Contains(t, err.Error(), "datasource") + }) + + t.Run("query clickhouse enforce limits", func(t *testing.T) { + ctx := newTestContext() + + // Test default limit + result, err := queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "SELECT number FROM system.numbers", + Limit: 0, // Should use default + }) + require.NoError(t, err) + // Default limit is enforced by max_result_rows parameter + assert.NotNil(t, result) + + // Test max limit enforcement + result, err = queryClickHouse(ctx, QueryClickHouseParams{ + DatasourceUID: "clickhouse", + Query: "SELECT number FROM system.numbers", + Limit: 2000, // Should be capped to MaxClickHouseLimit (1000) + }) + require.NoError(t, err) + assert.NotNil(t, result) + }) +} + +// Test helper functions + +func TestClickHouseHelperFunctions(t *testing.T) { + t.Run("toString helper", func(t *testing.T) { + assert.Equal(t, "hello", toString("hello")) + assert.Equal(t, "123", toString(123)) + assert.Equal(t, "123.45", toString(123.45)) + assert.Equal(t, "", toString(nil)) + }) + + t.Run("toUint64 helper", func(t *testing.T) { + assert.Equal(t, uint64(123), toUint64(123)) + assert.Equal(t, uint64(123), toUint64(int64(123))) + assert.Equal(t, uint64(123), toUint64(float64(123.0))) + assert.Equal(t, uint64(123), toUint64(uint64(123))) + assert.Equal(t, uint64(123), toUint64("123")) + assert.Equal(t, uint64(0), toUint64("invalid")) + assert.Equal(t, uint64(0), toUint64(nil)) + }) + + t.Run("enforceClickHouseLimit", func(t *testing.T) { + assert.Equal(t, DefaultClickHouseLimit, enforceClickHouseLimit(0)) + assert.Equal(t, DefaultClickHouseLimit, enforceClickHouseLimit(-1)) + assert.Equal(t, 50, enforceClickHouseLimit(50)) + assert.Equal(t, MaxClickHouseLimit, enforceClickHouseLimit(2000)) + }) +} diff --git a/tools/clickhouse_unit_test.go b/tools/clickhouse_unit_test.go new file mode 100644 index 00000000..eb95bd0e --- /dev/null +++ b/tools/clickhouse_unit_test.go @@ -0,0 +1,77 @@ +package tools + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestClickHouseHelperFunctions(t *testing.T) { + t.Run("toString", func(t *testing.T) { + tests := []struct { + input interface{} + expected string + }{ + {"hello", "hello"}, + {123, "123"}, + {int64(456), "456"}, + {float64(123.45), "123.45"}, + {true, "true"}, + {nil, ""}, + } + + for _, test := range tests { + assert.Equal(t, test.expected, toString(test.input), "toString(%v)", test.input) + } + }) + + t.Run("toUint64", func(t *testing.T) { + tests := []struct { + input interface{} + expected uint64 + }{ + {123, uint64(123)}, + {int64(456), uint64(456)}, + {float64(789.0), uint64(789)}, + {float64(789.9), uint64(789)}, // Should truncate + {uint64(999), uint64(999)}, + {"123", uint64(123)}, + {"456.78", uint64(0)}, // Invalid string should return 0 + {"invalid", uint64(0)}, + {nil, uint64(0)}, + {true, uint64(0)}, // Unsupported type should return 0 + } + + for _, test := range tests { + assert.Equal(t, test.expected, toUint64(test.input), "toUint64(%v)", test.input) + } + }) + + t.Run("enforceClickHouseLimit", func(t *testing.T) { + tests := []struct { + input int + expected int + }{ + {0, DefaultClickHouseLimit}, + {-1, DefaultClickHouseLimit}, + {50, 50}, + {DefaultClickHouseLimit, DefaultClickHouseLimit}, + {MaxClickHouseLimit, MaxClickHouseLimit}, + {MaxClickHouseLimit + 100, MaxClickHouseLimit}, // Should be capped + {2000, MaxClickHouseLimit}, // Should be capped + } + + for _, test := range tests { + assert.Equal(t, test.expected, enforceClickHouseLimit(test.input), "enforceClickHouseLimit(%d)", test.input) + } + }) +} + +func TestClickHouseConstants(t *testing.T) { + t.Run("constants are reasonable", func(t *testing.T) { + assert.Greater(t, DefaultClickHouseLimit, 0, "Default limit should be positive") + assert.Greater(t, MaxClickHouseLimit, DefaultClickHouseLimit, "Max limit should be greater than default") + assert.Equal(t, 100, DefaultClickHouseLimit, "Default limit should be 100") + assert.Equal(t, 1000, MaxClickHouseLimit, "Max limit should be 1000") + }) +}