Flight SQL Recipes

Some of these recipes are written againt a demo Flight SQL service backed by SQLite. You can run it yourself as follows:

$ go install github.com/apache/arrow/go/v${ARROW_MAJOR_VERSION}/arrow/flight/flightsql/example/cmd/sqlite_flightsql_server@latest
$ sqlite_flightsql_server -host 0.0.0.0 -port 8080

Other recipes work using the OSS version of Dremio:

$ docker run -p 9047:9047 -p 31010:31010 -p 45678:45678 dremio/dremio-oss

If you have the ADBC repository checked out and Docker Compose installed, you can use our configuration to run both services:

$ docker compose up --detach --wait dremio dremio-init flightsql-sqlite-test

Connect to an unsecured Flight SQL service

Recipe source: flightsql_sqlite_connect.py

To connect to an unsecured Flight SQL service, just provide the URI.

import os

import adbc_driver_flightsql.dbapi

uri = os.environ["ADBC_SQLITE_FLIGHTSQL_URI"]
conn = adbc_driver_flightsql.dbapi.connect(uri)

We can then execute a simple query.

with conn.cursor() as cur:
    cur.execute("SELECT 1")

    assert cur.fetchone() == (1,)

conn.close()

Connect to a Flight SQL service with username and password

Recipe source: flightsql_dremio_connect.py

Dremio requires a username and password. To connect to a Flight SQL service with authentication, provide the options at connection time.

import os

import adbc_driver_flightsql.dbapi
import adbc_driver_manager

uri = os.environ["ADBC_DREMIO_FLIGHTSQL_URI"]
username = os.environ["ADBC_DREMIO_FLIGHTSQL_USER"]
password = os.environ["ADBC_DREMIO_FLIGHTSQL_PASS"]
conn = adbc_driver_flightsql.dbapi.connect(
    uri,
    db_kwargs={
        adbc_driver_manager.DatabaseOptions.USERNAME.value: username,
        adbc_driver_manager.DatabaseOptions.PASSWORD.value: password,
    },
)

We can then execute a simple query.

with conn.cursor() as cur:
    cur.execute("SELECT 1")

    assert cur.fetchone() == (1,)

conn.close()

Set timeouts and other options

Recipe source: flightsql_sqlite_options.py

The Flight SQL driver supports various options.

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import ConnectionOptions, DatabaseOptions

uri = os.environ["ADBC_SQLITE_FLIGHTSQL_URI"]

We can enable cookie support, which some server implementations require.

conn = adbc_driver_flightsql.dbapi.connect(
    uri,
    db_kwargs={DatabaseOptions.WITH_COOKIE_MIDDLEWARE.value: "true"},
)

Other options are set on the connection or statement.

For example, we can add a custom header to all outgoing requests.

custom_header = f"{ConnectionOptions.RPC_CALL_HEADER_PREFIX.value}x-custom-header"
conn.adbc_connection.set_options(**{custom_header: "value"})

We can also set timeouts. These are in floating-point seconds.

conn.adbc_connection.set_options(
    **{
        ConnectionOptions.TIMEOUT_FETCH.value: 30.0,
        ConnectionOptions.TIMEOUT_QUERY.value: 30.0,
        ConnectionOptions.TIMEOUT_UPDATE.value: 30.0,
    }
)

These options will apply to all cursors we create.

with conn.cursor() as cur:
    cur.execute("SELECT 1")

    assert cur.fetchone() == (1,)

conn.close()

Set the max gRPC message size

Recipe source: flightsql_sqlite_max_msg_size.py

By default, the Flight SQL driver limits the size of incoming/outgoing messages. You might see an error like this if those limits are exceeded:

INTERNAL: [FlightSQL] grpc: received message larger than max

These limits can be adjusted to avoid this.

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import DatabaseOptions

uri = os.environ["ADBC_SQLITE_FLIGHTSQL_URI"]

This query generates about 16 MiB per batch, which will trip the default limit.

query = """
WITH RECURSIVE generate_series(value) AS (
  SELECT 1
  UNION ALL
  SELECT value + 1 FROM generate_series
   WHERE value + 1 <= 2048
)
SELECT printf('%.*c', 16384, 'x') FROM generate_series
"""

When we execute the query, we’ll get an error.

conn = adbc_driver_flightsql.dbapi.connect(uri)
with conn.cursor() as cur:
    cur.execute(query)

    try:
        cur.fetchallarrow()
    except adbc_driver_flightsql.dbapi.InternalError:
        # This exception is expected.
        pass
    else:
        assert False, "Did not raise expected exception"

conn.close()

We can instead change the limit when connecting.

conn = adbc_driver_flightsql.dbapi.connect(
    uri,
    db_kwargs={DatabaseOptions.WITH_MAX_MSG_SIZE.value: "2147483647"},
)
with conn.cursor() as cur:
    cur.execute(query)

    assert len(cur.fetchallarrow()) == 2048

conn.close()