allow overriding pinot cluster query timeout (#7398)

Author: pjain1

runtime/drivers/pinot/olap.go

@@ -45,7 +45,7 @@ func (c *connection) Exec(ctx context.Context, stmt *drivers.Statement) error {
 
 func (c *connection) Query(ctx context.Context, stmt *drivers.Statement) (*drivers.Result, error) {
 	if c.logQueries {
-		c.logger.Info("pinot query", zap.String("sql", stmt.Query), zap.Any("args", stmt.Args), observability.ZapCtx(ctx))
+		c.logger.Info("pinot query", zap.String("sql", stmt.Query), zap.Any("args", stmt.Args), zap.Int64("timeoutMS", c.timeoutMS), observability.ZapCtx(ctx))
 	}
 	if stmt.DryRun {
 		rows, err := c.db.QueryxContext(ctx, "EXPLAIN PLAN FOR "+stmt.Query, stmt.Args...)
@@ -61,6 +61,11 @@ func (c *connection) Query(ctx context.Context, stmt *drivers.Statement) (*drive
 		ctx, cancelFunc = context.WithTimeout(ctx, stmt.ExecutionTimeout)
 	}
 
+	// add timeout if configured to the sql to propagate it to the Pinot server to override the cluster timeout
+	if c.timeoutMS > 0 {
+		stmt.Query = fmt.Sprintf("SET timeoutMS=%d; %s", c.timeoutMS, stmt.Query)
+	}
+
 	rows, err := c.db.QueryxContext(ctx, stmt.Query, stmt.Args...)
 	if err != nil {
 		if cancelFunc != nil {

runtime/drivers/pinot/pinot.go

@@ -114,6 +114,8 @@ type configProperties struct {
 	LogQueries bool `mapstructure:"log_queries"`
 	// MaxOpenConns is the maximum number of open connections to the database. Set to 0 to use the default value or -1 for unlimited.
 	MaxOpenConns int `mapstructure:"max_open_conns"`
+	// TimeoutMS is the timeout in milliseconds for queries. Set to 0 to use the cluster default.
+	TimeoutMS int64 `mapstructure:"timeout_ms"`
 }
 
 // Open a connection to Apache Pinot using HTTP API.
@@ -205,6 +207,7 @@ func (d driver) Open(instanceID string, config map[string]any, st *storage.Clien
 		schemaURL:  controller,
 		headers:    headers,
 		logQueries: conf.LogQueries,
+		timeoutMS:  conf.TimeoutMS,
 		logger:     logger,
 	}
 	return conn, nil
@@ -229,6 +232,7 @@ type connection struct {
 	schemaURL  string
 	headers    map[string]string
 	logQueries bool
+	timeoutMS  int64 // timeout in milliseconds for queries, 0 means use cluster default
 	logger     *zap.Logger
 }