(arrow-go) branch main updated: fix(flightsql): stop row streaming after close (#1151)

[zeroshade]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git


The following commit(s) were added to refs/heads/main by this push:
     new 8b131cdb fix(flightsql): stop row streaming after close (#1151)
8b131cdb is described below

commit 8b131cdb3c8ebcd38dc5471158f6c7a6e23df78b
Author: Minh Vu <vuhoangmin...@gmail.com>
AuthorDate: Wed Aug 12 20:53:03 2026 +0200

    fix(flightsql): stop row streaming after close (#1151)
    
    ## What
    
    Rows.Close cancels the streaming context, but the producer could already
    be blocked sending a retained record into the full channel. This adds a
    context-aware send and releases the record when cancellation wins.
    
    ## Test
    
    - go test ./arrow/flight/flightsql/driver -run
    TestRowsSendRecordStopsWhenContextCancelled -count=1
---
 arrow/flight/flightsql/driver/driver.go    | 17 +++++++-
 arrow/flight/flightsql/driver/rows_test.go | 63 ++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

diff --git a/arrow/flight/flightsql/driver/driver.go 
b/arrow/flight/flightsql/driver/driver.go
index c80dfe9a..0e6a8abb 100644
--- a/arrow/flight/flightsql/driver/driver.go
+++ b/arrow/flight/flightsql/driver/driver.go
@@ -100,6 +100,16 @@ func (r *Rows) releaseRecord() {
        }
 }
 
+func (r *Rows) sendRecord(ctx context.Context, record arrow.RecordBatch) bool {
+       select {
+       case r.recordChan <- record:
+               return true
+       case <-ctx.Done():
+               record.Release()
+               return false
+       }
+}
+
 // Close closes the rows iterator.
 func (r *Rows) Close() error {
        r.ctxCancelFunc() // interrupting data streaming.
@@ -107,6 +117,9 @@ func (r *Rows) Close() error {
        r.currentRow = 0
 
        r.releaseRecord()
+       for record := range r.recordChan {
+               record.Release()
+       }
 
        return nil
 }
@@ -587,7 +600,9 @@ func (r *Rows) streamRecordset(ctx context.Context, c 
*flightsql.Client, endpoin
                                        continue
                                }
 
-                               r.recordChan <- record
+                               if !r.sendRecord(ctx, record) {
+                                       return
+                               }
 
                                go initializeOnceOnly.Do(func() { 
r.initializedChan <- true })
                        }
diff --git a/arrow/flight/flightsql/driver/rows_test.go 
b/arrow/flight/flightsql/driver/rows_test.go
new file mode 100644
index 00000000..47a2cc74
--- /dev/null
+++ b/arrow/flight/flightsql/driver/rows_test.go
@@ -0,0 +1,63 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package driver
+
+import (
+       "context"
+       "testing"
+
+       "github.com/apache/arrow-go/v18/arrow"
+       "github.com/apache/arrow-go/v18/arrow/array"
+       "github.com/apache/arrow-go/v18/arrow/memory"
+       "github.com/stretchr/testify/require"
+)
+
+func TestRowsCloseReleasesRetainedRecords(t *testing.T) {
+       mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+       defer mem.AssertSize(t, 0)
+
+       schema := arrow.NewSchema([]arrow.Field{{Name: "value", Type: 
arrow.PrimitiveTypes.Int64}}, nil)
+       builder := array.NewRecordBuilder(mem, schema)
+       builder.Field(0).(*array.Int64Builder).Append(1)
+       queued := builder.NewRecordBatch()
+       builder.Field(0).(*array.Int64Builder).Append(2)
+       pending := builder.NewRecordBatch()
+       builder.Release()
+
+       rows := newRows()
+       ctx, cancel := context.WithCancel(context.Background())
+       rows.ctxCancelFunc = cancel
+
+       queued.Retain()
+       pending.Retain()
+       rows.recordChan <- queued
+
+       done := make(chan struct{})
+       go func() {
+               defer close(done)
+               defer close(rows.recordChan)
+               rows.sendRecord(ctx, pending)
+       }()
+
+       queued.Release()
+       pending.Release()
+
+       require.Positive(t, mem.CurrentAlloc())
+       require.NoError(t, rows.Close())
+       <-done
+       require.Zero(t, mem.CurrentAlloc())
+}