Merge #96607

96607: sql: add ability to not apply statement timeout to DECLARE CURSOR r=rafiss a=otan

Epic: none

Currently, `DECLARE CURSOR ... FETCH` cycles which take longer than the statement timeout would error. As a workaround, we introduce a `declare_cursor_statement_timeout_enabled` session variable which disables the timeout for `DECLARE CURSOR` only.

Works around #96322.

Release note (sql change): Introduce the
`declare_cursor_statement_timeout_enabled` session variable which disables statement timeouts during FETCH when using DECLARE CURSOR.

Co-authored-by: Oliver Tan <[email protected]>

pkg/sql/exec_util.go

@@ -3372,6 +3372,10 @@ func (m *sessionDataMutator) SetCopyFromRetriesEnabled(val bool) {
 	m.data.CopyFromRetriesEnabled = val
 }
 
+func (m *sessionDataMutator) SetDeclareCursorStatementTimeoutEnabled(val bool) {
+	m.data.DeclareCursorStatementTimeoutEnabled = val
+}
+
 func (m *sessionDataMutator) SetEnforceHomeRegion(val bool) {
 	m.data.EnforceHomeRegion = val
 }

pkg/sql/logictest/testdata/logic_test/cursor

@@ -686,3 +686,35 @@ BEGIN;
 declare a cursor for select * from crdb_internal.gossip_network;
 FETCH 1 FROM a;
 COMMIT
+
+# Regression test for statement timeouts during DECLARE CURSOR.
+# We set timeout to 1s so that DECLARE CURSOR can execute, then sleep
+# 0.5s each to make up 1s and attempt to fetch.
+# Setting the timeout to 0.001ms means DECLARE CURSOR times out thus
+# voiding the test.
+statement ok
+SET declare_cursor_statement_timeout_enabled = false;
+BEGIN;
+SET statement_timeout = '1s';
+DECLARE a CURSOR FOR SELECT * FROM ( VALUES (1), (2) ) t(id);
+
+# Note we can't set pg_sleep to 1 or else it'll statement timeout!
+statement ok
+select pg_sleep(0.7)
+
+query I
+FETCH 1 FROM a
+----
+1
+
+statement ok
+select pg_sleep(0.7)
+
+query I
+FETCH 1 FROM a
+----
+2
+
+statement ok
+SET statement_timeout = 0;
+COMMIT

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -4922,6 +4922,7 @@ cost_scans_with_default_col_size                      off
 database                                              test
 datestyle                                             ISO, MDY
 datestyle_enabled                                     on
+declare_cursor_statement_timeout_enabled              on
 default_int_size                                      8
 default_table_access_method                           heap
 default_tablespace                                    ·

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -2566,6 +2566,7 @@ copy_from_retries_enabled                             off                 NULL
 cost_scans_with_default_col_size                      off                 NULL      NULL        NULL        string
 database                                              test                NULL      NULL        NULL        string
 datestyle                                             ISO, MDY            NULL      NULL        NULL        string
+declare_cursor_statement_timeout_enabled              on                  NULL      NULL        NULL        string
 default_int_size                                      8                   NULL      NULL        NULL        string
 default_table_access_method                           heap                NULL      NULL        NULL        string
 default_tablespace                                    ·                   NULL      NULL        NULL        string
@@ -2709,6 +2710,7 @@ copy_from_retries_enabled                             off                 NULL
 cost_scans_with_default_col_size                      off                 NULL  user     NULL      off                 off
 database                                              test                NULL  user     NULL      ·                   test
 datestyle                                             ISO, MDY            NULL  user     NULL      ISO, MDY            ISO, MDY
+declare_cursor_statement_timeout_enabled              on                  NULL  user     NULL      on                  on
 default_int_size                                      8                   NULL  user     NULL      8                   8
 default_table_access_method                           heap                NULL  user     NULL      heap                heap
 default_tablespace                                    ·                   NULL  user     NULL      ·                   ·
@@ -2848,6 +2850,7 @@ cost_scans_with_default_col_size                      NULL    NULL     NULL
 crdb_version                                          NULL    NULL     NULL     NULL        NULL
 database                                              NULL    NULL     NULL     NULL        NULL
 datestyle                                             NULL    NULL     NULL     NULL        NULL
+declare_cursor_statement_timeout_enabled              NULL    NULL     NULL     NULL        NULL
 default_int_size                                      NULL    NULL     NULL     NULL        NULL
 default_table_access_method                           NULL    NULL     NULL     NULL        NULL
 default_tablespace                                    NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -40,6 +40,7 @@ copy_from_retries_enabled                             off
 cost_scans_with_default_col_size                      off
 database                                              test
 datestyle                                             ISO, MDY
+declare_cursor_statement_timeout_enabled              on
 default_int_size                                      8
 default_table_access_method                           heap
 default_tablespace                                    ·

pkg/sql/sessiondatapb/local_only_session_data.proto

@@ -326,6 +326,9 @@ message LocalOnlySessionData {
   // CopyFromRetriesEnabled controls whether retries should be internally
   // attempted for retriable errors.
   bool copy_from_retries_enabled = 89;
+  // DeclareCursorStatementTimeoutEnabled controls whether statement timeouts
+  // apply during DECLARE CURSOR.
+  bool declare_cursor_statement_timeout_enabled = 90;
 
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //

pkg/sql/sql_cursor.go

@@ -48,7 +48,17 @@ func (p *planner) DeclareCursor(ctx context.Context, s *tree.DeclareCursor) (pla
 				return nil, pgerror.Newf(pgcode.NoActiveSQLTransaction, "DECLARE CURSOR can only be used in transaction blocks")
 			}
 
-			ie := p.ExecCfg().InternalDB.NewInternalExecutor(p.SessionData())
+			sd := p.SessionData()
+			// This session variable was introduced as a workaround to #96322.
+			// Today, if a timeout is set, FETCH's timeout is from the point
+			// DECLARE CURSOR is executed rather than the FETCH itself.
+			// The setting allows us to override the setting without affecting
+			// third-party applications.
+			if !p.SessionData().DeclareCursorStatementTimeoutEnabled {
+				sd = sd.Clone()
+				sd.StmtTimeout = 0
+			}
+			ie := p.ExecCfg().InternalDB.NewInternalExecutor(sd)
 			if cursor := p.sqlCursors.getCursor(s.Name); cursor != nil {
 				return nil, pgerror.Newf(pgcode.DuplicateCursor, "cursor %q already exists", s.Name)
 			}

pkg/sql/vars.go

@@ -2340,6 +2340,23 @@ var varGen = map[string]sessionVar{
 		GlobalDefault: globalFalse,
 	},
 
+	// CockroachDB extension.
+	`declare_cursor_statement_timeout_enabled`: {
+		GetStringVal: makePostgresBoolGetStringValFn(`declare_cursor_statement_timeout_enabled`),
+		Set: func(_ context.Context, m sessionDataMutator, s string) error {
+			b, err := paramparse.ParseBoolVar("declare_cursor_statement_timeout_enabled", s)
+			if err != nil {
+				return err
+			}
+			m.SetDeclareCursorStatementTimeoutEnabled(b)
+			return nil
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			return formatBoolAsPostgresSetting(evalCtx.SessionData().DeclareCursorStatementTimeoutEnabled), nil
+		},
+		GlobalDefault: globalTrue,
+	},
+
 	// CockroachDB extension.
 	`enforce_home_region`: {
 		GetStringVal: makePostgresBoolGetStringValFn(`enforce_home_region`),