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 f1bd3c43 feat(flightsql): allow configuring the driver memory
allocator (#1076)
f1bd3c43 is described below
commit f1bd3c437750dbf14682e8f5cd76430b39348fff
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 5 20:11:14 2026 +0200
feat(flightsql): allow configuring the driver memory allocator (#1076)
Fixes #59
## Problem
The FlightSQL database driver always used `memory.DefaultAllocator` when
building prepared-statement parameter records. Applications configuring
the driver programmatically could not route these allocations through a
checked, pooled, or otherwise custom allocator.
## Change
Add an optional `Allocator` field to `DriverConfig`.
`Connector.Configure` installs it on the FlightSQL client, and
prepared-statement parameter builders use the client allocator.
Compatibility is unchanged:
- a nil allocator falls back to `memory.DefaultAllocator`
- DSN-based connections continue using the default because an allocator
cannot be encoded in a connection string
- programmatic connector configuration can supply a custom allocator
## Coverage
The prepared-statement test uses a concurrency-safe recording allocator
and verifies that parameter construction allocates through the
configured instance.
## Validation
`go test ./arrow/flight/flightsql/driver`
---
arrow/flight/flightsql/driver/driver.go | 24 ++++++++++++++++++++----
arrow/flight/flightsql/driver/driver_test.go | 23 +++++++++++++++++++++--
2 files changed, 41 insertions(+), 6 deletions(-)
diff --git a/arrow/flight/flightsql/driver/driver.go
b/arrow/flight/flightsql/driver/driver.go
index ee78614b..c80dfe9a 100644
--- a/arrow/flight/flightsql/driver/driver.go
+++ b/arrow/flight/flightsql/driver/driver.go
@@ -314,7 +314,7 @@ func (s *Stmt) setParameters(args []driver.NamedValue)
error {
schema = arrow.NewSchema(fields, nil)
}
- recBuilder := array.NewRecordBuilder(memory.DefaultAllocator, schema)
+ recBuilder := array.NewRecordBuilder(s.client.Alloc, schema)
defer recBuilder.Release()
for i, arg := range args {
@@ -388,16 +388,31 @@ func (d *Driver) OpenConnector(name string)
(driver.Connector, error) {
}
type Connector struct {
- addr string
- timeout time.Duration
- options []grpc.DialOption
+ addr string
+ timeout time.Duration
+ options []grpc.DialOption
+ allocator memory.Allocator
}
// Configure the driver with the corresponding config
func (c *Connector) Configure(config *DriverConfig) error {
+ return c.configure(config, nil)
+}
+
+// ConfigureWithAllocator configures the driver with a custom memory allocator.
+// DSN-based connections continue to use memory.DefaultAllocator.
+func (c *Connector) ConfigureWithAllocator(config *DriverConfig, allocator
memory.Allocator) error {
+ return c.configure(config, allocator)
+}
+
+func (c *Connector) configure(config *DriverConfig, allocator
memory.Allocator) error {
// Set the driver properties
c.addr = config.Address
c.timeout = config.Timeout
+ c.allocator = allocator
+ if allocator == nil {
+ c.allocator = memory.DefaultAllocator
+ }
c.options = []grpc.DialOption{}
// Create GRPC options necessary for the backend
@@ -434,6 +449,7 @@ func (c *Connector) Connect(ctx context.Context)
(driver.Conn, error) {
if err != nil {
return nil, err
}
+ client.Alloc = c.allocator
return &Connection{
client: client,
diff --git a/arrow/flight/flightsql/driver/driver_test.go
b/arrow/flight/flightsql/driver/driver_test.go
index 82c66753..425fd8fd 100644
--- a/arrow/flight/flightsql/driver/driver_test.go
+++ b/arrow/flight/flightsql/driver/driver_test.go
@@ -28,6 +28,7 @@ import (
"os"
"strings"
"sync"
+ "sync/atomic"
"testing"
"time"
@@ -43,6 +44,21 @@ import (
"github.com/apache/arrow-go/v18/arrow/memory"
)
+type recordingAllocator struct {
+ memory.Allocator
+ allocated atomic.Bool
+}
+
+func (a *recordingAllocator) Allocate(size int) []byte {
+ a.allocated.Store(true)
+ return a.Allocator.Allocate(size)
+}
+
+func (a *recordingAllocator) Reallocate(size int, b []byte) []byte {
+ a.allocated.Store(true)
+ return a.Allocator.Reallocate(size, b)
+}
+
const defaultTableName = "drivertest"
var defaultStatements = map[string]string{
@@ -1691,12 +1707,14 @@ func TestPreparedStatementSchema(t *testing.T) {
defer server.Shutdown()
// Configure client
+ alloc := &recordingAllocator{Allocator: memory.DefaultAllocator}
cfg := driver.DriverConfig{
Timeout: 5 * time.Second,
Address: server.Addr().String(),
}
- db, err := sql.Open("flightsql", cfg.DSN())
- require.NoError(t, err)
+ connector := &driver.Connector{}
+ require.NoError(t, connector.ConfigureWithAllocator(&cfg, alloc))
+ db := sql.OpenDB(connector)
defer db.Close()
// Do query
@@ -1713,6 +1731,7 @@ func TestPreparedStatementSchema(t *testing.T) {
rows, err := stmt.Query("master")
require.NoError(t, err)
require.NotNil(t, rows)
+ require.True(t, alloc.allocated.Load())
}
func TestPreparedStatementNoSchema(t *testing.T) {
