Merge pull request #153796 from yuzefovich/blathers/backport-release-25.2-153783

release-25.2: sql: fix a few edge case bugs with CHECK EXTERNAL CONNECTION

Author: yuzefovich

pkg/sql/BUILD.bazel

@@ -638,6 +638,7 @@ go_test(
         "backfill_test.go",
         "builtin_mem_usage_test.go",
         "builtin_test.go",
+        "check_external_connection_test.go",
         "check_test.go",
         "closed_session_cache_test.go",
         "comment_on_column_test.go",

pkg/sql/check_external_connection.go

@@ -50,7 +50,7 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 		return err
 	}
 
-	ctx, span := tracing.ChildSpan(params.ctx, "CheckExternalConnection")
+	ctx, span := tracing.ChildSpan(params.ctx, "CheckExternalConnection-planning")
 	defer span.Finish()
 
 	store, err := params.ExecCfg().DistSQLSrv.ExternalStorageFromURI(ctx, n.loc, params.p.User())
@@ -90,7 +90,7 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 			time.Duration(tree.MustBeDInt(nanos)),
 		))
 	}
-	n.rows = make(chan tree.Datums, int64(len(sqlInstanceIDs))*n.params.Concurrency)
+	n.rows = make(chan tree.Datums, len(sqlInstanceIDs)*getCloudCheckConcurrency(n.params))
 	rowWriter := NewCallbackResultWriter(func(ctx context.Context, row tree.Datums) error {
 		// collapse the two pairs of bytes+time to a single string rate each.
 		res := make(tree.Datums, len(row)-1)
@@ -103,9 +103,16 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 		return nil
 	})
 
-	grp := ctxgroup.WithContext(ctx)
-	n.execGrp = grp
-	grp.GoCtx(func(ctx context.Context) error {
+	workerStarted := make(chan struct{})
+	n.execGrp = ctxgroup.WithContext(params.ctx)
+	n.execGrp.GoCtx(func(ctx context.Context) error {
+		// Derive a separate tracing span since the planning one will be
+		// finished when the main goroutine exits from startExec.
+		ctx, span := tracing.ChildSpan(ctx, "CheckExternalConnection-execution")
+		defer span.Finish()
+		// Unblock the main goroutine after having created the tracing span.
+		close(workerStarted)
+
 		recv := MakeDistSQLReceiver(
 			ctx,
 			rowWriter,
@@ -124,13 +131,18 @@ func (n *checkExternalConnectionNode) startExec(params runParams) error {
 		return nil
 	})
 
+	// Block until the worker goroutine has started. This allows us to guarantee
+	// that params.ctx contains a tracing span that hasn't been finished.
+	// TODO(yuzefovich): this is a bit hacky. The issue is that
+	// planNodeToRowSource has already created a new tracing span for this
+	// checkExternalConnectionNode and has updated params.ctx accordingly; then,
+	// if the query is canceled before the worker goroutine starts, the tracing
+	// span is finished, yet it will have already been captured by the ctxgroup.
+	<-workerStarted
 	return nil
 }
 
 func (n *checkExternalConnectionNode) Next(params runParams) (bool, error) {
-	if n.rows == nil {
-		return false, nil
-	}
 	select {
 	case <-params.ctx.Done():
 		return false, params.ctx.Err()
@@ -149,7 +161,6 @@ func (n *checkExternalConnectionNode) Values() tree.Datums {
 
 func (n *checkExternalConnectionNode) Close(_ context.Context) {
 	_ = n.execGrp.Wait()
-	n.rows = nil
 }
 
 func (n *checkExternalConnectionNode) parseParams(params runParams) error {

pkg/sql/check_external_connection_test.go

@@ -0,0 +1,45 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package sql_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+func TestCheckExternalConnection(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	rng, _ := randutil.NewTestRand()
+	dir, dirCleanupFn := testutils.TempDir(t)
+	defer dirCleanupFn()
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		ExternalIODir: dir,
+	})
+	defer s.Stopper().Stop(ctx)
+
+	runner := sqlutils.MakeSQLRunner(sqlDB)
+	runner.Exec(t, "CREATE EXTERNAL CONNECTION foo_conn AS 'nodelocal://1/foo';")
+	query := "CHECK EXTERNAL CONNECTION 'nodelocal://1/foo';"
+	// Should execute successfully without a statement timeout.
+	runner.Exec(t, query)
+	// Run with a random statement timeout which will likely make the query
+	// fail. We don't care whether it does nor which error is returned as long
+	// as the process doesn't crash.
+	runner.Exec(t, fmt.Sprintf("SET statement_timeout='%dms'", rng.Intn(100)+1))
+	_, _ = sqlDB.Exec(query)
+}

pkg/sql/cloud_check_processor.go

@@ -203,20 +203,27 @@ func newCloudCheckProcessor(
 	return p, nil
 }
 
-// Start is part of the RowSource interface.
-func (p *proc) Start(ctx context.Context) {
-	p.StartInternal(ctx, "cloudcheck.proc")
-
-	concurrency := int(p.spec.Params.Concurrency)
+func getCloudCheckConcurrency(params CloudCheckParams) int {
+	concurrency := int(params.Concurrency)
 	if concurrency < 1 {
 		concurrency = 1
 	}
+	return concurrency
+}
 
+// Start is part of the RowSource interface.
+func (p *proc) Start(ctx context.Context) {
+	p.StartInternal(ctx, "cloudcheck.proc")
+
+	concurrency := getCloudCheckConcurrency(p.spec.Params)
 	p.results = make(chan result, concurrency)
 
 	if err := p.FlowCtx.Stopper().RunAsyncTask(p.Ctx(), "cloudcheck.proc", func(ctx context.Context) {
 		defer close(p.results)
-		if err := ctxgroup.GroupWorkers(ctx, concurrency, func(ctx context.Context, _ int) error {
+		// We're ignoring the context cancellation error (which is the only one
+		// possible) because the main goroutine will observe it on its own
+		// anyway in Next.
+		_ = ctxgroup.GroupWorkers(ctx, concurrency, func(ctx context.Context, _ int) error {
 			select {
 			case p.results <- checkURI(
 				ctx,
@@ -229,9 +236,7 @@ func (p *proc) Start(ctx context.Context) {
 			case <-ctx.Done():
 				return ctx.Err()
 			}
-		}); err != nil {
-			p.MoveToDraining(err)
-		}
+		})
 	}); err != nil {
 		p.MoveToDraining(err)
 	}

pkg/sql/execinfra/processorsbase.go

@@ -521,6 +521,9 @@ const (
 //
 // An error can be optionally passed. It will be the first piece of metadata
 // returned by DrainHelper().
+//
+// MoveToDraining should only be called from the main goroutine of the
+// processor.
 func (pb *ProcessorBaseNoHelper) MoveToDraining(err error) {
 	if pb.State != StateRunning {
 		// Calling MoveToDraining in any state is allowed in order to facilitate the