From e44129869ca69f33b7e31aadf59f695b778c0f96 Mon Sep 17 00:00:00 2001
From: goswamisandeep <sandeep.goswami@nutanix.com>
Date: Tue, 6 Dec 2022 19:54:19 +0530
Subject: [PATCH] feat(inputs.sqlserver): introduce timeout for query execution
 (#12319)

---
 plugins/inputs/sqlserver/README.md    |  5 +++++
 plugins/inputs/sqlserver/sample.conf  |  5 +++++
 plugins/inputs/sqlserver/sqlserver.go | 12 +++++++++++-
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/plugins/inputs/sqlserver/README.md b/plugins/inputs/sqlserver/README.md
index d6b0f8ea9..f5bb3500e 100644
--- a/plugins/inputs/sqlserver/README.md
+++ b/plugins/inputs/sqlserver/README.md
@@ -134,6 +134,11 @@ See the [CONFIGURATION.md][CONFIGURATION.md] for more details.
     "Server=192.168.1.10;Port=1433;User Id=<user>;Password=<pw>;app name=telegraf;log=1;",
   ]
 
+  ## Timeout for query execution operation
+  ## Note that the timeout for queries is per query not per gather.
+  ## 0 value means no timeout
+  # query_timeout = "0s"
+
   ## Authentication method
   ## valid methods: "connection_string", "AAD"
   # auth_method = "connection_string"
diff --git a/plugins/inputs/sqlserver/sample.conf b/plugins/inputs/sqlserver/sample.conf
index ad19f28f1..2c1044998 100644
--- a/plugins/inputs/sqlserver/sample.conf
+++ b/plugins/inputs/sqlserver/sample.conf
@@ -11,6 +11,11 @@
     "Server=192.168.1.10;Port=1433;User Id=<user>;Password=<pw>;app name=telegraf;log=1;",
   ]
 
+  ## Timeout for query execution operation
+  ## Note that the timeout for queries is per query not per gather.
+  ## 0 value means no timeout
+  # query_timeout = "0s"
+
   ## Authentication method
   ## valid methods: "connection_string", "AAD"
   # auth_method = "connection_string"
diff --git a/plugins/inputs/sqlserver/sqlserver.go b/plugins/inputs/sqlserver/sqlserver.go
index 96e3c3ac0..282f7ea74 100644
--- a/plugins/inputs/sqlserver/sqlserver.go
+++ b/plugins/inputs/sqlserver/sqlserver.go
@@ -2,6 +2,7 @@
 package sqlserver
 
 import (
+	"context"
 	"database/sql"
 	_ "embed"
 	"errors"
@@ -14,6 +15,7 @@ import (
 	mssql "github.com/denisenkom/go-mssqldb"
 
 	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/config"
 	"github.com/influxdata/telegraf/filter"
 	"github.com/influxdata/telegraf/plugins/inputs"
 )
@@ -24,6 +26,7 @@ var sampleConfig string
 // SQLServer struct
 type SQLServer struct {
 	Servers      []string        `toml:"servers"`
+	QueryTimeout config.Duration `toml:"query_timeout"`
 	AuthMethod   string          `toml:"auth_method"`
 	QueryVersion int             `toml:"query_version" deprecated:"1.16.0;use 'database_type' instead"`
 	AzureDB      bool            `toml:"azuredb" deprecated:"1.16.0;use 'database_type' instead"`
@@ -294,7 +297,14 @@ func (s *SQLServer) Stop() {
 
 func (s *SQLServer) gatherServer(pool *sql.DB, query Query, acc telegraf.Accumulator, connectionString string) error {
 	// execute query
-	rows, err := pool.Query(query.Script)
+	ctx := context.Background()
+	// Use the query timeout if any
+	if s.QueryTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(s.QueryTimeout))
+		defer cancel()
+	}
+	rows, err := pool.QueryContext(ctx, query.Script)
 	if err != nil {
 		serverName, databaseName := getConnectionIdentifiers(connectionString)