Every statement on the pipeline's DuckDB connection holds the shared …

internal/cli/examples_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -166,15 +167,18 @@ func TestConfigValidation_ExampleConfigsBuildRealComponents(t *testing.T) {
 			// nothing the commands above may have failed to create. Building
 			// it second meant a handler that skipped for a missing ATTACHed
 			// table took the sink check down with it.
-			_, err = sinks.New(ctx, conf.Pipeline.Sink, conn)
+			// The run command hands every sink the connection's lock, and a
+			// sqlcommand sink refuses to build without one.
+			lock := &sync.Mutex{}
+			_, err = sinks.New(ctx, conf.Pipeline.Sink, conn, sinks.WithConnLock(lock))
 			checkBuildError(t, "sink", err)
 
 			if conf.Tables != nil {
 				for _, table := range conf.Tables.SQL {
 					if table.Window == nil {
 						continue
 					}
-					_, err := sinks.New(ctx, table.Window.Sink, conn)
+					_, err := sinks.New(ctx, table.Window.Sink, conn, sinks.WithConnLock(lock))
 					checkBuildError(t, "window sink for "+table.Name, err)
 				}
 			}

internal/cli/run/managers.go

@@ -110,10 +110,16 @@ func buildManagedTables(
 				"table %q: opening the window sink's connection", table.Name)
 		}
 		conns = append(conns, sinkConn)
+		// The sink's connection is the window's alone, and the manager runs
+		// the sink from one goroutine, so the lock a sqlcommand sink needs
+		// has no other party. It is a lock of its own rather than the
+		// pipeline's: taking the pipeline's here would stall the consume
+		// loop for the length of every window write.
 		sink, err := sinks.New(ctx, table.Window.Sink, sinkConn,
 			sinks.WithMeterProvider(mp),
 			sinks.WithSinkRole("manager"),
-			sinks.WithRetryEvents(events))
+			sinks.WithRetryEvents(events),
+			sinks.WithConnLock(&sync.Mutex{}))
 		if err != nil {
 			return nil, closeConns, fmt.Errorf("table %q window sink: %w", table.Name, err)
 		}

internal/cli/run/root.go

@@ -37,6 +37,7 @@ func newErrorPolicies(
 	ctx context.Context,
 	conf *config.Conf,
 	conn adbc.Connection,
+	lock *sync.Mutex,
 	mp metric.MeterProvider,
 	events sinks.RetryEvents,
 ) (core.PipelineErrorPolicies, error) {
@@ -63,7 +64,8 @@ func newErrorPolicies(
 		dlqSink, err := sinks.New(ctx, *onError.DLQ, conn,
 			sinks.WithMeterProvider(mp),
 			sinks.WithSinkRole("dlq"),
-			sinks.WithRetryEvents(events))
+			sinks.WithRetryEvents(events),
+			sinks.WithConnLock(lock))
 		if err != nil {
 			return policies, fmt.Errorf("pipeline.on_error dlq: %w", err)
 		}
@@ -73,6 +75,23 @@ func newErrorPolicies(
 	return policies, nil
 }
 
+// newPipelineSink builds the pipeline's own sink. A function of its own so the
+// shared-connection test builds it exactly as the run command does.
+func newPipelineSink(
+	ctx context.Context,
+	conf *config.Conf,
+	conn adbc.Connection,
+	lock *sync.Mutex,
+	mp metric.MeterProvider,
+	events sinks.RetryEvents,
+) (core.Sink, error) {
+	return sinks.New(ctx, conf.Pipeline.Sink, conn,
+		sinks.WithMeterProvider(mp),
+		sinks.WithSinkRole(core.SinkRolePipeline),
+		sinks.WithRetryEvents(events),
+		sinks.WithConnLock(lock))
+}
+
 func NewCommand() *cobra.Command {
 	var configPath string
 	var maxMsgs int
@@ -375,10 +394,7 @@ func NewCommand() *cobra.Command {
 
 			// The signal context, so a SIGTERM arriving while the sink dials
 			// its destination stops the start instead of waiting it out.
-			sink, err := sinks.New(ctx, conf.Pipeline.Sink, conn,
-				sinks.WithMeterProvider(meterProvider),
-				sinks.WithSinkRole(core.SinkRolePipeline),
-				sinks.WithRetryEvents(retryEvents))
+			sink, err := newPipelineSink(ctx, conf, conn, lock, meterProvider, retryEvents)
 			if err != nil {
 				return err
 			}
@@ -401,7 +417,7 @@ func NewCommand() *cobra.Command {
 				}()
 			}
 
-			errorPolicies, err := newErrorPolicies(ctx, conf, conn, meterProvider, retryEvents)
+			errorPolicies, err := newErrorPolicies(ctx, conf, conn, lock, meterProvider, retryEvents)
 			if err != nil {
 				return err
 			}

internal/cli/run/row_roles_test.go

@@ -3,6 +3,7 @@ package run
 import (
 	"context"
 	"os"
+	"sync"
 	"testing"
 
 	"github.com/apache/arrow-adbc/go/adbc"
@@ -112,7 +113,7 @@ func TestDLQRowsCarryTheDLQRole(t *testing.T) {
 		DLQ:    &config.Sink{Type: "console"},
 	}
 
-	policies, err := newErrorPolicies(context.Background(), conf, nil, mp, sinks.RetryEvents{})
+	policies, err := newErrorPolicies(context.Background(), conf, nil, &sync.Mutex{}, mp, sinks.RetryEvents{})
 	assert.NoError(t, err)
 	assert.That(t, policies.DLQSink != nil)
 

internal/cli/run/shared_conn_test.go

@@ -0,0 +1,168 @@
+package run
+
+import (
+	"context"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/apache/arrow-adbc/go/adbc"
+	"github.com/apache/arrow-go/v18/arrow/array"
+	"github.com/turbolytics/sql-flow/internal/config"
+	"github.com/turbolytics/sql-flow/internal/core"
+	"github.com/turbolytics/sql-flow/internal/coverage"
+	"github.com/turbolytics/sql-flow/internal/duckdb"
+	"github.com/turbolytics/sql-flow/internal/handlers"
+	"github.com/turbolytics/sql-flow/internal/sinks"
+	"github.com/zeebo/assert"
+	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
+	"go.uber.org/zap"
+)
+
+// sliceSource delivers fixed batches and then closes its stream.
+type sliceSource struct {
+	batches [][]core.Message
+}
+
+func (s *sliceSource) Start() error  { return nil }
+func (s *sliceSource) Commit() error { return nil }
+func (s *sliceSource) Close() error  { return nil }
+
+func (s *sliceSource) Stream() <-chan []core.Message {
+	ch := make(chan []core.Message, len(s.batches))
+	for _, b := range s.batches {
+		ch <- b
+	}
+	close(ch)
+	return ch
+}
+
+func sharedConnCount(t *testing.T, conn adbc.Connection, table string) int64 {
+	t.Helper()
+	stmt, err := conn.NewStatement()
+	assert.NoError(t, err)
+	defer stmt.Close()
+	assert.NoError(t, stmt.SetSqlQuery("SELECT count(*) FROM "+table))
+	reader, _, err := stmt.ExecuteQuery(context.Background())
+	assert.NoError(t, err)
+	defer reader.Release()
+	assert.That(t, reader.Next())
+	return reader.Record().Column(0).(*array.Int64).Value(0)
+}
+
+// Every party that runs statements on the pipeline's DuckDB connection holds
+// the shared lock while it does: the handler, the pipeline sink, the DLQ sink,
+// the progress store, and each table manager with its own sink.
+//
+// DuckDB closes a pending result the moment another statement runs on its
+// connection. A party that skips the lock fails whichever query another party
+// has in flight, with "Attempting to execute an unsuccessful or closed pending
+// query result". #280: the handler reset and the progress write ran outside
+// the lock, the Bluesky demo failed a window poll at batch 500, and at batch 1
+// the failure recurred every few minutes. On main a failed poll stops the
+// process.
+//
+// The race has no deterministic reproduction, so this checks the invariant on
+// every statement instead. The components are built by the same functions the
+// run command calls, over a connection that records any execution finding the
+// lock free. Nothing else touches the connection while the loop and the poll
+// run, so a free lock means a party let go of it.
+//
+// It covers the two places that build a sqlcommand sink on the shared
+// connection, the pipeline's and the DLQ's. sinks.New refuses to build one
+// without the lock; this test is what shows the sink then uses it. The window
+// manager is built beside them and polls, but it runs on connections of its
+// own since #281, so it is not a party to the shared one: the lock checker
+// never sees it, and that is the point.
+func TestPipelineSharedConnection_EveryPartyHoldsTheLock(t *testing.T) {
+	coverage.Covers(t, "manager.window", "sink.sqlcommand", "error.dlq", "handler.inferred_mem")
+	ctx := context.Background()
+
+	db, raw := rowsTestDB(t)
+	rowsTestExec(t, raw, "CREATE TABLE win (bucket TIMESTAMPTZ, n BIGINT)")
+	rowsTestExec(t, raw, "CREATE TABLE published (bucket TIMESTAMPTZ, n BIGINT)")
+	rowsTestExec(t, raw, "CREATE TABLE dead (error VARCHAR, message VARCHAR, phase VARCHAR, timestamp VARCHAR)")
+
+	lock := &sync.Mutex{}
+	conn := duckdb.NewLockChecked(raw, lock)
+	mp := sdkmetric.NewMeterProvider()
+
+	conf := &config.Conf{
+		Tables: &config.Tables{
+			SQL: []config.TableSQL{{
+				Name: "win",
+				// One-second buckets with no grace: a bucket closes once a
+				// later one exists.
+				Window: &config.Window{
+					TimeColumn:       "bucket",
+					SizeSeconds:      1,
+					LateRows:         "drop",
+					PollIntervalSecs: 3600,
+					EmitSQL:          "SELECT bucket, sum(n)::BIGINT AS n FROM closed GROUP BY ALL",
+					Sink: config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{
+						SQL: "INSERT INTO published SELECT bucket, n FROM sqlflow_sink_batch",
+					}},
+				},
+			}},
+		},
+	}
+	conf.Pipeline.Handler = config.Handler{
+		Type: "handlers.InferredMemBatch",
+		SQL:  "SELECT to_timestamp(a) AS bucket, count(*)::BIGINT AS n FROM batch GROUP BY ALL",
+	}
+	conf.Pipeline.Sink = config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{
+		SQL: "INSERT INTO win SELECT bucket, n FROM sqlflow_sink_batch",
+	}}
+	conf.Pipeline.OnError = &config.OnError{
+		Policy: "DLQ",
+		DLQ: &config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{
+			SQL: "INSERT INTO dead SELECT * FROM sqlflow_sink_batch",
+		}},
+	}
+
+	// Startup runs before any second goroutine exists, and the run command
+	// does it without the lock. Holding it here keeps the check about the
+	// running pipeline.
+	progress := core.NewProgressStore(conn)
+	lock.Lock()
+	assert.NoError(t, progress.Init(ctx))
+	assert.NoError(t, initWindowStores(ctx, conf, conn))
+	lock.Unlock()
+
+	sink, err := newPipelineSink(ctx, conf, conn, lock, mp, sinks.RetryEvents{})
+	assert.NoError(t, err)
+	policies, err := newErrorPolicies(ctx, conf, conn, lock, mp, sinks.RetryEvents{})
+	assert.NoError(t, err)
+	handler, err := handlers.New(conn, conf.Pipeline.Handler, zap.NewNop())
+	assert.NoError(t, err)
+	managed, closeConns, err := buildManagedTables(ctx, conf, db, zap.NewNop(), mp,
+		nil, sinks.RetryEvents{})
+	assert.NoError(t, err)
+	defer closeConns()
+	assert.Equal(t, 1, len(managed))
+
+	msg := func(v string) core.Message { return core.Message{Value: []byte(v)} }
+	src := &sliceSource{batches: [][]core.Message{
+		{msg(`{"a": 1}`), msg(`{"a": 1}`)},
+		{msg(`not json`), msg(`{"a": 2}`)},
+		{msg(`{"a": 2}`), msg(`{"a": 3}`)},
+	}}
+	tb := core.NewTurbine(src, handler, sink, 2, time.Hour, lock, policies,
+		core.WithProgressStore(progress))
+
+	_, err = tb.ConsumeLoop(ctx, 6)
+	assert.NoError(t, err)
+	assert.NoError(t, managed[0].Poll(ctx))
+
+	// Every party ran at least once, or a clean result proves nothing. The
+	// window closed the two buckets a later one exists for, and the newest
+	// stays open.
+	assert.Equal(t, int64(2), sharedConnCount(t, raw, "published"))
+	assert.Equal(t, int64(1), sharedConnCount(t, raw, "dead"))
+
+	if v := conn.Violations(); len(v) > 0 {
+		t.Fatalf("%d statements ran on the shared connection without the lock:\n%s",
+			len(v), strings.Join(v, "\n"))
+	}
+}

internal/core/lock_test.go

@@ -0,0 +1,96 @@
+package core
+
+import (
+	"context"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/apache/arrow-adbc/go/adbc"
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/turbolytics/sql-flow/internal/coverage"
+	"github.com/turbolytics/sql-flow/internal/duckdb"
+	"github.com/zeebo/assert"
+)
+
+// connHandler is a fakeHandler that touches the connection the way the real
+// handlers do: Init resets a batch table and Invoke queries one. The pipeline
+// decides whether either runs under the lock, and the fake handler that
+// touches nothing hides that decision.
+type connHandler struct {
+	fakeHandler
+	conn adbc.Connection
+}
+
+func (h *connHandler) exec(ctx context.Context, sql string) error {
+	stmt, err := h.conn.NewStatement()
+	if err != nil {
+		return err
+	}
+	defer stmt.Close()
+	if err := stmt.SetSqlQuery(sql); err != nil {
+		return err
+	}
+	_, err = stmt.ExecuteUpdate(ctx)
+	return err
+}
+
+func (h *connHandler) Init(ctx context.Context) error {
+	if err := h.fakeHandler.Init(ctx); err != nil {
+		return err
+	}
+	return h.exec(ctx, "DROP TABLE IF EXISTS batch")
+}
+
+func (h *connHandler) Invoke(ctx context.Context) (arrow.Table, error) {
+	if err := h.exec(ctx, "CREATE OR REPLACE TABLE batch AS SELECT 1 AS n"); err != nil {
+		return nil, err
+	}
+	return h.fakeHandler.Invoke(ctx)
+}
+
+// Every statement the pipeline runs on the shared connection runs under the
+// shared lock. The table managers collect on the same connection under that
+// lock, and DuckDB closes a pending result the moment another statement runs
+// on its connection, so a statement outside the lock fails whichever poll is
+// in flight. #280: the handler reset after every batch and the progress write
+// both ran outside the lock, and the Bluesky demo logged "closed pending
+// query result" once at batch 500 and seven times in ten minutes at batch 1.
+//
+// The race cannot be reproduced on demand. The invariant can be checked on
+// every statement: the wrapper records any execution that finds the lock
+// free. The loop runs with no other goroutine touching the connection, so a
+// free lock means the pipeline itself let go of it.
+func TestCoreConsumeLoop_EveryStatementOnTheSharedConnectionHoldsTheLock(t *testing.T) {
+	coverage.Covers(t, "state.durability")
+	ctx := context.Background()
+	db, err := duckdb.OpenPath(ctx, "")
+	assert.NoError(t, err)
+	defer db.Close()
+	raw, err := db.Connect(ctx)
+	assert.NoError(t, err)
+	defer raw.Close()
+
+	lock := &sync.Mutex{}
+	conn := duckdb.NewLockChecked(raw, lock)
+
+	// Startup runs before the managers exist, and root.go does it without
+	// the lock. Holding it here keeps the check about the running pipeline.
+	progress := NewProgressStore(conn)
+	lock.Lock()
+	assert.NoError(t, progress.Init(ctx))
+	lock.Unlock()
+
+	src := &fakeSource{batches: [][]Message{messages(4), messages(4), messages(4)}}
+	tb := NewTurbine(src, &connHandler{conn: conn}, &fakeSink{}, 4, time.Second,
+		lock, PipelineErrorPolicies{}, WithProgressStore(progress))
+	stats, err := tb.ConsumeLoop(ctx, 12)
+	assert.NoError(t, err)
+	assert.Equal(t, int64(12), stats.MessagesConsumed())
+
+	if v := conn.Violations(); len(v) > 0 {
+		t.Fatalf("%d statements ran on the shared connection without the lock:\n%s",
+			len(v), strings.Join(v, "\n"))
+	}
+}