feat: support per query settings (partially)

Author: Yury Yurochko

bootstrap.go

@@ -149,13 +149,19 @@ func open(dsn string) (*clickhouse, error) {
 		connOpenStrategy = connOpenInOrder
 	}
 
+	settings, err := makeQuerySettings(query)
+	if err != nil {
+		return nil, err
+	}
+
 	if v, err := strconv.ParseBool(query.Get("compress")); err == nil {
 		compress = v
 	}
 
 	var (
 		ch = clickhouse{
 			logf:      func(string, ...interface{}) {},
+			settings:  settings,
 			compress:  compress,
 			blockSize: blockSize,
 			ServerInfo: data.ServerInfo{

clickhouse.go

@@ -47,6 +47,7 @@ type clickhouse struct {
 	buffer        *bufio.Writer
 	decoder       *binary.Decoder
 	encoder       *binary.Encoder
+	settings      *querySettings
 	compress      bool
 	blockSize     int
 	inTransaction bool

clickhouse_send_query.go

@@ -29,7 +29,15 @@ func (ch *clickhouse) sendQuery(query string) error {
 		ch.encoder.String("")
 	}
 
-	if err := ch.encoder.String(""); err != nil { // settings
+	// the settings are written as list of contiguous name-value pairs, finished with empty name
+	if !ch.settings.IsEmpty() {
+		ch.logf("[query settings] %s", ch.settings.settingsStr)
+		if err := ch.settings.Serialize(ch.encoder); err != nil {
+			return err
+		}
+	}
+	// empty string is a marker of the end of the settings
+	if err := ch.encoder.String(""); err != nil {
 		return err
 	}
 	if err := ch.encoder.Uvarint(protocol.StateComplete); err != nil {

issues_test.go

@@ -234,3 +234,18 @@ func TestNullableEnumWithoutLeadZero(t *testing.T) {
 		}
 	}
 }
+
+func TestQuerySettings(t *testing.T) {
+	settings := "max_memory_usage=1000&max_execution_time=10&max_execution_speed=10"
+	connect, err := sql.Open(
+		"clickhouse",
+		"tcp://127.0.0.1:9000?debug=true"+"&"+settings,
+	)
+	require.Nil(t, err)
+	require.Nil(t, connect.Ping())
+	defer connect.Close()
+
+	_, err = connect.Query(`SELECT * FROM system.parts`)
+	require.NotNil(t, err)
+	require.Contains(t, err.Error(), "Memory limit")
+}

query_settings.go

@@ -0,0 +1,88 @@
+package clickhouse
+
+import (
+	"fmt"
+	"net/url"
+	"strconv"
+
+	"github.com/kshvakov/clickhouse/lib/binary"
+)
+
+type querySettingType int
+
+// all possible query setting's data type
+// TODO: support remaining data types
+const (
+	uintQS querySettingType = iota + 1
+)
+
+// description of single query setting
+type querySettingInfo struct {
+	name   string
+	qsType querySettingType
+}
+
+// all possible query settings
+// TODO: support remaining query serrings
+var querySettingList = []querySettingInfo{
+	{"max_memory_usage", uintQS},
+	{"max_execution_time", uintQS},
+	{"max_execution_speed", uintQS},
+}
+
+type querySettingValueEncoder func(enc *binary.Encoder) error
+
+type querySettings struct {
+	settings    map[string]querySettingValueEncoder
+	settingsStr string // used for debug output
+}
+
+func makeQuerySettings(query url.Values) (*querySettings, error) {
+	qs := &querySettings{
+		settings:    make(map[string]querySettingValueEncoder),
+		settingsStr: "",
+	}
+
+	for _, info := range querySettingList {
+		valueStr := query.Get(info.name)
+		if valueStr == "" {
+			continue
+		}
+
+		switch info.qsType {
+		case uintQS:
+			value, err := strconv.ParseUint(valueStr, 10, 64)
+			if err != nil {
+				return nil, err
+			}
+			qs.settings[info.name] = func(enc *binary.Encoder) error { return enc.Uvarint(value) }
+		default:
+			err := fmt.Errorf("query setting %s has unsupported data type", info.name)
+			return nil, err
+		}
+
+		if qs.settingsStr != "" {
+			qs.settingsStr += "&"
+		}
+		qs.settingsStr += info.name + "=" + valueStr
+	}
+
+	return qs, nil
+}
+
+func (qs *querySettings) IsEmpty() bool {
+	return len(qs.settings) == 0
+}
+
+func (qs *querySettings) Serialize(enc *binary.Encoder) error {
+	for name, fn := range qs.settings {
+		if err := enc.String(name); err != nil {
+			return err
+		}
+		if err := fn(enc); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}