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

其他方案使用开源版本的 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()

使用 OAuth 2.0 客户端凭据连接

方案来源: flightsql_oauth_client_credentials.py

Flight SQL 驱动程序支持 OAuth 2.0 身份验证。此示例展示了如何使用客户端凭据流（RFC 6749）进行连接，这适用于无需用户交互的机器对机器身份验证。

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

使用 OAuth 2.0 客户端凭据流连接。驱动程序将自动获取并刷新访问令牌。

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,

可选地，请求特定的作用域 (scopes)

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

为了使用自签名证书进行测试，请跳过 TLS 验证。在生产环境中，您应该提供适当的 TLS 证书。

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)

然后我们可以照常执行查询。

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

    result = cur.fetchone()
    print(result)

conn.close()

使用 OAuth 2.0 令牌交换连接

方案来源: flightsql_oauth_token_exchange.py

Flight SQL 驱动程序支持 OAuth 2.0 令牌交换 (RFC 8693)。这允许将现有令牌（例如，来自身份提供商的 JWT）交换为可用于访问 Flight SQL 服务的新令牌。

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

这通常是来自您身份提供商的 JWT 或其他令牌

subject_token = os.environ["ADBC_OAUTH_SUBJECT_TOKEN"]

为了使用自签名证书进行测试，请跳过 TLS 验证。在生产环境中，您应该提供适当的 TLS 证书。

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

使用 OAuth 2.0 令牌交换流连接。驱动程序会将主体令牌 (subject 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,

指定正在交换的主体令牌的类型

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

可选地，指定您想要接收的令牌类型

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

可选地，指定预期的受众 (audience)

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

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

然后我们可以照常执行查询。

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

    result = cur.fetchone()
    print(result)

conn.close()