Add support for optional query_timeout param in dsn

Author: aalan3

README.md

@@ -239,6 +239,16 @@ trino.RegisterCustomClient("otel", otelClient)
 db, err := sql.Open("trino", "https://user@localhost:8080?custom_client=otel")
 ```
 
+##### `queryTimeout`
+
+```
+Type:           time.Duration
+Valid values:   duration string
+Default:        nil
+```
+
+The `queryTimeout` parameter sets a timeout for the query. If the query takes longer than the timeout, it will be cancelled. If it is not set the default context timeout will be used.
+
 #### Examples
 
 ```

trino/trino.go

@@ -85,7 +85,7 @@ func init() {
 
 var (
 	// DefaultQueryTimeout is the default timeout for queries executed without a context.
-	DefaultQueryTimeout = 60 * time.Second
+	DefaultQueryTimeout = 10 * time.Hour
 
 	// DefaultCancelQueryTimeout is the timeout for the request to cancel queries in Trino.
 	DefaultCancelQueryTimeout = 30 * time.Second
@@ -187,6 +187,7 @@ type Config struct {
 	SSLCert                    string            // The SSL cert for TLS verification (optional)
 	AccessToken                string            // An access token (JWT) for authentication (optional)
 	ForwardAuthorizationHeader bool              // Allow forwarding the `accessToken` named query parameter in the authorization header, overwriting the `AccessToken` option, if set (optional)
+	QueryTimeout               *time.Duration    // Configurable timeout for query (optional)
 }
 
 // FormatDSN returns a DSN string from the configuration.
@@ -266,6 +267,10 @@ func (c *Config) FormatDSN() (string, error) {
 	sort.Strings(sessionkv)
 	sort.Strings(credkv)
 
+	if c.QueryTimeout != nil {
+		query.Add("query_timeout", c.QueryTimeout.String())
+	}
+
 	for k, v := range map[string]string{
 		"catalog":            c.Catalog,
 		"schema":             c.Schema,
@@ -295,6 +300,7 @@ type Conn struct {
 	progressUpdaterPeriod      queryProgressCallbackPeriod
 	useExplicitPrepare         bool
 	forwardAuthorizationHeader bool
+	queryTimeout               *time.Duration
 }
 
 var (
@@ -369,6 +375,15 @@ func newConn(dsn string) (*Conn, error) {
 		}
 	}
 
+	var queryTimeout *time.Duration
+	if timeoutStr := query.Get("query_timeout"); timeoutStr != "" {
+		d, err := time.ParseDuration(timeoutStr)
+		if err != nil {
+			return nil, fmt.Errorf("trino: invalid timeout: %w", err)
+		}
+		queryTimeout = &d
+	}
+
 	c := &Conn{
 		baseURL:                    serverURL.Scheme + "://" + serverURL.Host,
 		httpClient:                 *httpClient,
@@ -378,6 +393,7 @@ func newConn(dsn string) (*Conn, error) {
 		kerberosRemoteServiceName:  query.Get(kerberosRemoteServiceNameConfig),
 		useExplicitPrepare:         useExplicitPrepare,
 		forwardAuthorizationHeader: forwardAuthorizationHeader,
+		queryTimeout:               queryTimeout,
 	}
 
 	var user string
@@ -963,9 +979,12 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 	}
 
 	var cancel context.CancelFunc = func() {}
-	if _, ok := ctx.Deadline(); !ok {
+	if st.conn.queryTimeout != nil {
+		ctx, cancel = context.WithTimeout(ctx, *st.conn.queryTimeout)
+	} else if _, ok := ctx.Deadline(); !ok {
 		ctx, cancel = context.WithTimeout(ctx, DefaultQueryTimeout)
 	}
+
 	req, err := st.conn.newRequest(ctx, "POST", st.conn.baseURL+"/v1/statement", strings.NewReader(query), hs)
 	if err != nil {
 		cancel()

trino/trino_test.go

@@ -265,6 +265,19 @@ func TestRegisterCustomClientReserved(t *testing.T) {
 	}
 }
 
+func TestQueryTimeout(t *testing.T) {
+	timeout := 10 * time.Second
+	c := &Config{
+		ServerURI:    "https://foobar@localhost:8090",
+		QueryTimeout: &timeout,
+	}
+	dsn, err := c.FormatDSN()
+	require.NoError(t, err)
+
+	want := "https://foobar@localhost:8090?query_timeout=10s&source=trino-go-client"
+	assert.Equal(t, want, dsn)
+}
+
 func TestRoundTripRetryQueryError(t *testing.T) {
 	testcases := []struct {
 		Name                string
@@ -1973,3 +1986,45 @@ func TestForwardAuthorizationHeader(t *testing.T) {
 
 	assert.NoError(t, db.Close())
 }
+
+func TestQueryTimeoutDeadline(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(200 * time.Millisecond) // Simulate slow response
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer ts.Close()
+
+	testcases := []struct {
+		name          string
+		queryTimeout  string
+		expectedError string
+	}{
+		{
+			name:          "with timeout",
+			queryTimeout:  "100ms",
+			expectedError: "context deadline exceeded",
+		},
+		{
+			name:          "without timeout",
+			queryTimeout:  "10s",
+			expectedError: "EOF", // Default server response
+		},
+		{
+			name:          "bad timeout",
+			queryTimeout:  "abc",
+			expectedError: "trino: invalid timeout", // Default server response
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			println(ts.URL + "?query_timeout=" + tc.queryTimeout)
+			db, err := sql.Open("trino", ts.URL+"?query_timeout="+tc.queryTimeout)
+			require.NoError(t, err)
+			defer db.Close()
+
+			_, err = db.Query("SELECT 1")
+			assert.ErrorContains(t, err, tc.expectedError)
+		})
+	}
+}