100 lines
2.5 KiB
Go
100 lines
2.5 KiB
Go
package clickhouse
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"example.com/siem-greenfield/internal/config"
|
|
"example.com/siem-greenfield/internal/contracts"
|
|
)
|
|
|
|
type Client struct {
|
|
base, db, user, pass string
|
|
hc *http.Client
|
|
}
|
|
|
|
func New(cfg config.Config) *Client {
|
|
return &Client{base: cfg.ClickHouseURL, db: cfg.ClickHouseDB, user: cfg.ClickHouseUser, pass: cfg.ClickHousePassword, hc: &http.Client{Timeout: 30 * time.Second}}
|
|
}
|
|
func (c *Client) request(ctx context.Context, query string, body io.Reader) (*http.Response, error) {
|
|
u := c.base + "/?query=" + url.QueryEscape(query)
|
|
req, e := http.NewRequestWithContext(ctx, http.MethodPost, u, body)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
req.SetBasicAuth(c.user, c.pass)
|
|
return c.hc.Do(req)
|
|
}
|
|
func (c *Client) Exec(ctx context.Context, q string) error {
|
|
r, e := c.request(ctx, q, nil)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
defer r.Body.Close()
|
|
if r.StatusCode/100 != 2 {
|
|
b, _ := io.ReadAll(io.LimitReader(r.Body, 8192))
|
|
return fmt.Errorf("clickhouse %s: %s", r.Status, string(b))
|
|
}
|
|
return nil
|
|
}
|
|
func (c *Client) InsertEvents(ctx context.Context, events []contracts.CanonicalEvent) error {
|
|
if len(events) == 0 {
|
|
return nil
|
|
}
|
|
var b bytes.Buffer
|
|
enc := json.NewEncoder(&b)
|
|
enc.SetEscapeHTML(false)
|
|
for i := range events {
|
|
if e := enc.Encode(events[i]); e != nil {
|
|
return e
|
|
}
|
|
}
|
|
q := `INSERT INTO ` + Ident(c.db) + `.events FORMAT JSONEachRow`
|
|
r, e := c.request(ctx, q, &b)
|
|
if e != nil {
|
|
return e
|
|
}
|
|
defer r.Body.Close()
|
|
if r.StatusCode/100 != 2 {
|
|
x, _ := io.ReadAll(io.LimitReader(r.Body, 16384))
|
|
return fmt.Errorf("clickhouse insert %s: %s", r.Status, string(x))
|
|
}
|
|
return nil
|
|
}
|
|
func (c *Client) QueryJSON(ctx context.Context, q string) ([]map[string]any, error) {
|
|
if !strings.Contains(strings.ToUpper(q), "FORMAT") {
|
|
q += " FORMAT JSONEachRow"
|
|
}
|
|
r, e := c.request(ctx, q, nil)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
defer r.Body.Close()
|
|
if r.StatusCode/100 != 2 {
|
|
x, _ := io.ReadAll(io.LimitReader(r.Body, 16384))
|
|
return nil, fmt.Errorf("clickhouse query %s: %s", r.Status, string(x))
|
|
}
|
|
var out []map[string]any
|
|
s := bufio.NewScanner(r.Body)
|
|
buf := make([]byte, 0, 64*1024)
|
|
s.Buffer(buf, 4*1024*1024)
|
|
for s.Scan() {
|
|
var m map[string]any
|
|
if e := json.Unmarshal(s.Bytes(), &m); e != nil {
|
|
return nil, e
|
|
}
|
|
out = append(out, m)
|
|
}
|
|
return out, s.Err()
|
|
}
|
|
func Q(s string) string { return "'" + strings.ReplaceAll(s, "'", "''") + "'" }
|
|
func Ident(s string) string { return "`" + strings.ReplaceAll(s, "`", "``") + "`" }
|