67 lines
1.8 KiB
Go
67 lines
1.8 KiB
Go
package investigation
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
|
|
)
|
|
|
|
// ClickHouseStore keeps the latest finished record per (ip, mode) as JSON.
|
|
// No TTL: what an address turned out to be is worth keeping.
|
|
type ClickHouseStore struct {
|
|
conn driver.Conn
|
|
db string
|
|
}
|
|
|
|
func NewClickHouseStore(ctx context.Context, conn driver.Conn, db string) (*ClickHouseStore, error) {
|
|
err := conn.Exec(ctx, fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s.ip_investigations (
|
|
ip String,
|
|
mode LowCardinality(String),
|
|
status LowCardinality(String),
|
|
started_at DateTime64(3, 'UTC'),
|
|
finished_at DateTime64(3, 'UTC'),
|
|
record String
|
|
) ENGINE = ReplacingMergeTree(finished_at)
|
|
ORDER BY (ip, mode)`, db))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &ClickHouseStore{conn: conn, db: db}, nil
|
|
}
|
|
|
|
func (s *ClickHouseStore) Save(ctx context.Context, rec Record) error {
|
|
body, err := json.Marshal(rec)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
finished := time.Now().UTC()
|
|
if rec.FinishedAt != nil {
|
|
finished = *rec.FinishedAt
|
|
}
|
|
return s.conn.Exec(ctx, fmt.Sprintf(`INSERT INTO %s.ip_investigations (ip, mode, status, started_at, finished_at, record) VALUES (?, ?, ?, ?, ?, ?)`, s.db),
|
|
rec.IP, rec.Mode, rec.Status, rec.StartedAt.UTC(), finished, string(body))
|
|
}
|
|
|
|
func (s *ClickHouseStore) Load(ctx context.Context, ip string) ([]Record, error) {
|
|
rows, err := s.conn.Query(ctx, fmt.Sprintf(`SELECT record FROM %s.ip_investigations FINAL WHERE ip = ?`, s.db), ip)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []Record
|
|
for rows.Next() {
|
|
var body string
|
|
if err := rows.Scan(&body); err != nil {
|
|
return nil, err
|
|
}
|
|
var rec Record
|
|
if err := json.Unmarshal([]byte(body), &rec); err == nil {
|
|
out = append(out, rec)
|
|
}
|
|
}
|
|
return out, rows.Err()
|
|
}
|