From 3957d91a9bb356d97e83eb97a3640e80780c3963 Mon Sep 17 00:00:00 2001
From: Nityananda Gohain <nityanandagohain@gmail.com>
Date: Mon, 6 May 2024 15:01:53 +0530
Subject: [PATCH] fix: add read-in-order config (#4918)

---
 pkg/query-service/app/clickhouseReader/reader.go  | 12 ++++++++++++
 pkg/query-service/app/clickhouseReader/wrapper.go |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/pkg/query-service/app/clickhouseReader/reader.go b/pkg/query-service/app/clickhouseReader/reader.go
index ff5d754843..efad720f74 100644
--- a/pkg/query-service/app/clickhouseReader/reader.go
+++ b/pkg/query-service/app/clickhouseReader/reader.go
@@ -163,12 +163,24 @@ func NewReaderFromClickhouseConnection(
 		os.Exit(1)
 	}
 
+	regex := os.Getenv("ClickHouseOptimizeReadInOrderRegex")
+	var regexCompiled *regexp.Regexp
+	if regex != "" {
+		regexCompiled, err = regexp.Compile(regex)
+		if err != nil {
+			zap.L().Error("Incorrect regex for ClickHouseOptimizeReadInOrderRegex")
+			os.Exit(1)
+		}
+	}
+
 	wrap := clickhouseConnWrapper{
 		conn: db,
 		settings: ClickhouseQuerySettings{
 			MaxExecutionTimeLeaf:                os.Getenv("ClickHouseMaxExecutionTimeLeaf"),
 			TimeoutBeforeCheckingExecutionSpeed: os.Getenv("ClickHouseTimeoutBeforeCheckingExecutionSpeed"),
 			MaxBytesToRead:                      os.Getenv("ClickHouseMaxBytesToRead"),
+			OptimizeReadInOrderRegex:            os.Getenv("ClickHouseOptimizeReadInOrderRegex"),
+			OptimizeReadInOrderRegexCompiled:    regexCompiled,
 		},
 	}
 
diff --git a/pkg/query-service/app/clickhouseReader/wrapper.go b/pkg/query-service/app/clickhouseReader/wrapper.go
index eac6721b4a..2f08167534 100644
--- a/pkg/query-service/app/clickhouseReader/wrapper.go
+++ b/pkg/query-service/app/clickhouseReader/wrapper.go
@@ -3,6 +3,7 @@ package clickhouseReader
 import (
 	"context"
 	"encoding/json"
+	"regexp"
 	"strings"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
@@ -13,6 +14,8 @@ type ClickhouseQuerySettings struct {
 	MaxExecutionTimeLeaf                string
 	TimeoutBeforeCheckingExecutionSpeed string
 	MaxBytesToRead                      string
+	OptimizeReadInOrderRegex            string
+	OptimizeReadInOrderRegexCompiled    *regexp.Regexp
 }
 
 type clickhouseConnWrapper struct {
@@ -58,6 +61,11 @@ func (c clickhouseConnWrapper) addClickHouseSettings(ctx context.Context, query
 		settings["timeout_before_checking_execution_speed"] = c.settings.TimeoutBeforeCheckingExecutionSpeed
 	}
 
+	// only list queries of
+	if c.settings.OptimizeReadInOrderRegex != "" && c.settings.OptimizeReadInOrderRegexCompiled.Match([]byte(query)) {
+		settings["optimize_read_in_order"] = 0
+	}
+
 	ctx = clickhouse.Context(ctx, clickhouse.WithSettings(settings))
 	return ctx
 }