Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 108 additions & 9 deletions adbc.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package gospice

import (
"context"
"errors"
"fmt"
"log"
"strings"
Expand Down Expand Up @@ -122,17 +123,60 @@ func (c *SpiceClient) closeADBC() error {
// reader, err := client.SqlWithParams(ctx, "SELECT * FROM table WHERE id = $1 AND name = $2", 123, "test")
// reader, err := client.SqlWithParams(ctx, "SELECT * FROM table WHERE ts = $1", TimestampParam(ts, arrow.Microsecond, "UTC"))
func (c *SpiceClient) SqlWithParams(ctx context.Context, sql string, params ...any) (array.RecordReader, error) {
// Capture the ADBC connection to use under the mutex so a concurrent
// re-authentication cannot swap it out from under this query.
used, err := c.ensureADBC()
if err != nil {
return nil, fmt.Errorf("ADBC client is not initialized and failed to initialize: %w", err)
}

rdr, err := c.execADBCWithBackoff(ctx, used, sql, params...)
if err != nil && isADBCAuthError(err) {
// The ADBC connection authenticates only once, when it is opened: a
// Basic-auth handshake yields a server-side session token that is then
// reused for every prepared statement on that connection. That session
// can be invalidated server-side (e.g. expired after a period of
// inactivity), after which the cached token is rejected on every
// subsequent request and the connection cannot recover on its own.
// Re-open the connection to perform a fresh handshake, then retry once.
fresh, reinitErr := c.reinitADBC(used)
if reinitErr != nil {
return nil, fmt.Errorf("ADBC re-authentication failed: %w (original error: %v)", reinitErr, err)
}
rdr, err = c.execADBCWithBackoff(ctx, fresh, sql, params...)
}
if err != nil {
return nil, err
}

return rdr, nil
}

// ensureADBC returns the current ADBC connection, lazily initializing it if
// necessary. The read and initialization run under adbcMu so the returned
// pointer is a consistent snapshot even while another goroutine may be
// re-authenticating via reinitADBC.
func (c *SpiceClient) ensureADBC() (*ADBCClient, error) {
c.adbcMu.Lock()
defer c.adbcMu.Unlock()

if c.adbcClient == nil {
// Try lazy initialization
if err := c.initADBC(); err != nil {
return nil, fmt.Errorf("ADBC client is not initialized and failed to initialize: %w", err)
return nil, err
}
}
return c.adbcClient, nil
}

// execADBCWithBackoff runs a parameterized ADBC query, retrying transient
// (e.g. Unavailable / Internal) failures with the client's backoff policy.
// Authentication failures are returned as-is (not retried here) so the caller
// can re-establish the connection before retrying.
func (c *SpiceClient) execADBCWithBackoff(ctx context.Context, client *ADBCClient, sql string, params ...any) (array.RecordReader, error) {
var rdr array.RecordReader
err := backoff.Retry(func() error {
var err error
rdr, err = c.queryADBCWithParams(ctx, sql, params...)
rdr, err = c.queryADBCWithParams(ctx, client, sql, params...)
if err != nil {
st, ok := status.FromError(err)
if ok {
Expand All @@ -158,14 +202,69 @@ func (c *SpiceClient) SqlWithParams(ctx context.Context, sql string, params ...a
return rdr, nil
}

// isADBCAuthError reports whether err indicates the ADBC connection's
// credentials/session were rejected by the server (as opposed to a transient
// or query error). Such failures are not recoverable on the existing
// connection and require re-opening it to perform a fresh handshake.
func isADBCAuthError(err error) bool {
if err == nil {
return false
}
// Match both the value and pointer forms of adbc.Error, since the driver
// or wrapping layers may return either.
var adbcErr adbc.Error
if errors.As(err, &adbcErr) && isADBCAuthStatus(adbcErr.Code) {
return true
}
var adbcErrPtr *adbc.Error
if errors.As(err, &adbcErrPtr) && adbcErrPtr != nil && isADBCAuthStatus(adbcErrPtr.Code) {
return true
}
// Fall back to matching the message in case the typed error is not
// propagated through the wrapping chain.
msg := err.Error()
return strings.Contains(msg, "Unauthenticated") || strings.Contains(msg, "Invalid credentials")
Comment thread
lukekim marked this conversation as resolved.
}

// isADBCAuthStatus reports whether an ADBC status code indicates the server
// rejected the connection's credentials or session.
func isADBCAuthStatus(code adbc.Status) bool {
return code == adbc.StatusUnauthenticated || code == adbc.StatusUnauthorized
}

// reinitADBC closes and re-opens the ADBC connection so that the next query
// performs a fresh authentication handshake, returning the connection to use
// for the retry. The stale argument is the connection the caller observed
// failing; if another goroutine has already replaced it with a live one, that
// connection is reused so the connection is only re-opened once.
func (c *SpiceClient) reinitADBC(stale *ADBCClient) (*ADBCClient, error) {
c.adbcMu.Lock()
defer c.adbcMu.Unlock()

// Another caller may have already re-opened the connection we observed as
// stale; if so, reuse theirs rather than churning the connection again.
if c.adbcClient != stale && c.adbcClient != nil {
return c.adbcClient, nil
}

if closeErr := c.closeADBC(); closeErr != nil {
log.Printf("warning: failed to close stale ADBC connection during re-authentication: %v", closeErr)
}
c.adbcClient = nil
if err := c.initADBC(); err != nil {
return nil, err
}
return c.adbcClient, nil
}

// queryADBCWithParams executes a parameterized query using ADBC with prepare/execute pattern
func (c *SpiceClient) queryADBCWithParams(ctx context.Context, sql string, params ...any) (array.RecordReader, error) {
if c.adbcClient == nil || c.adbcClient.conn == nil {
func (c *SpiceClient) queryADBCWithParams(ctx context.Context, client *ADBCClient, sql string, params ...any) (array.RecordReader, error) {
if client == nil || client.conn == nil {
return nil, fmt.Errorf("ADBC connection is not initialized")
}

// Create a prepared statement
stmt, err := c.adbcClient.conn.NewStatement()
stmt, err := client.conn.NewStatement()
if err != nil {
return nil, fmt.Errorf("error creating statement: %w", err)
}
Expand All @@ -188,7 +287,7 @@ func (c *SpiceClient) queryADBCWithParams(ctx context.Context, sql string, param

// Bind parameters if provided
if len(params) > 0 {
if err := c.bindParameters(stmt, params...); err != nil {
if err := c.bindParameters(client, stmt, params...); err != nil {
return nil, fmt.Errorf("error binding parameters: %w", err)
}
}
Expand All @@ -203,7 +302,7 @@ func (c *SpiceClient) queryADBCWithParams(ctx context.Context, sql string, param
}

// bindParameters binds parameters to an ADBC statement
func (c *SpiceClient) bindParameters(stmt adbc.Statement, params ...any) error {
func (c *SpiceClient) bindParameters(client *ADBCClient, stmt adbc.Statement, params ...any) error {
if len(params) == 0 {
return nil
}
Expand Down Expand Up @@ -249,7 +348,7 @@ func (c *SpiceClient) bindParameters(stmt adbc.Statement, params ...any) error {
schema := arrow.NewSchema(fields, nil)

// Create a record builder using the reusable allocator
bldr := array.NewRecordBuilder(c.adbcClient.mem, schema)
bldr := array.NewRecordBuilder(client.mem, schema)
defer bldr.Release()

// Add values to the builders
Expand Down
73 changes: 73 additions & 0 deletions adbc_reauth_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
package gospice

import (
"errors"
"fmt"
"testing"

"github.com/apache/arrow-adbc/go/adbc"
)

func TestIsADBCAuthError(t *testing.T) {
tests := []struct {
name string
err error
want bool
}{
{
name: "nil",
err: nil,
want: false,
},
{
name: "adbc unauthenticated",
err: adbc.Error{Code: adbc.StatusUnauthenticated, Msg: "[FlightSQL] Invalid credentials"},
want: true,
},
{
name: "adbc unauthorized",
err: adbc.Error{Code: adbc.StatusUnauthorized, Msg: "[FlightSQL] forbidden"},
want: true,
},
Comment thread
lukekim marked this conversation as resolved.
{
name: "wrapped adbc unauthenticated (the prepared-statement failure shape)",
err: fmt.Errorf("error preparing statement: %w",
adbc.Error{Code: adbc.StatusUnauthenticated, Msg: "[FlightSQL] Invalid credentials (Unauthenticated; Prepare)"}),
want: true,
},
{
name: "pointer adbc unauthorized",
err: &adbc.Error{Code: adbc.StatusUnauthorized, Msg: "[FlightSQL] forbidden"},
want: true,
},
{
name: "wrapped pointer adbc unauthenticated",
err: fmt.Errorf("error preparing statement: %w",
&adbc.Error{Code: adbc.StatusUnauthenticated, Msg: "[FlightSQL] Invalid credentials"}),
want: true,
},
{
name: "adbc internal is not an auth error",
err: adbc.Error{Code: adbc.StatusInternal, Msg: "boom"},
want: false,
},
{
name: "plain error mentioning Unauthenticated (fallback)",
err: errors.New("rpc error: code = Unauthenticated desc = invalid"),
want: true,
},
{
name: "unrelated transient error",
err: errors.New("connection refused"),
want: false,
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := isADBCAuthError(tt.err); got != tt.want {
t.Errorf("isADBCAuthError(%v) = %v, want %v", tt.err, got, tt.want)
}
})
}
}
2 changes: 2 additions & 0 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"net/http"
"os"
"strings"
"sync"
"time"

"github.com/apache/arrow-go/v18/arrow/flight"
Expand Down Expand Up @@ -40,6 +41,7 @@ type SpiceClient struct {

flightClient flight.Client
adbcClient *ADBCClient
adbcMu sync.Mutex // guards re-initialization of adbcClient
httpClient http.Client
backoffPolicy backoff.BackOff
maxRetries uint
Expand Down
Loading