Skip to content

Commit 3194e44

Browse files
authored
fix(arrow-go/flightsql): route QueryContext through active transaction (#692)
I used GPT-5.3-Codex to generate this patch to fix #667. I have reviewed the outputs. --- ### Rationale for this change `database/sql` queries executed via `tx.QueryContext`/`tx.QueryRowContext` were not consistently bound to the active Flight SQL transaction when no query arguments were provided. The driver always used `c.client.Execute(...)`, which bypassed transaction context and could not see uncommitted transactional state (e.g., tables created inside the transaction). ### What changes are included in this PR? - Updated `Connection.QueryContext` in the Flight SQL Go driver to use: - `c.txn.Execute(...)` when `c.txn` is active and valid. - `c.client.Execute(...)` otherwise. - This aligns zero-argument query execution semantics with transactional expectations in `database/sql`. ### Are these changes tested? Yes. - Added/used regression coverage in driver tests: - `TestTxQueryContextUsesTransaction` - Validation confirms: - test fails when the fix is removed (cannot see tx-local table), - test passes when the fix is applied. ### Are there any user-facing changes? Yes, behavior is corrected for users of the Go Flight SQL `database/sql` driver: - `tx.QueryContext` / `tx.QueryRowContext` (including zero-arg queries) now execute within the active transaction as expected. - No API or configuration changes.
1 parent 6b44167 commit 3194e44

2 files changed

Lines changed: 55 additions & 1 deletion

File tree

arrow/flight/flightsql/driver/driver.go

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -505,7 +505,15 @@ func (c *Connection) QueryContext(ctx context.Context, query string, args []driv
505505
defer cancel()
506506
}
507507

508-
info, err := c.client.Execute(execCtx, query)
508+
var (
509+
info *flight.FlightInfo
510+
err error
511+
)
512+
if c.txn != nil && c.txn.ID().IsValid() {
513+
info, err = c.txn.Execute(execCtx, query)
514+
} else {
515+
info, err = c.client.Execute(execCtx, query)
516+
}
509517
if err != nil {
510518
return nil, err
511519
}

arrow/flight/flightsql/driver/driver_test.go

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1505,6 +1505,52 @@ func (s *SqlTestSuite) TestTxRollback() {
15051505
wg.Wait()
15061506
}
15071507

1508+
func (s *SqlTestSuite) TestTxQueryContextUsesTransaction() {
1509+
t := s.T()
1510+
1511+
// Create and start the server
1512+
server, addr, err := s.createServer()
1513+
require.NoError(t, err)
1514+
1515+
var wg sync.WaitGroup
1516+
wg.Add(1)
1517+
go func() {
1518+
defer wg.Done()
1519+
require.NoError(s.T(), s.startServer(server))
1520+
}()
1521+
defer s.stopServer(server)
1522+
time.Sleep(100 * time.Millisecond)
1523+
1524+
// Configure client
1525+
cfg := s.Config
1526+
cfg.Address = addr
1527+
db, err := sql.Open("flightsql", cfg.DSN())
1528+
require.NoError(t, err)
1529+
defer db.Close()
1530+
1531+
ctx := context.Background()
1532+
tx, err := db.BeginTx(ctx, nil)
1533+
require.NoError(t, err)
1534+
1535+
// Create table and insert one row inside the transaction.
1536+
_, err = tx.ExecContext(ctx, fmt.Sprintf(s.Statements["create table"], s.TableName))
1537+
require.NoError(t, err)
1538+
_, err = tx.ExecContext(ctx, fmt.Sprintf(s.Statements["insert"], s.TableName, "inside-tx", 123))
1539+
require.NoError(t, err)
1540+
1541+
// This zero-argument query must execute inside the active transaction.
1542+
var count int
1543+
err = tx.QueryRowContext(ctx, fmt.Sprintf("SELECT COUNT(*) FROM %s", s.TableName)).Scan(&count)
1544+
require.NoError(t, err)
1545+
require.Equal(t, 1, count)
1546+
1547+
require.NoError(t, tx.Rollback())
1548+
1549+
// Tear-down server
1550+
s.stopServer(server)
1551+
wg.Wait()
1552+
}
1553+
15081554
func (s *SqlTestSuite) TestTxCommit() {
15091555
t := s.T()
15101556

0 commit comments

Comments
 (0)