cloud-services/pkg/clickhouse/clickhouse.go

package clickhouse

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/fiskerinc/cloud-services/pkg/common"
	"github.com/fiskerinc/cloud-services/pkg/utils/envtool"

	"github.com/ClickHouse/clickhouse-go/v2"
	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
	"github.com/pkg/errors"
)

var (
	TIMEOUT               = envtool.GetEnvInt("CLICKHOUSE_TIMEOUT", 1)
	MAX_CONNS             = envtool.GetEnvInt("CLICKHOUSE_MAX_CONNS", 1)
	CLICKHOUSE_HOST       = envtool.GetEnv("CLICKHOUSE_HOST", "localhost")
	CLICKHOUSE_PORT       = envtool.GetEnv("CLICKHOUSE_PORT", "9000")
	CLICKHOUSE_DB         = envtool.GetEnv("CLICKHOUSE_DB", "default")
	CLICKHOUSE_USER       = envtool.GetEnv("CLICKHOUSE_USER", "")
	CLICKHOUSE_PASS       = envtool.GetEnv("CLICKHOUSE_PASS", "")
	VEHICLE_FILTERS_TABLE = envtool.GetEnv("CLICKHOUSE_VEHICLE_FILTERS_TABLE", "can_filter_list_vin")
	DEFAULT_FILTERS_TABLE = envtool.GetEnv("CLICKHOUSE_DEFAULT_FILTERS_TABLE", "can_filter_list_all")
)

func NewClient(conn ConnInterface) (ClientInterface, error) {
	return &Client{conn: conn}, nil
}

func NewConn() (clickhouse.Conn, error) {
	return clickhouse.Open(&clickhouse.Options{
		Addr: []string{fmt.Sprintf("%s:%s", CLICKHOUSE_HOST, CLICKHOUSE_PORT)},
		Auth: clickhouse.Auth{
			Database: CLICKHOUSE_DB,
			Username: CLICKHOUSE_USER,
			Password: CLICKHOUSE_PASS,
		},
		DialTimeout:     time.Second * 60,
		MaxOpenConns:    MAX_CONNS,
		MaxIdleConns:    MAX_CONNS,
		ConnMaxLifetime: 24 * time.Hour,
		Compression: &clickhouse.Compression{
			Method: clickhouse.CompressionLZ4,
		},
	})
}

type ClientInterface interface {
	Select(result interface{}, query string) error
	RetrieveDefaultFilters() ([]CANFilterSchema, error)
	RetrieveFiltersForVehicle(vin string) ([]CANFilterSchema, error)
	SaveDBCInfo(dbc common.DBCDesc) error
	SaveDBCMessages(ms []common.MessageDesc) error
	SaveDBCSignals(signals []common.SignalDesc) error
	SelectDBCsByVersions(versions []string) ([]string, error)
	SelectDBCSignals(dbc string, options PageQueryOptions) ([]common.SignalDescWithECU, int, error)
	TruncateDBCDescs() error

	SetConn(conn ConnInterface)
	Exec(query string) error
}

type ConnInterface interface {
	Select(ctx context.Context, dest interface{}, query string, args ...interface{}) error
	PrepareBatch(ctx context.Context, query string) (driver.Batch, error)
	AsyncInsert(ctx context.Context, query string, wait bool) error
	QueryRow(ctx context.Context, query string, args ...interface{}) driver.Row
	Query(ctx context.Context, query string, args ...interface{}) (driver.Rows, error)
	Exec(ctx context.Context, query string, args ...interface{}) error
}

type Client struct {
	conn ConnInterface
}

func (c *Client) Select(result interface{}, query string) error {
	ctx := context.Background()

	err := c.conn.Select(ctx, result, query)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (c *Client) RetrieveDefaultFilters() ([]CANFilterSchema, error) {
	var result []CANFilterSchema

	if err := c.Select(&result, fmt.Sprintf("SELECT ID, Period FROM %s", DEFAULT_FILTERS_TABLE)); err != nil {
		return result, err
	}

	return result, nil
}

func (c *Client) RetrieveFiltersForVehicle(vin string) ([]CANFilterSchema, error) {
	var result []CANFilterSchema

	if err := c.Select(&result, fmt.Sprintf("SELECT ID, Period FROM %s WHERE VIN='%s'", VEHICLE_FILTERS_TABLE, vin)); err != nil {
		return result, err
	}

	return result, nil
}

func (c *Client) SelectDBCsByVersions(versions []string) ([]string, error) {
	if len(versions) == 0 {
		return nil, nil
	}

	q := fmt.Sprintf("'%s'", strings.Join(versions, "','"))

	var result []string
	if err := c.Select(&result, fmt.Sprintf("SELECT dbc_name FROM dbcs where dbc_hash IN (%s)", q)); err != nil {
		return result, err
	}

	return result, nil
}

func (c *Client) TruncateDBCDescs() error {
	err := c.Exec("TRUNCATE TABLE dbc_signals_shard ON CLUSTER default")
	if err != nil {
		return err
	}

	err = c.Exec("TRUNCATE TABLE dbc_messages_shard ON CLUSTER default")
	if err != nil {
		return err
	}

	err = c.Exec("TRUNCATE TABLE dbcs_shard ON CLUSTER default")
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) SaveDBCInfo(dbc common.DBCDesc) error {
	query := fmt.Sprintf(`INSERT INTO dbcs (dbc_hash, dbc_name) VALUES ('%s', '%s')`, dbc.Hash, dbc.Name)
	return errors.WithStack(c.conn.AsyncInsert(context.Background(), query, true))
}

func (c *Client) SaveDBCMessages(ms []common.MessageDesc) error {
	batch, err := c.conn.PrepareBatch(context.Background(),
		`INSERT INTO dbc_messages (dbc_hash, message_name, message_id, is_extended,
					send_type, length, description, sender_node, cycle_time_ns, delay_time_ns, ecu_name)`)
	if err != nil {
		return errors.WithStack(err)
	}

	for _, m := range ms {
		err = batch.Append(m.DBCHash, m.Name, m.ID, m.IsExtended, m.SendType, m.Length, m.Description,
			m.SenderNode, m.CycleTime, m.DelayTime, m.ECUName)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return errors.WithStack(batch.Send())
}

func (c *Client) SaveDBCSignals(signals []common.SignalDesc) error {
	batch, err := c.conn.PrepareBatch(context.Background(),
		`INSERT INTO dbc_signals (
                         dbc_hash, message_id, signal_name, start,
                         length, big_endian, signed, multiplexer, multiplexed,
                         multiplexer_value, offset, scale, min, max, unit,
                         description, value_descriptions, receiver_nodes, default_value, ecu_name)`)
	if err != nil {
		return errors.WithStack(err)
	}

	for _, s := range signals {
		err = batch.Append(s.DBCHash, s.MessageID, s.Name, s.Start, s.Length, s.IsBigEndian, s.IsSigned,
			s.IsMultiplexer, s.IsMultiplexed, s.MultiplexerValue, s.Offset, s.Scale, s.Min, s.Max,
			s.Unit, s.Description, s.ValueDescriptions, s.ReceiverNodes, s.DefaultValue, s.ECUName)
		if err != nil {
			return errors.WithStack(err)
		}
	}

	return batch.Send()
}

func (c *Client) SelectDBCSignals(dbc string, options PageQueryOptions) ([]common.SignalDescWithECU, int, error) {
	var result []common.SignalDescWithECU
	chCtx := clickhouse.Context(
		context.Background(),
		clickhouse.WithParameters(clickhouse.Parameters{
			"dbc": dbc,

			// we cannot use keywords like "offset" and "limit" as parameter names
			"lim":  fmt.Sprint(options.Limit),
			"offs": fmt.Sprint(options.Offset),
		}))

	query := CreateDBCSignalQuery(options)

	if err := c.conn.Select(chCtx, &result, query); err != nil {
		return nil, 0, errors.WithStack(err)
	}

	var count uint64
	if err := c.conn.QueryRow(chCtx, "SELECT COUNT() FROM dbc_signals a WHERE a.dbc_hash = {dbc:String}").Scan(&count); err != nil {
		return nil, 0, errors.WithStack(err)
	}

	return result, int(count), nil
}

func (c *Client) SetConn(conn ConnInterface) {
	c.conn = conn
}

func (c *Client) Exec(query string) error {
	ctx := context.Background()

	err := c.conn.Exec(ctx, query)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

type CANFilterSchema struct {
	ID     int16 `ch:"ID"`
	Period int32 `ch:"Period"`
}

type MigrationSchema struct {
	Version  int64  `ch:"version"`
	Dirty    uint8  `ch:"dirty"`
	Sequence uint64 `ch:"sequence"`
}

func CreateDBCSignalQuery(options PageQueryOptions) string {

	initQuery := `select * from dbc_signals where dbc_hash = {dbc:String}`

	query := initQuery
	if options.Limit != 0 {
		query += ` LIMIT {lim:UInt64}`
	}

	if options.Offset != 0 {
		query += ` OFFSET {offs:UInt64}`
	}

	return query
}