Reject concurrent worker session operations (#711)

* Reject concurrent worker session operations

* Fix abandoned session operation cleanup

* Allow cleanup of abandoned session continuations

* Avoid mutating rejected transaction cleanup

* Clean up abandoned metadata operations on rollback

* Keep rejected transaction cleanup from consuming metadata

* Fix lint after main merge

Author: bill-ph

controlplane/provisioning/api_test.go

@@ -511,6 +511,7 @@ func TestProvisionCnpgShard(t *testing.T) {
 	w := store.warehouses["shardco"]
 	if w == nil {
 		t.Fatal("expected warehouse to be created")
+		return
 	}
 	if w.MetadataStore.Kind != configstore.MetadataStoreKindCnpgShard {
 		t.Errorf("metadata store kind = %q, want cnpg-shard", w.MetadataStore.Kind)
@@ -545,6 +546,7 @@ func TestProvisionIcebergExternal(t *testing.T) {
 	w := store.warehouses["extice"]
 	if w == nil {
 		t.Fatal("expected warehouse to be created")
+		return
 	}
 	if w.MetadataStore.Kind != configstore.MetadataStoreKindExternal {
 		t.Errorf("metadata store kind = %q, want external", w.MetadataStore.Kind)
@@ -632,6 +634,7 @@ func TestProvisionDuckLakeExternal(t *testing.T) {
 	w := store.warehouses["extdl"]
 	if w == nil {
 		t.Fatal("expected warehouse to be created")
+		return
 	}
 	if w.MetadataStore.Kind != configstore.MetadataStoreKindExternal {
 		t.Errorf("metadata store kind = %q, want external", w.MetadataStore.Kind)

duckdbservice/activation_test.go

@@ -355,6 +355,7 @@ func TestSessionPoolActivateTenantAllowsSameOwnerSameEpochCredentialRetry(t *tes
 	current := pool.currentActivation()
 	if current == nil {
 		t.Fatal("expected activation to remain present")
+		return
 	}
 	if current.payload.DuckLake.S3AccessKey != "NEW_ACCESS_KEY" {
 		t.Fatalf("expected activation payload to be refreshed, got %q", current.payload.DuckLake.S3AccessKey)

duckdbservice/copy_from_stdin.go

@@ -48,6 +48,11 @@ func (h *FlightSQLHandler) doCopyFromStdin(
 	if err != nil {
 		return err
 	}
+	finishOperation, ok := session.beginOperation()
+	if busyErr := sessionBusyStatus(ok); busyErr != nil {
+		return busyErr
+	}
+	defer finishOperation()
 	endConnWork, ok := session.beginConnWork()
 	if !ok {
 		return status.Error(codes.NotFound, "session closed")

duckdbservice/copy_from_stdin_test.go

@@ -186,6 +186,32 @@ func TestDoCopyFromStdinIngestsCSVAndRunsCOPY(t *testing.T) {
 	}
 }
 
+func TestDoCopyFromStdinRejectsConcurrentSessionOperation(t *testing.T) {
+	session, cleanup := newSessionWithInMemoryDuckDB(t)
+	defer cleanup()
+	finishOperation, ok := session.beginOperation()
+	if !ok {
+		t.Fatal("expected operation to start")
+	}
+	defer finishOperation()
+
+	pool := &SessionPool{sessions: map[string]*Session{session.ID: session}}
+	handler := &FlightSQLHandler{pool: pool}
+	ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs("x-duckgres-session", session.ID))
+
+	first := &flight.FlightData{
+		FlightDescriptor: &flight.FlightDescriptor{
+			Type: flight.DescriptorPATH,
+			Path: []string{flightclient.CopyFromStdinDescriptorPath},
+			Cmd:  []byte("COPY t FROM '" + flightclient.CopyFromStdinPathPlaceholder + "' (FORMAT CSV)"),
+		},
+	}
+	err := handler.doCopyFromStdin(ctx, first, &fakeDoPutStream{ctx: ctx})
+	if status.Code(err) != codes.FailedPrecondition {
+		t.Fatalf("expected concurrent COPY to be rejected, got %v", err)
+	}
+}
+
 func TestDoCopyFromStdinRejectsMissingPlaceholder(t *testing.T) {
 	session, cleanup := newSessionWithInMemoryDuckDB(t)
 	defer cleanup()