135 lines
3.5 KiB
Go
135 lines
3.5 KiB
Go
package clickhouse
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/url"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
// Client 通过 HTTP 接口执行只读查询（SELECT），返回 JSONEachRow 解析为 map 或结构体
|
||
type Client struct {
|
||
cfg *Config
|
||
httpCli *http.Client
|
||
}
|
||
|
||
// NewClient 使用共享配置创建客户端
|
||
func NewClient() *Client {
|
||
cfg := SharedConfig()
|
||
return &Client{
|
||
cfg: cfg,
|
||
httpCli: &http.Client{
|
||
Timeout: 30 * time.Second,
|
||
},
|
||
}
|
||
}
|
||
|
||
// IsConfigured 是否已配置
|
||
func (c *Client) IsConfigured() bool {
|
||
return c.cfg != nil && c.cfg.IsConfigured()
|
||
}
|
||
|
||
// Query 执行 SELECT，将每行 JSON 解析到 dest 切片；dest 元素类型需为 *struct 或 map
|
||
func (c *Client) Query(ctx context.Context, query string, dest interface{}) error {
|
||
if !c.IsConfigured() {
|
||
return fmt.Errorf("clickhouse: not configured")
|
||
}
|
||
// 强制 JSONEachRow 便于解析
|
||
q := strings.TrimSpace(query)
|
||
if !strings.HasSuffix(strings.ToUpper(q), "FORMAT JSONEACHROW") {
|
||
query = q + " FORMAT JSONEachRow"
|
||
}
|
||
u := c.buildURL(query)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if c.cfg.User != "" || c.cfg.Password != "" {
|
||
req.SetBasicAuth(c.cfg.User, c.cfg.Password)
|
||
}
|
||
resp, err := c.httpCli.Do(req)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("clickhouse HTTP %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
dec := json.NewDecoder(resp.Body)
|
||
return decodeRows(dec, dest)
|
||
}
|
||
|
||
// QueryRow 执行仅返回一行的查询，将结果解析到 dest（*struct 或 *map）
|
||
func (c *Client) QueryRow(ctx context.Context, query string, dest interface{}) error {
|
||
if !c.IsConfigured() {
|
||
return fmt.Errorf("clickhouse: not configured")
|
||
}
|
||
q := strings.TrimSpace(query)
|
||
if !strings.HasSuffix(strings.ToUpper(q), "FORMAT JSONEACHROW") {
|
||
query = q + " FORMAT JSONEachRow"
|
||
}
|
||
u := c.buildURL(query)
|
||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if c.cfg.User != "" || c.cfg.Password != "" {
|
||
req.SetBasicAuth(c.cfg.User, c.cfg.Password)
|
||
}
|
||
resp, err := c.httpCli.Do(req)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer resp.Body.Close()
|
||
if resp.StatusCode != http.StatusOK {
|
||
body, _ := io.ReadAll(resp.Body)
|
||
return fmt.Errorf("clickhouse HTTP %d: %s", resp.StatusCode, string(body))
|
||
}
|
||
dec := json.NewDecoder(resp.Body)
|
||
return decodeOneRow(dec, dest)
|
||
}
|
||
|
||
func (c *Client) buildURL(query string) string {
|
||
rawURL := fmt.Sprintf("http://%s:%d/?query=%s&database=%s",
|
||
c.cfg.Host, c.cfg.Port, url.QueryEscape(query), url.QueryEscape(c.cfg.Database))
|
||
return rawURL
|
||
}
|
||
|
||
// decodeRows 将 JSONEachRow 流解析到 slice；元素类型须为 *struct 或 *map[string]interface{}
|
||
func decodeRows(dec *json.Decoder, dest interface{}) error {
|
||
// dest 应为 *[]*SomeStruct 或 *[]map[string]interface{}
|
||
switch d := dest.(type) {
|
||
case *[]map[string]interface{}:
|
||
*d = (*d)[:0]
|
||
for {
|
||
var row map[string]interface{}
|
||
if err := dec.Decode(&row); err != nil {
|
||
if err == io.EOF {
|
||
return nil
|
||
}
|
||
return err
|
||
}
|
||
*d = append(*d, row)
|
||
}
|
||
default:
|
||
return fmt.Errorf("clickhouse: unsupported dest type for Query (use *[]map[string]interface{} or implement decoder)")
|
||
}
|
||
}
|
||
|
||
func decodeOneRow(dec *json.Decoder, dest interface{}) error {
|
||
switch d := dest.(type) {
|
||
case *map[string]interface{}:
|
||
if err := dec.Decode(d); err != nil {
|
||
return err
|
||
}
|
||
return nil
|
||
default:
|
||
return fmt.Errorf("clickhouse: unsupported dest type for QueryRow (use *map[string]interface{})")
|
||
}
|
||
}
|