Flight SQL Recipes

Some of these recipes are written against 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()

Connect with OAuth 2.0 Client Credentials

Recipe source: flightsql_oauth_client_credentials.py

The Flight SQL driver supports OAuth 2.0 authentication. This example shows how to connect using the Client Credentials flow (RFC 6749), which is suitable for machine-to-machine authentication without user interaction.

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import DatabaseOptions, OAuthFlowType

uri = os.environ["ADBC_TEST_FLIGHTSQL_URI"]
token_uri = os.environ["ADBC_OAUTH_TOKEN_URI"]
client_id = os.environ["ADBC_OAUTH_CLIENT_ID"]
client_secret = os.environ["ADBC_OAUTH_CLIENT_SECRET"]

Connect using OAuth 2.0 Client Credentials flow. The driver will automatically obtain and refresh access tokens.

db_kwargs = {
    DatabaseOptions.OAUTH_FLOW.value: OAuthFlowType.CLIENT_CREDENTIALS.value,
    DatabaseOptions.OAUTH_TOKEN_URI.value: token_uri,
    DatabaseOptions.OAUTH_CLIENT_ID.value: client_id,
    DatabaseOptions.OAUTH_CLIENT_SECRET.value: client_secret,

Optionally, request specific scopes

    # DatabaseOptions.OAUTH_SCOPE.value: "dremio.all",
}

For testing with self-signed certificates, skip TLS verification. In production, you should provide proper TLS certificates.

if os.environ.get("ADBC_OAUTH_SKIP_VERIFY", "true").lower() in ("1", "true"):
    db_kwargs[DatabaseOptions.TLS_SKIP_VERIFY.value] = "true"

conn = adbc_driver_flightsql.dbapi.connect(uri, db_kwargs=db_kwargs)

We can then execute queries as usual.

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

    result = cur.fetchone()
    print(result)

conn.close()

Connect with OAuth 2.0 Token Exchange

Recipe source: flightsql_oauth_token_exchange.py

The Flight SQL driver supports OAuth 2.0 Token Exchange (RFC 8693). This allows exchanging an existing token (e.g., a JWT from an identity provider) for a new token that can be used to access the Flight SQL service.

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import DatabaseOptions, OAuthFlowType, OAuthTokenType

uri = os.environ["ADBC_TEST_FLIGHTSQL_URI"]
token_uri = os.environ["ADBC_OAUTH_TOKEN_URI"]

This is typically a JWT or other token from your identity provider

subject_token = os.environ["ADBC_OAUTH_SUBJECT_TOKEN"]

For testing with self-signed certificates, skip TLS verification. In production, you should provide proper TLS certificates.

db_kwargs = {}
if os.environ.get("ADBC_OAUTH_SKIP_VERIFY", "true").lower() in ("1", "true"):
    db_kwargs[DatabaseOptions.TLS_SKIP_VERIFY.value] = "true"

Connect using OAuth 2.0 Token Exchange flow. The driver will exchange the subject token for an access token.

db_kwargs.update(
    {
        DatabaseOptions.OAUTH_FLOW.value: OAuthFlowType.TOKEN_EXCHANGE.value,
        DatabaseOptions.OAUTH_TOKEN_URI.value: token_uri,
        DatabaseOptions.OAUTH_EXCHANGE_SUBJECT_TOKEN.value: subject_token,

Specify the type of the subject token being exchanged

        DatabaseOptions.OAUTH_EXCHANGE_SUBJECT_TOKEN_TYPE.value: (
            OAuthTokenType.JWT.value
        ),

Optionally, specify the type of token you want to receive

        # DatabaseOptions.OAUTH_EXCHANGE_REQUESTED_TOKEN_TYPE.value:
        #   OAuthTokenType.ACCESS_TOKEN.value,

Optionally, specify the intended audience

        # DatabaseOptions.OAUTH_EXCHANGE_AUD.value: "my-service",
    }
)

conn = adbc_driver_flightsql.dbapi.connect(uri, db_kwargs=db_kwargs)

We can then execute queries as usual.

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

    result = cur.fetchone()
    print(result)

conn.close()