This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new b3d1da9 feat(go/adbc/driver/flightsql): bump max incoming message
size (#402)
b3d1da9 is described below
commit b3d1da9ffaf38152382a84a0b42eea4333393a84
Author: David Li <[email protected]>
AuthorDate: Thu Jan 26 14:53:01 2023 -0500
feat(go/adbc/driver/flightsql): bump max incoming message size (#402)
Fixes #400.
---
docs/source/driver/go/flight_sql.rst | 33 ++++----
go/adbc/driver/flightsql/flightsql_adbc.go | 59 +++++++++++++-
go/adbc/driver/flightsql/flightsql_adbc_test.go | 102 +++++++++++++++++++++++-
go/adbc/driver/flightsql/flightsql_statement.go | 2 +-
4 files changed, 176 insertions(+), 20 deletions(-)
diff --git a/docs/source/driver/go/flight_sql.rst
b/docs/source/driver/go/flight_sql.rst
index a598ce4..bc90a8a 100644
--- a/docs/source/driver/go/flight_sql.rst
+++ b/docs/source/driver/go/flight_sql.rst
@@ -113,28 +113,33 @@ Client Options
The options used for creating the Flight RPC client can be customized.
These options map 1:1 with the options in FlightClientOptions:
-``arrow.flight.sql.client_option.mtls_cert_chain``
+``adbc.flight.sql.client_option.mtls_cert_chain``
The certificate chain to use for mTLS.
-``arrow.flight.sql.client_option.mtls_private_key``
+``adbc.flight.sql.client_option.mtls_private_key``
The private key to use for mTLS.
-``arrow.flight.sql.client_option.tls_override_hostname``
+``adbc.flight.sql.client_option.tls_override_hostname``
Override the hostname used to verify the server's TLS certificate.
-``arrow.flight.sql.client_option.tls_skip_verify``
+``adbc.flight.sql.client_option.tls_skip_verify``
Disable verification of the server's TLS certificate. Value
should be ``true`` or ``false``.
-``arrow.flight.sql.client_option.tls_root_certs``
+``adbc.flight.sql.client_option.tls_root_certs``
Override the root certificates used to validate the server's TLS
certificate.
-``arrow.flight.sql.client_option.generic_int_option.<OPTION_NAME>``
- Option prefixes used to specify generic transport-layer options.
+``adbc.flight.sql.client_option.with_block``
+ Whether connections should wait until connections are established,
+ or connect lazily when used. The latter is gRPC's default
+ behavior, but the driver defaults to eager connection to surface
+ errors earlier. Value should be ``true`` or ``false``.
-``arrow.flight.sql.client_option.generic_string_option.<OPTION_NAME>``
- Option prefixes used to specify generic transport-layer options.
+``adbc.flight.sql.client_option.with_max_msg_size``
+ The maximum message size to accept from the server. The driver
+ defaults to 16 MiB since Flight services tend to return larger
+ reponse payloads. Should be a positive integer number of bytes.
Custom Call Headers
-------------------
@@ -143,7 +148,7 @@ Custom HTTP headers can be attached to requests via options
that apply
to :cpp:class:`AdbcDatabase`, :cpp:class:`AdbcConnection`, and
:cpp:class:`AdbcStatement`.
-``arrow.flight.sql.rpc.call_header.<HEADER NAME>``
+``adbc.flight.sql.rpc.call_header.<HEADER NAME>``
Add the header ``<HEADER NAME>`` to outgoing requests with the given
value.
@@ -170,7 +175,7 @@ of the partitions.
The queue size can be changed by setting an option on the
:cpp:class:`AdbcStatement`:
-``arrow.flight.sql.rpc.queue_size``
+``adbc.flight.sql.rpc.queue_size``
The number of batches to queue per partition. Defaults to 5.
Metadata
@@ -203,14 +208,14 @@ special options on :cpp:class:`AdbcConnection`. In
general, it is
best practice to set timeouts to avoid unexpectedly getting stuck.
The options are as follows:
-``arrow.flight.sql.rpc.timeout_seconds.fetch``
+``adbc.flight.sql.rpc.timeout_seconds.fetch``
A timeout (in floating-point seconds) for any API calls that fetch
data. This corresponds to Flight ``DoGet`` calls.
For example, this controls the timeout of the underlying Flight
calls that fetch more data as a result set is consumed.
-``arrow.flight.sql.rpc.timeout_seconds.query``
+``adbc.flight.sql.rpc.timeout_seconds.query``
A timeout (in floating-point seconds) for any API calls that
execute a query. This corresponds to Flight ``GetFlightInfo``
calls.
@@ -218,7 +223,7 @@ The options are as follows:
For example, this controls the timeout of the underlying Flight
calls that implement :func:`AdbcStatementExecuteQuery`.
-``arrow.flight.sql.rpc.timeout_seconds.update``
+``adbc.flight.sql.rpc.timeout_seconds.update``
A timeout (in floating-point seconds) for any API calls that
upload data or perform other updates.
diff --git a/go/adbc/driver/flightsql/flightsql_adbc.go
b/go/adbc/driver/flightsql/flightsql_adbc.go
index 4c51a7b..5ffa88b 100644
--- a/go/adbc/driver/flightsql/flightsql_adbc.go
+++ b/go/adbc/driver/flightsql/flightsql_adbc.go
@@ -70,6 +70,8 @@ const (
OptionSSLOverrideHostname =
"adbc.flight.sql.client_option.tls_override_hostname"
OptionSSLSkipVerify =
"adbc.flight.sql.client_option.tls_skip_verify"
OptionSSLRootCerts =
"adbc.flight.sql.client_option.tls_root_certs"
+ OptionWithBlock = "adbc.flight.sql.client_option.with_block"
+ OptionWithMaxMsgSize =
"adbc.flight.sql.client_option.with_max_msg_size"
OptionAuthorizationHeader = "adbc.flight.sql.authorization_header"
OptionTimeoutFetch = "adbc.flight.sql.rpc.timeout_seconds.fetch"
OptionTimeoutQuery = "adbc.flight.sql.rpc.timeout_seconds.query"
@@ -129,15 +131,37 @@ func (d Driver) NewDatabase(opts map[string]string)
(adbc.Database, error) {
return nil, adbc.Error{Msg: err.Error(), Code:
adbc.StatusInvalidArgument}
}
+ // Use WithBlock to surface connection errors eagerly
+ // Use WithMaxMsgSize(16 MiB) since Flight services tend to send large
messages
+ db.dialOpts.block = true
+ db.dialOpts.maxMsgSize = 16 * 1024 * 1024
+
return db, db.SetOptions(opts)
}
+type dbDialOpts struct {
+ opts []grpc.DialOption
+ block bool
+ maxMsgSize int
+}
+
+func (d *dbDialOpts) rebuild() {
+ d.opts = []grpc.DialOption{
+
grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(d.maxMsgSize),
+ grpc.MaxCallSendMsgSize(d.maxMsgSize)),
+ }
+ if d.block {
+ d.opts = append(d.opts, grpc.WithBlock())
+ }
+}
+
type database struct {
uri *url.URL
creds credentials.TransportCredentials
user, pass string
hdrs metadata.MD
timeout timeoutOption
+ dialOpts dbDialOpts
alloc memory.Allocator
}
@@ -262,6 +286,38 @@ func (d *database) SetOptions(cnOptions map[string]string)
error {
}
}
+ if val, ok := cnOptions[OptionWithBlock]; ok {
+ if val == adbc.OptionValueEnabled {
+ d.dialOpts.block = true
+ } else if val == adbc.OptionValueDisabled {
+ d.dialOpts.block = false
+ } else {
+ return adbc.Error{
+ Msg: fmt.Sprintf("Invalid value for database
option '%s': '%s'", OptionWithBlock, val),
+ Code: adbc.StatusInvalidArgument,
+ }
+ }
+ delete(cnOptions, OptionWithBlock)
+ }
+ if val, ok := cnOptions[OptionWithMaxMsgSize]; ok {
+ var err error
+ var size int
+ if size, err = strconv.Atoi(val); err != nil {
+ return adbc.Error{
+ Msg: fmt.Sprintf("Invalid value for database
option '%s': '%s' is not a positive integer", OptionWithMaxMsgSize, val),
+ Code: adbc.StatusInvalidArgument,
+ }
+ } else if size <= 0 {
+ return adbc.Error{
+ Msg: fmt.Sprintf("Invalid value for database
option '%s': '%s' is not a positive integer", OptionWithMaxMsgSize, val),
+ Code: adbc.StatusInvalidArgument,
+ }
+ }
+ d.dialOpts.maxMsgSize = size
+ delete(cnOptions, OptionWithMaxMsgSize)
+ }
+ d.dialOpts.rebuild()
+
for key, val := range cnOptions {
if strings.HasPrefix(key, OptionRPCCallHeaderPrefix) {
d.hdrs.Append(strings.TrimPrefix(key,
OptionRPCCallHeaderPrefix), val)
@@ -354,8 +410,9 @@ func getFlightClient(ctx context.Context, loc string, d
*database) (*flightsql.C
if uri.Scheme == "grpc" || uri.Scheme == "grpc+tcp" {
creds = insecure.NewCredentials()
}
+ dialOpts := append(d.dialOpts.opts,
grpc.WithTransportCredentials(creds))
- cl, err := flightsql.NewClient(uri.Host, nil, middleware,
grpc.WithTransportCredentials(creds))
+ cl, err := flightsql.NewClient(uri.Host, nil, middleware, dialOpts...)
if err != nil {
return nil, adbc.Error{
Msg: err.Error(),
diff --git a/go/adbc/driver/flightsql/flightsql_adbc_test.go
b/go/adbc/driver/flightsql/flightsql_adbc_test.go
index 398f379..38a2f00 100644
--- a/go/adbc/driver/flightsql/flightsql_adbc_test.go
+++ b/go/adbc/driver/flightsql/flightsql_adbc_test.go
@@ -222,6 +222,7 @@ func TestADBCFlightSQL(t *testing.T) {
suite.Run(t, &validation.ConnectionTests{Quirks: q})
suite.Run(t, &validation.StatementTests{Quirks: q})
+ suite.Run(t, &DefaultDialOptionsTests{Quirks: q})
suite.Run(t, &PartitionTests{Quirks: q})
suite.Run(t, &SSLTests{Quirks: q})
suite.Run(t, &StatementTests{Quirks: q})
@@ -231,6 +232,99 @@ func TestADBCFlightSQL(t *testing.T) {
// Driver-specific tests
+type DefaultDialOptionsTests struct {
+ suite.Suite
+
+ Driver adbc.Driver
+ Quirks validation.DriverQuirks
+
+ ctx context.Context
+ DB adbc.Database
+}
+
+func (suite *DefaultDialOptionsTests) SetupSuite() {
+ suite.Driver = suite.Quirks.SetupDriver(suite.T())
+
+ var err error
+ suite.ctx = context.Background()
+ suite.DB, err = suite.Driver.NewDatabase(suite.Quirks.DatabaseOptions())
+ suite.NoError(err)
+
+ cnxn, err := suite.DB.Open(suite.ctx)
+ suite.NoError(err)
+ defer cnxn.Close()
+
+ stmt, err := cnxn.NewStatement()
+ suite.NoError(err)
+ defer stmt.Close()
+
+ // Construct huge table
+ suite.NoError(stmt.SetSqlQuery("CREATE TABLE huge (str)"))
+ _, err = stmt.ExecuteUpdate(suite.ctx)
+ suite.NoError(err)
+
+ // 4 KiB
+ suite.NoError(stmt.SetSqlQuery("INSERT INTO huge (str) VALUES
(printf('%.*c', 4096, '!'))"))
+ _, err = stmt.ExecuteUpdate(suite.ctx)
+ suite.NoError(err)
+
+ // 4 MiB
+ suite.NoError(stmt.SetSqlQuery("INSERT INTO huge (str) SELECT * FROM
huge"))
+ for i := 0; i < 10; i++ {
+ _, err = stmt.ExecuteUpdate(suite.ctx)
+ suite.NoError(err)
+ }
+}
+
+func (suite *DefaultDialOptionsTests) TearDownSuite() {
+ suite.Quirks.TearDownDriver(suite.T(), suite.Driver)
+ suite.DB = nil
+ suite.Driver = nil
+}
+
+func (suite *DefaultDialOptionsTests) TestMaxIncomingMessageSizeDefault() {
+ opts := suite.Quirks.DatabaseOptions()
+ opts["adbc.flight.sql.client_option.with_max_msg_size"] = "1000000"
+ db, err := suite.Driver.NewDatabase(opts)
+ suite.NoError(err)
+
+ cnxn, err := db.Open(suite.ctx)
+ suite.NoError(err)
+ defer cnxn.Close()
+
+ stmt, err := cnxn.NewStatement()
+ suite.NoError(err)
+ defer stmt.Close()
+
+ suite.NoError(stmt.SetSqlQuery("SELECT * FROM huge"))
+ reader, _, err := stmt.ExecuteQuery(suite.ctx)
+ suite.NoError(err)
+ defer reader.Release()
+
+ for reader.Next() {
+ }
+ suite.ErrorContains(reader.Err(), "received message larger than max")
+}
+
+func (suite *DefaultDialOptionsTests) TestMaxIncomingMessageSizeLow() {
+ cnxn, err := suite.DB.Open(suite.ctx)
+ suite.NoError(err)
+ defer cnxn.Close()
+
+ stmt, err := cnxn.NewStatement()
+ suite.NoError(err)
+ defer stmt.Close()
+
+ suite.NoError(stmt.SetSqlQuery("SELECT * FROM huge"))
+ reader, _, err := stmt.ExecuteQuery(suite.ctx)
+ suite.NoError(err)
+ defer reader.Release()
+
+ for reader.Next() {
+ }
+ suite.NoError(reader.Err())
+}
+
type SSLTests struct {
suite.Suite
@@ -388,16 +482,16 @@ func (suite *StatementTests) TearDownTest() {
func (suite *StatementTests) TestQueueSizeOption() {
var err error
- option := "arrow.flight.sql.rpc.queue_size"
+ option := "adbc.flight.sql.rpc.queue_size"
err = suite.Stmt.SetOption(option, "")
- suite.Require().ErrorContains(err, "Invalid value for statement option
'arrow.flight.sql.rpc.queue_size': '' is not a positive integer")
+ suite.Require().ErrorContains(err, "Invalid value for statement option
'adbc.flight.sql.rpc.queue_size': '' is not a positive integer")
err = suite.Stmt.SetOption(option, "foo")
- suite.Require().ErrorContains(err, "Invalid value for statement option
'arrow.flight.sql.rpc.queue_size': 'foo' is not a positive integer")
+ suite.Require().ErrorContains(err, "Invalid value for statement option
'adbc.flight.sql.rpc.queue_size': 'foo' is not a positive integer")
err = suite.Stmt.SetOption(option, "-1")
- suite.Require().ErrorContains(err, "Invalid value for statement option
'arrow.flight.sql.rpc.queue_size': '-1' is not a positive integer")
+ suite.Require().ErrorContains(err, "Invalid value for statement option
'adbc.flight.sql.rpc.queue_size': '-1' is not a positive integer")
err = suite.Stmt.SetOption(option, "1")
suite.Require().NoError(err)
diff --git a/go/adbc/driver/flightsql/flightsql_statement.go
b/go/adbc/driver/flightsql/flightsql_statement.go
index ccd0598..45010d3 100644
--- a/go/adbc/driver/flightsql/flightsql_statement.go
+++ b/go/adbc/driver/flightsql/flightsql_statement.go
@@ -35,7 +35,7 @@ import (
)
const (
- OptionStatementQueueSize = "arrow.flight.sql.rpc.queue_size"
+ OptionStatementQueueSize = "adbc.flight.sql.rpc.queue_size"
)
type statement struct {
