Flight SQL 示例

其中一些示例是针对由 SQLite 支持的演示 Flight SQL 服务编写的。 您可以按如下方式自己运行它

$ 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

其他示例使用 OSS 版本的 Dremio

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

如果您检出了 ADBC 存储库并安装了 Docker Compose，则可以使用我们的配置来运行这两项服务

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

连接到不安全的 Flight SQL 服务

示例源码: flightsql_sqlite_connect.py

要连接到不安全的 Flight SQL 服务，只需提供 URI。

import os

import adbc_driver_flightsql.dbapi

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

然后我们可以执行一个简单的查询。

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

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

conn.close()

使用用户名和密码连接到 Flight SQL 服务

示例源码: flightsql_dremio_connect.py

Dremio 需要用户名和密码。 要连接到需要身份验证的 Flight SQL 服务，请在连接时提供选项。

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,
    },
)

然后我们可以执行一个简单的查询。

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

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

conn.close()

设置超时和其他选项

示例源码: flightsql_sqlite_options.py

Flight SQL 驱动支持各种选项。

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import ConnectionOptions, DatabaseOptions

uri = os.environ["ADBC_SQLITE_FLIGHTSQL_URI"]

我们可以启用 cookie 支持，某些服务器实现需要这样做。

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

其他选项在连接或语句中设置。

例如，我们可以将自定义标头添加到所有传出请求。

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

我们还可以设置超时。 这些以浮点秒为单位。

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

这些选项将应用于我们创建的所有游标。

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

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

conn.close()

设置最大 gRPC 消息大小

示例源码: flightsql_sqlite_max_msg_size.py

默认情况下，Flight SQL 驱动会限制传入/传出消息的大小。 如果超过这些限制，您可能会看到如下错误

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

可以调整这些限制以避免这种情况。

import os

import adbc_driver_flightsql.dbapi
from adbc_driver_flightsql import DatabaseOptions

uri = os.environ["ADBC_SQLITE_FLIGHTSQL_URI"]

此查询每个批次生成大约 16 MiB，这将触发默认限制。

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
"""

当我们执行查询时，我们会收到一个错误。

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()

我们可以在连接时更改限制。

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()