mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat: add a LanceDB command line (#4600)
Add a LanceDB command line with the following functionality: - `lancedb auth`: log in with OAuth, check cached OAuth session status, or log out. - `lancedb profile`: Create, edit, list, show, select, and delete connection profiles. - When no subcommand is given, we will attempt to create a SQL shell. Update ci/check_cargo_lock.py and its regression tests to allow pruning dependency edges while their target packages remain in Cargo.lock. Removing the CLI from the Python source distribution can disable optional features such as Axum tracing, even though other SDK dependencies still require the tracing package.
This commit is contained in:
Colin Patrick McCabe
@@ -0,0 +1,45 @@
|
||||
name: CLI
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
paths:
|
||||
- 'rust/lancedb/**'
|
||||
- 'Cargo.toml'
|
||||
- 'Cargo.lock'
|
||||
- '.github/workflows/cli.yml'
|
||||
permissions:
|
||||
contents: read
|
||||
jobs:
|
||||
test:
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-24.04, macos-14]
|
||||
runs-on: ${{ matrix.os }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
persist-credentials: false
|
||||
- uses: actions-rust-lang/setup-rust-toolchain@166cdcfd11aee3cb47222f9ddb555ce30ddb9659 # v1
|
||||
with:
|
||||
components: clippy
|
||||
- uses: Swatinem/rust-cache@d1527c8f67ef9af357290475d4af4b0fab8f12e0 # v2
|
||||
- name: Install Linux prerequisites
|
||||
if: runner.os == 'Linux'
|
||||
run: sudo apt-get update && sudo apt-get install -y protobuf-compiler libssl-dev clang
|
||||
- name: Install macOS prerequisites
|
||||
if: runner.os == 'macOS'
|
||||
run: brew install protobuf
|
||||
- name: Build standalone CLI binary
|
||||
run: cargo build --locked --profile ci -p lancedb-cli --bin lancedb
|
||||
- name: Test standalone CLI
|
||||
run: cargo test --locked --profile ci -p lancedb-cli
|
||||
- name: Test SDK SQL
|
||||
run: cargo test --locked --profile ci -p lancedb --features remote --lib remote::sql
|
||||
- name: Lint CLI
|
||||
run: cargo clippy --locked --profile ci -p lancedb-cli --all-targets --all-features -- -D warnings
|
||||
- name: Check SDK without remote support
|
||||
# Feature-disabled SDK code has existing unused-code warnings; this is a compile check.
|
||||
env:
|
||||
RUSTFLAGS: ""
|
||||
run: cargo check --locked --profile ci -p lancedb --no-default-features --tests
|
||||
@@ -187,12 +187,14 @@ jobs:
|
||||
command: sdist
|
||||
working-directory: python
|
||||
args: --out dist
|
||||
# maturin ships the nodejs-only pins in Cargo.lock. Remove this step when
|
||||
# the pinned maturin fixes https://github.com/PyO3/maturin/issues/2609.
|
||||
# maturin ships pins used only by other workspace members in Cargo.lock.
|
||||
# Remove this step when the pinned maturin fixes
|
||||
# https://github.com/PyO3/maturin/issues/2609.
|
||||
- name: Sync Cargo.lock to the packaged workspace
|
||||
shell: bash
|
||||
run: |
|
||||
# A repo lock in sync with every manifest means the prune only drops pins.
|
||||
# A synced repo lock lets the prune drop packages and optional edges
|
||||
# without changing the remaining pins or adding dependency edges.
|
||||
cargo metadata --locked --format-version 1 > /dev/null
|
||||
mkdir "$RUNNER_TEMP/sdist-src"
|
||||
tar -xzf python/dist/lancedb-*.tar.gz -C "$RUNNER_TEMP/sdist-src"
|
||||
|
||||
@@ -404,10 +404,18 @@ jobs:
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Downgrade dependencies
|
||||
# aws-smithy-types 1.6.1+ requires rustc 1.94.1
|
||||
run: cargo update -p aws-smithy-types --precise 1.5.0
|
||||
- name: Resolve MSRV-compatible dependencies
|
||||
# Resolve the AWS SDK and Smithy crates together: newer SDK releases
|
||||
# require newer Smithy versions, so downgrading only Smithy can fail.
|
||||
run: cargo +${{ matrix.msrv }} update --config 'resolver.incompatible-rust-versions="fallback"'
|
||||
- name: Keep Xet dependencies compatible with the MSRV
|
||||
# Xet 1.7 does not declare its MSRV, but requires sysinfo 0.39 (Rust 1.95).
|
||||
# Downgrade parents before their dependencies to satisfy their constraints.
|
||||
run: |
|
||||
for crate in hf-xet xet-data xet-client xet-core-structures xet-runtime; do
|
||||
cargo +${{ matrix.msrv }} update -p "$crate" --precise 1.6.0
|
||||
done
|
||||
- name: cargo +${{ matrix.msrv }} check
|
||||
env:
|
||||
RUSTUP_TOOLCHAIN: ${{ matrix.msrv }}
|
||||
run: cargo check --profile ci --workspace --tests --benches --all-features
|
||||
run: cargo check --locked --profile ci --workspace --tests --benches --all-features
|
||||
|
||||
Generated
+1516
-1157
File diff suppressed because it is too large
Load Diff
@@ -1,5 +1,6 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"rust/lancedb/cli",
|
||||
"rust/lancedb",
|
||||
"nodejs",
|
||||
"python",
|
||||
@@ -35,6 +36,7 @@ lance-encoding = { "version" = "=15.0.0-beta.3", "tag" = "v15.0.0-beta.3", "git"
|
||||
lance-arrow = { version = "=15.0.0-beta.3", path = "python/vendor/lance-arrow" }
|
||||
lancedb = { path = "rust/lancedb", default-features = false }
|
||||
ahash = "0.8"
|
||||
anyhow = "1"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
arrow-array = "58.0.0"
|
||||
@@ -49,6 +51,7 @@ arrow-flight = { version = "58.0.0", features = ["flight-sql-experimental"] }
|
||||
async-trait = "0"
|
||||
# Smithy JSON 0.63 requires the pre-1.7 Document representation; allow the MSRV pin.
|
||||
aws-smithy-types = ">=1.3.6, <1.7"
|
||||
base64 = "0.22"
|
||||
bytes = "1"
|
||||
datafusion = { version = "54.0.0", default-features = false }
|
||||
datafusion-catalog = "54.0.0"
|
||||
@@ -64,12 +67,15 @@ half = { "version" = "2.7.1", default-features = false, features = [
|
||||
"num-traits",
|
||||
] }
|
||||
futures = "0.3"
|
||||
hf-hub = { version = "0.4.1", default-features = false }
|
||||
http = "1"
|
||||
log = "0.4"
|
||||
metrics = "0.24"
|
||||
metrics-util = "0.19"
|
||||
moka = { version = "0.12", features = ["future"] }
|
||||
object_store = "0.14.1"
|
||||
pin-project = "1.0.7"
|
||||
prost = "0.14"
|
||||
rand = "0.9"
|
||||
snafu = "0.8"
|
||||
url = "2"
|
||||
@@ -79,6 +85,8 @@ regex = "1.10"
|
||||
semver = "1.0.25"
|
||||
serde = "1"
|
||||
serde_json = { version = "1", features = ["preserve_order"] }
|
||||
serial_test = "3"
|
||||
sha2 = "0.10"
|
||||
tempfile = "3.5.0"
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] }
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
"""Fail if pruning Cargo.lock did more than remove packages and the edges to them.
|
||||
"""Fail if pruning Cargo.lock changes pins or adds dependency edges.
|
||||
|
||||
Removing workspace members can disable optional dependency edges even when
|
||||
their target packages remain in use elsewhere in the packaged workspace.
|
||||
|
||||
Usage: python ci/check_cargo_lock.py <lock-before> <lock-after>
|
||||
"""
|
||||
@@ -45,7 +48,7 @@ def main() -> None:
|
||||
for package, deps in edges_after.items():
|
||||
if package not in edges_before:
|
||||
errors.append(f"added or changed {package}")
|
||||
elif deps != edges_before[package] - removed:
|
||||
elif not deps <= edges_before[package]:
|
||||
errors.append(f"dependencies of {package[:2]} changed")
|
||||
if errors:
|
||||
sys.exit("error: the prune changed Cargo.lock:\n" + "\n".join(errors))
|
||||
|
||||
@@ -37,14 +37,15 @@ APP = package("app", "1.0.0", ["libc", "untrusted"], source=None)
|
||||
|
||||
|
||||
class CheckCargoLockTest(unittest.TestCase):
|
||||
def run_check(self, after, header="version = 4\n\n"):
|
||||
def run_check(self, after, header="version = 4\n\n", before=BEFORE):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
before_path = Path(tmp, "before.lock")
|
||||
after_path = Path(tmp, "after.lock")
|
||||
before_path.write_text(BEFORE)
|
||||
before_path.write_text(before)
|
||||
after_path.write_text(header + after)
|
||||
return subprocess.run(
|
||||
[sys.executable, SCRIPT, before_path, after_path],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
@@ -81,9 +82,36 @@ class CheckCargoLockTest(unittest.TestCase):
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn("dependencies of ('libc'", result.stderr)
|
||||
|
||||
def test_removed_dependency_fails(self):
|
||||
def test_removed_dependency_passes(self):
|
||||
app = package("app", "1.0.0", ["untrusted"], source=None)
|
||||
result = self.run_check(app + LIBC + UNTRUSTED)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
|
||||
def test_disabled_optional_dependency_still_used_elsewhere_passes(self):
|
||||
# Removing the CLI disables axum's tracing feature, but other SDK
|
||||
# dependencies still need tracing. Its package therefore stays pinned.
|
||||
app = package("app", "1.0.0", ["axum", "tracing"], source=None)
|
||||
tracing = package("tracing", "0.1.44")
|
||||
before = (
|
||||
"version = 4\n\n" + app + package("axum", "0.8.9", ["tracing"]) + tracing
|
||||
)
|
||||
after = app + package("axum", "0.8.9") + tracing
|
||||
result = self.run_check(after, before=before)
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
|
||||
def test_dependency_redirected_to_an_existing_version_fails(self):
|
||||
before = (
|
||||
"version = 4\n\n"
|
||||
+ package("app", "1.0.0", ["untrusted 0.7.1"], source=None)
|
||||
+ package("untrusted", "0.7.1")
|
||||
+ UNTRUSTED
|
||||
)
|
||||
after = (
|
||||
package("app", "1.0.0", ["untrusted 0.9.0"], source=None)
|
||||
+ package("untrusted", "0.7.1")
|
||||
+ UNTRUSTED
|
||||
)
|
||||
result = self.run_check(after, before=before)
|
||||
self.assertNotEqual(result.returncode, 0)
|
||||
self.assertIn("dependencies of ('app'", result.stderr)
|
||||
|
||||
|
||||
@@ -28,6 +28,9 @@ endpoint, returning results as Arrow record batches:
|
||||
any tool with its own Flight SQL integration. A LanceDB-specific SDK is not
|
||||
required to submit SQL.
|
||||
|
||||
From the command line, use `lancedb` to open the [SQL shell](/tools/cli)
|
||||
or execute a statement or file.
|
||||
|
||||
The [SQL reference](/sql#connecting) shows a query from Python. See
|
||||
[SDKs](/basics/sdk) for the LanceDB language clients.
|
||||
|
||||
|
||||
@@ -288,6 +288,12 @@
|
||||
"deployment/support"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "Tools",
|
||||
"pages": [
|
||||
"tools/cli"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "Concepts",
|
||||
"pages": [
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
---
|
||||
title: "LanceDB CLI"
|
||||
sidebarTitle: "Command line"
|
||||
description: "Connect to LanceDB, run SQL, and manage connection profiles with the lancedb command-line client."
|
||||
icon: terminal
|
||||
keywords: ["cli", "command line", "sql", "profiles", "authentication"]
|
||||
---
|
||||
|
||||
Use the `lancedb` command-line client to open an interactive SQL shell, run a
|
||||
statement, or execute a SQL file. SQL requires a running deployment with an
|
||||
[Arrow Flight SQL endpoint](/basics/sql#sql-is-remote).
|
||||
|
||||
## Install
|
||||
|
||||
With Rust and Git installed, build the CLI from the LanceDB repository:
|
||||
|
||||
```sh
|
||||
git clone https://github.com/lancedb/lancedb.git
|
||||
cd lancedb
|
||||
cargo install --locked --path rust/lancedb/cli
|
||||
```
|
||||
|
||||
Check the available options:
|
||||
|
||||
```sh
|
||||
lancedb --help
|
||||
```
|
||||
|
||||
## Connect to a deployment
|
||||
|
||||
Get the Flight SQL endpoint, HTTP API endpoint, database name, and credentials
|
||||
from your deployment administrator. Create a named profile to save the
|
||||
connection settings:
|
||||
|
||||
```sh
|
||||
lancedb profile create production \
|
||||
--sql-server-url grpcs://sql.example.com:10025 \
|
||||
--http-server-url https://query.example.com \
|
||||
--database analytics
|
||||
```
|
||||
|
||||
The SQL endpoint executes statements. The HTTP endpoint supports shell commands
|
||||
that inspect table and index metadata. Replace both example URLs with the
|
||||
endpoints for your deployment; they may use different hosts or ports.
|
||||
|
||||
For API-key authentication, set the credential in your environment:
|
||||
|
||||
```sh
|
||||
export LANCEDB_API_KEY="YOUR_API_KEY"
|
||||
lancedb --profile production
|
||||
```
|
||||
|
||||
Use `grpcs://` or `https://` for a TLS-protected SQL connection. A bare
|
||||
`host:port` or a `grpc://` address selects a plaintext connection.
|
||||
|
||||
### Set the default profile
|
||||
|
||||
```sh
|
||||
lancedb profile set production
|
||||
lancedb
|
||||
```
|
||||
|
||||
<Note>
|
||||
Running `lancedb` opens the SQL shell when the default profile has a Flight SQL
|
||||
endpoint. If no endpoint is configured, it prints usage and connection setup
|
||||
guidance. Command and file execution also require a configured endpoint.
|
||||
</Note>
|
||||
|
||||
To connect to an existing server on localhost, explicitly select the built-in
|
||||
local profile:
|
||||
|
||||
```sh
|
||||
lancedb --profile local
|
||||
```
|
||||
|
||||
This profile uses Flight SQL on port `10025` and HTTP on port `10024`.
|
||||
|
||||
## Run SQL
|
||||
|
||||
Execute a statement and exit:
|
||||
|
||||
```sh
|
||||
lancedb --profile production --command 'SELECT 1' --format json
|
||||
```
|
||||
|
||||
To execute a file, separate its statements with semicolons. For example, save
|
||||
the following as `queries.sql`:
|
||||
|
||||
```sql
|
||||
SELECT COUNT(*) FROM documents;
|
||||
SELECT id, text FROM documents LIMIT 10;
|
||||
```
|
||||
|
||||
Run the file with:
|
||||
|
||||
```sh
|
||||
lancedb --profile production --file queries.sql
|
||||
```
|
||||
|
||||
Quoted strings and comments in SQL files are preserved when statements are
|
||||
split. `--command` and `--file` are mutually exclusive.
|
||||
|
||||
### Choose an output format
|
||||
|
||||
For query results, `--format` accepts the following values:
|
||||
|
||||
| Format | Output |
|
||||
| --- | --- |
|
||||
| `table` | A display table, the default |
|
||||
| `csv` | Comma-separated values with column headers |
|
||||
| `json` | A JSON array of rows |
|
||||
| `ndjson` | One JSON row per line |
|
||||
|
||||
```sh
|
||||
lancedb --profile production --command 'SELECT id FROM documents' --format csv
|
||||
```
|
||||
|
||||
## Use the SQL shell
|
||||
|
||||
End SQL statements with a semicolon. Enter `\?` to see the shell's commands.
|
||||
|
||||
| Command | Purpose |
|
||||
| --- | --- |
|
||||
| `\dt` | List tables |
|
||||
| `\d documents` | Show a table's columns and indexes |
|
||||
| `\di` | List indexes |
|
||||
| `\dj` | List jobs |
|
||||
| `\c analytics` | Connect to another database |
|
||||
| `\e` | Edit the query buffer in an external editor |
|
||||
| `\watch 2` | Repeat the last query or listing every two seconds |
|
||||
| `\q` | Exit the shell |
|
||||
|
||||
Add `+` to `\dt`, `\d`, or `\di` for more detail. Use Ctrl+C to cancel the
|
||||
current input or stop watching a query, and Ctrl+D to exit.
|
||||
|
||||
Large interactive results use a pager. Set `NO_PAGER=1` to disable paging.
|
||||
|
||||
### Inspect jobs with SQL
|
||||
|
||||
```sh
|
||||
lancedb --profile production --command 'SHOW JOBS'
|
||||
lancedb --profile production --command "SHOW ERRORS ON documents FOR JOB 'JOB_ID'"
|
||||
lancedb --profile production --command "CANCEL JOB 'JOB_ID'"
|
||||
```
|
||||
|
||||
See the [job SQL reference](/sql/ddl/job) for syntax and deployment requirements,
|
||||
and the [SQL reference](/sql) for other supported statements.
|
||||
|
||||
## Manage profiles
|
||||
|
||||
Profiles are stored in `~/.lancedb/config.toml`. The CLI reuses existing
|
||||
Enterprise profiles and cached OAuth sessions.
|
||||
|
||||
```sh
|
||||
lancedb profile list
|
||||
lancedb profile show production
|
||||
lancedb profile edit production --database staging
|
||||
lancedb profile delete production
|
||||
```
|
||||
|
||||
Use `lancedb profile create NAME` without connection flags for interactive
|
||||
setup. For a single SQL invocation, flags such as `--database` and
|
||||
`--sql-server-url` override the selected profile.
|
||||
|
||||
Put each command's flags after that command. For example, use
|
||||
`lancedb auth login --profile production` or `lancedb profile edit production
|
||||
--database staging`.
|
||||
|
||||
## Authenticate with OAuth
|
||||
|
||||
To use OAuth, configure an `oidc` table in the selected profile with your
|
||||
identity provider's settings. For example, add the following to
|
||||
`~/.lancedb/config.toml` for a device-code login:
|
||||
|
||||
```toml
|
||||
[profiles.production.oidc]
|
||||
flow = "device_code"
|
||||
issuer_url = "https://identity.example.com"
|
||||
client_id = "YOUR_CLIENT_ID"
|
||||
scopes = ["openid", "offline_access"]
|
||||
```
|
||||
|
||||
Choose one authentication method per profile. An OAuth profile cannot also
|
||||
contain `api_key` or `oidc_access_token_file`. Environment credentials override
|
||||
the profile, so clear API-key overrides before connecting with OAuth:
|
||||
|
||||
```sh
|
||||
unset LANCEDB_API_KEY
|
||||
lancedb auth login --profile production
|
||||
lancedb auth status --profile production
|
||||
lancedb --profile production
|
||||
```
|
||||
|
||||
`auth login` follows the profile's OAuth flow. Browser and device-code sessions
|
||||
cache refresh tokens for later commands. Use `lancedb auth logout --profile
|
||||
production` to clear the profile's cached session.
|
||||
|
||||
See [authentication](/enterprise/authentication) for identity provider setup and
|
||||
available authentication methods.
|
||||
|
||||
## Configure TLS
|
||||
|
||||
To trust a private certificate authority, pass its PEM certificate:
|
||||
|
||||
```sh
|
||||
lancedb --profile production --ca-cert /path/to/ca.pem
|
||||
```
|
||||
|
||||
If the server requires mutual TLS, supply both a client certificate and its
|
||||
PKCS#8 private key. Save these settings in the profile with:
|
||||
|
||||
```sh
|
||||
lancedb profile edit production \
|
||||
--ca-cert /path/to/ca.pem \
|
||||
--client-cert /path/to/client.pem \
|
||||
--client-key /path/to/client-key.pem
|
||||
```
|
||||
@@ -8,6 +8,7 @@ repository.workspace = true
|
||||
keywords.workspace = true
|
||||
categories.workspace = true
|
||||
rust-version.workspace = true
|
||||
exclude = ["cli/**"]
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
[dependencies]
|
||||
@@ -22,7 +23,7 @@ arrow-ord = { workspace = true }
|
||||
arrow-cast = { workspace = true }
|
||||
arrow-ipc.workspace = true
|
||||
arrow-flight = { workspace = true, optional = true }
|
||||
prost = { version = "0.14", optional = true }
|
||||
prost = { workspace = true, optional = true }
|
||||
chrono = { workspace = true }
|
||||
datafusion-catalog.workspace = true
|
||||
datafusion-common.workspace = true
|
||||
@@ -86,16 +87,16 @@ reqwest = { version = "0.13.0", default-features = false, features = [
|
||||
"system-proxy",
|
||||
], optional = true }
|
||||
tonic = { workspace = true, optional = true }
|
||||
http = { version = "1", optional = true } # Matching what is in reqwest
|
||||
http = { workspace = true, optional = true } # Matching what is in reqwest
|
||||
oauth2 = { workspace = true, optional = true }
|
||||
urlencoding = { version = "2", optional = true }
|
||||
base64 = { version = "0.22", optional = true }
|
||||
base64 = { workspace = true, optional = true }
|
||||
fs4 = { version = "0.13", optional = true }
|
||||
webbrowser = { version = "1", optional = true }
|
||||
uuid = { workspace = true, features = ["v5"] }
|
||||
polars-arrow = { version = ">=0.37,<0.40.0", optional = true }
|
||||
polars = { version = ">=0.37,<0.40.0", optional = true }
|
||||
hf-hub = { version = "0.4.1", optional = true, default-features = false, features = [
|
||||
hf-hub = { workspace = true, optional = true, features = [
|
||||
"rustls-tls",
|
||||
"tokio",
|
||||
"ureq",
|
||||
@@ -106,10 +107,10 @@ candle-nn = { version = "0.9.1", optional = true }
|
||||
tokenizers = { version = "0.19.1", optional = true }
|
||||
semver = { workspace = true }
|
||||
roaring = "0.11.4"
|
||||
sha2 = "0.10"
|
||||
sha2.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
anyhow = "1"
|
||||
anyhow.workspace = true
|
||||
lance-testing = { workspace = true }
|
||||
# RestAdapter is only used by cfg(test) code. Keep rest-adapter off the
|
||||
# production remote feature so Python wheels do not resolve its Axum 0.7 stack.
|
||||
@@ -130,7 +131,7 @@ datafusion.workspace = true
|
||||
http-body = "1" # Matching reqwest
|
||||
rstest = "0.23.0"
|
||||
test-log = "0.2"
|
||||
serial_test = "3"
|
||||
serial_test.workspace = true
|
||||
[target.'cfg(unix)'.dev-dependencies]
|
||||
pprof = { version = "0.14", features = ["flamegraph"] }
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
[package]
|
||||
name = "lancedb-cli"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
publish = false
|
||||
|
||||
[[bin]]
|
||||
name = "lancedb"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
tempfile.workspace = true
|
||||
which = "6"
|
||||
anyhow.workspace = true
|
||||
arrow-array.workspace = true
|
||||
arrow-buffer.workspace = true
|
||||
arrow-cast.workspace = true
|
||||
arrow-json = "58.0.0"
|
||||
arrow-schema.workspace = true
|
||||
async-trait.workspace = true
|
||||
clap = { version = "4.5.13", features = ["derive", "env"] }
|
||||
directories = "5"
|
||||
futures.workspace = true
|
||||
http.workspace = true
|
||||
lancedb = { workspace = true, features = ["remote"] }
|
||||
rustyline = "14"
|
||||
rustls = { version = "0.23", features = ["ring"] }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_ignored = "0.1.14"
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
sqlparser = { version = "0.62", default-features = false, features = ["std"] }
|
||||
termsize = "0.1"
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
toml = "0.8"
|
||||
tracing.workspace = true
|
||||
unicode-width = "0.2"
|
||||
|
||||
[dev-dependencies]
|
||||
arrow-flight.workspace = true
|
||||
prost.workspace = true
|
||||
rcgen = "0.12"
|
||||
tokio-stream = { version = "0.1", features = ["net"] }
|
||||
half.workspace = true
|
||||
axum = "0.8"
|
||||
base64.workspace = true
|
||||
csv = "1"
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
tonic.workspace = true
|
||||
url.workspace = true
|
||||
@@ -0,0 +1,419 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! CLI configuration and transport adapters for LanceDB's OAuth implementation.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use anyhow::{anyhow, ensure, Context, Result};
|
||||
use clap::{Args, Subcommand};
|
||||
use lancedb::remote::{
|
||||
AuthorizationCodeOptions, ClientAuthMethod, HeaderProvider, OAuthConfig, OAuthFlow,
|
||||
OAuthHeaderProvider, OAuthSession, TokenCacheOptions,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
use crate::server_config::{Authentication, Config, DEFAULT_PROFILE};
|
||||
|
||||
const CLIENT_SECRET_ENV: &str = "LANCEDB_OIDC_CLIENT_SECRET";
|
||||
pub const OIDC_CREDENTIAL_HEADER: &str = "x-lancedb-credential-type";
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum OidcFlow {
|
||||
ClientCredentials,
|
||||
AuthorizationCode,
|
||||
DeviceCode,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for OidcFlow {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(match self {
|
||||
Self::ClientCredentials => "client_credentials",
|
||||
Self::AuthorizationCode => "authorization_code",
|
||||
Self::DeviceCode => "device_code",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum OidcClientAuthMethod {
|
||||
None,
|
||||
ClientSecretBasic,
|
||||
ClientSecretPost,
|
||||
}
|
||||
|
||||
impl From<OidcClientAuthMethod> for ClientAuthMethod {
|
||||
fn from(method: OidcClientAuthMethod) -> Self {
|
||||
match method {
|
||||
OidcClientAuthMethod::None => Self::None,
|
||||
OidcClientAuthMethod::ClientSecretBasic => Self::ClientSecretBasic,
|
||||
OidcClientAuthMethod::ClientSecretPost => Self::ClientSecretPost,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Nonsecret configuration stored inside a connection profile.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct OidcConfig {
|
||||
pub flow: OidcFlow,
|
||||
pub issuer_url: String,
|
||||
pub client_id: String,
|
||||
pub scopes: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_secret_file: Option<PathBuf>,
|
||||
/// When omitted, the SDK selects authentication based on the client secret.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_auth_method: Option<OidcClientAuthMethod>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub resource: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub audience: Option<String>,
|
||||
/// Registered loopback callback URI for authorization-code login.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub redirect_uri: Option<String>,
|
||||
/// Defaults to the SDK's callback port (8400).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub callback_port: Option<u16>,
|
||||
/// Defaults to the SDK's refresh buffer (300 seconds).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub refresh_buffer_secs: Option<u64>,
|
||||
}
|
||||
|
||||
/// One SDK provider is shared by all transports using the resolved profile.
|
||||
#[derive(Debug)]
|
||||
pub struct ManagedOidc {
|
||||
config: OidcConfig,
|
||||
cache_dir: PathBuf,
|
||||
provider: OnceCell<OAuthHeaderProvider>,
|
||||
}
|
||||
|
||||
impl PartialEq for ManagedOidc {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
// Acquired tokens do not change which authentication configuration was selected.
|
||||
self.config == other.config && self.cache_dir == other.cache_dir
|
||||
}
|
||||
}
|
||||
impl Eq for ManagedOidc {}
|
||||
|
||||
/// A resolved credential, with no provider state attached.
|
||||
pub enum Credential {
|
||||
ApiKey(String),
|
||||
OidcAccessToken(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Credential {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(match self {
|
||||
Self::ApiKey(_) => "ApiKey(<redacted>)",
|
||||
Self::OidcAccessToken(_) => "OidcAccessToken(<redacted>)",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Credential {
|
||||
pub fn headers(&self) -> Result<http::HeaderMap> {
|
||||
let mut headers = http::HeaderMap::new();
|
||||
let (name, value) = match self {
|
||||
Self::ApiKey(key) => ("x-api-key", key.clone()),
|
||||
Self::OidcAccessToken(token) => {
|
||||
headers.insert(
|
||||
OIDC_CREDENTIAL_HEADER,
|
||||
http::HeaderValue::from_static("oidc"),
|
||||
);
|
||||
("authorization", format!("Bearer {token}"))
|
||||
}
|
||||
};
|
||||
let mut value = http::HeaderValue::from_str(&value)
|
||||
.map_err(|_| anyhow!("Invalid authentication credential: not a valid HTTP header"))?;
|
||||
value
|
||||
.to_str()
|
||||
.map_err(|_| anyhow!("Invalid authentication credential: expected an ASCII header"))?;
|
||||
value.set_sensitive(true);
|
||||
headers.insert(name, value);
|
||||
Ok(headers)
|
||||
}
|
||||
}
|
||||
|
||||
impl Authentication {
|
||||
pub async fn headers(&self) -> Result<http::HeaderMap> {
|
||||
match self {
|
||||
Self::ApiKey(key) => Credential::ApiKey(key.clone()).headers(),
|
||||
Self::OidcAccessToken(token) => Credential::OidcAccessToken(token.clone()).headers(),
|
||||
Self::OidcAccessTokenFile(path) => {
|
||||
let token = tokio::fs::read_to_string(path).await.with_context(|| {
|
||||
format!("Failed to read OIDC access token file {}", path.display())
|
||||
})?;
|
||||
let token = token.trim();
|
||||
ensure!(!token.is_empty(), "OIDC access token file is empty");
|
||||
Credential::OidcAccessToken(token.to_owned()).headers()
|
||||
}
|
||||
Self::ManagedOidc(session) => session.headers().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl HeaderProvider for Authentication {
|
||||
async fn get_headers(&self) -> lancedb::Result<HashMap<String, String>> {
|
||||
let headers = self
|
||||
.headers()
|
||||
.await
|
||||
.map_err(|error| lancedb::Error::Runtime {
|
||||
message: format!("Authentication failed: {error:#}"),
|
||||
})?;
|
||||
Ok(headers
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
(
|
||||
key.as_str().to_owned(),
|
||||
value
|
||||
.to_str()
|
||||
.expect("validated credential header")
|
||||
.to_owned(),
|
||||
)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl ManagedOidc {
|
||||
pub fn for_profile(profile: &str, config: OidcConfig) -> Result<Self> {
|
||||
let root = Config::default_config_path()?;
|
||||
// Profile names are arbitrary strings; hashing prevents paths from escaping the cache.
|
||||
let profile_key = format!("{:x}", Sha256::digest(profile.as_bytes()));
|
||||
Self::new(
|
||||
config,
|
||||
root.parent().unwrap().join("auth").join(profile_key),
|
||||
)
|
||||
}
|
||||
|
||||
fn new(config: OidcConfig, cache_dir: PathBuf) -> Result<Self> {
|
||||
ensure!(
|
||||
!config.client_id.trim().is_empty(),
|
||||
"OIDC client_id cannot be empty"
|
||||
);
|
||||
ensure!(
|
||||
!config.scopes.is_empty(),
|
||||
"OIDC scopes must contain at least one scope"
|
||||
);
|
||||
ensure!(
|
||||
config.flow == OidcFlow::AuthorizationCode
|
||||
|| (config.redirect_uri.is_none() && config.callback_port.is_none()),
|
||||
"redirect_uri and callback_port require flow = authorization_code"
|
||||
);
|
||||
Ok(Self {
|
||||
config,
|
||||
cache_dir,
|
||||
provider: OnceCell::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn flow(&self) -> &OidcFlow {
|
||||
&self.config.flow
|
||||
}
|
||||
|
||||
async fn sdk_config(&self) -> Result<OAuthConfig> {
|
||||
let flow = match self.config.flow {
|
||||
OidcFlow::ClientCredentials => OAuthFlow::ClientCredentials,
|
||||
OidcFlow::DeviceCode => OAuthFlow::DeviceCode,
|
||||
OidcFlow::AuthorizationCode => {
|
||||
let mut options = AuthorizationCodeOptions::new();
|
||||
if let Some(uri) = &self.config.redirect_uri {
|
||||
options = options.redirect_uri(uri);
|
||||
}
|
||||
if let Some(port) = self.config.callback_port {
|
||||
options = options.callback_port(port);
|
||||
}
|
||||
OAuthFlow::AuthorizationCode(options)
|
||||
}
|
||||
};
|
||||
Ok(OAuthConfig {
|
||||
issuer_url: self.config.issuer_url.clone(),
|
||||
client_id: self.config.client_id.clone(),
|
||||
client_secret: self.client_secret().await?,
|
||||
scopes: self.config.scopes.clone(),
|
||||
resource: self.config.resource.clone(),
|
||||
audience: self.config.audience.clone(),
|
||||
flow,
|
||||
client_auth_method: self.config.client_auth_method.map(Into::into),
|
||||
refresh_buffer_secs: self.config.refresh_buffer_secs,
|
||||
token_cache: (self.config.flow != OidcFlow::ClientCredentials)
|
||||
.then(|| TokenCacheOptions::new().cache_dir(&self.cache_dir)),
|
||||
})
|
||||
}
|
||||
|
||||
async fn provider(&self) -> Result<&OAuthHeaderProvider> {
|
||||
self.provider
|
||||
.get_or_try_init(|| async {
|
||||
OAuthHeaderProvider::new(self.sdk_config().await?).map_err(sdk_error)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn session(&self) -> Result<OAuthSession> {
|
||||
OAuthSession::new(self.sdk_config().await?).map_err(sdk_error)
|
||||
}
|
||||
|
||||
async fn headers(&self) -> Result<http::HeaderMap> {
|
||||
let values = self
|
||||
.provider()
|
||||
.await?
|
||||
.get_headers()
|
||||
.await
|
||||
.map_err(sdk_error)?;
|
||||
let mut headers = http::HeaderMap::new();
|
||||
for (name, value) in values {
|
||||
let name = http::HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|_| anyhow!("Invalid authentication header name"))?;
|
||||
let mut value = http::HeaderValue::from_str(&value).map_err(|_| {
|
||||
anyhow!("Invalid authentication credential: not a valid HTTP header")
|
||||
})?;
|
||||
value.to_str().map_err(|_| {
|
||||
anyhow!("Invalid authentication credential: expected an ASCII header")
|
||||
})?;
|
||||
if name == http::header::AUTHORIZATION || name == "x-api-key" {
|
||||
value.set_sensitive(true);
|
||||
}
|
||||
headers.insert(name, value);
|
||||
}
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
async fn client_secret(&self) -> Result<Option<String>> {
|
||||
if let Some(secret) = std::env::var_os(CLIENT_SECRET_ENV) {
|
||||
let secret = secret
|
||||
.into_string()
|
||||
.map_err(|_| anyhow!("{CLIENT_SECRET_ENV} must be UTF-8"))?;
|
||||
ensure!(!secret.is_empty(), "{CLIENT_SECRET_ENV} cannot be empty");
|
||||
return Ok(Some(secret));
|
||||
}
|
||||
let Some(path) = self.config.client_secret_file.as_ref() else {
|
||||
ensure!(
|
||||
self.config.flow != OidcFlow::ClientCredentials,
|
||||
"Client credentials require LANCEDB_OIDC_CLIENT_SECRET or oidc.client_secret_file"
|
||||
);
|
||||
return Ok(None);
|
||||
};
|
||||
let secret = tokio::fs::read_to_string(path)
|
||||
.await
|
||||
.context("Failed to read OIDC client secret file")?;
|
||||
let secret = secret.trim();
|
||||
ensure!(!secret.is_empty(), "OIDC client secret file is empty");
|
||||
Ok(Some(secret.to_owned()))
|
||||
}
|
||||
}
|
||||
|
||||
fn sdk_error(error: lancedb::Error) -> anyhow::Error {
|
||||
// SDK runtime errors can contain provider response bodies and credentials.
|
||||
// Configuration errors are local; never attach the original runtime error as a cause.
|
||||
match error {
|
||||
lancedb::Error::InvalidInput { message } => anyhow!("Invalid OIDC configuration: {message}"),
|
||||
_ => anyhow!("OIDC authentication failed; check the issuer, client credentials, scopes, and provider availability. For a cached interactive session, try lancedb auth logout followed by lancedb auth login."),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Args, Debug)]
|
||||
pub struct AuthArgs {
|
||||
#[command(subcommand)]
|
||||
pub command: AuthCommand,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum AuthCommand {
|
||||
/// Log in using the profile's OAuth flow; browser/device sessions cache refresh tokens.
|
||||
Login(ProfileArgs),
|
||||
/// Report whether an interactive session is cached, without contacting the provider.
|
||||
Status(ProfileArgs),
|
||||
/// Clear the profile's local interactive session; does not revoke provider credentials.
|
||||
Logout(ProfileArgs),
|
||||
}
|
||||
|
||||
#[derive(Args, Debug)]
|
||||
pub struct ProfileArgs {
|
||||
#[arg(short = 'p', long, value_name = "PROFILE")]
|
||||
pub profile: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn run(args: AuthArgs, parent_profile: Option<String>) -> Result<()> {
|
||||
let profile_args = match &args.command {
|
||||
AuthCommand::Login(args) | AuthCommand::Status(args) | AuthCommand::Logout(args) => args,
|
||||
};
|
||||
let config = Config::load_default()?;
|
||||
let profile_name = profile_args
|
||||
.profile
|
||||
.clone()
|
||||
.or(parent_profile)
|
||||
.or(config.default_profile.clone())
|
||||
.unwrap_or_else(|| DEFAULT_PROFILE.to_owned());
|
||||
let profile = config
|
||||
.profiles
|
||||
.get(&profile_name)
|
||||
.context("Profile does not exist")?;
|
||||
profile.validate_authentication()?;
|
||||
let oidc = profile
|
||||
.oidc
|
||||
.clone()
|
||||
.context("Profile has no managed OIDC configuration")?;
|
||||
let managed = ManagedOidc::for_profile(&profile_name, oidc)?;
|
||||
match args.command {
|
||||
AuthCommand::Login(_) => {
|
||||
if managed.config.flow == OidcFlow::ClientCredentials {
|
||||
managed.headers().await?;
|
||||
println!("Authenticated service account for profile '{profile_name}'. Access tokens are cached only within each process.");
|
||||
} else {
|
||||
let status = managed.session().await?.login().await.map_err(sdk_error)?;
|
||||
println!(
|
||||
"Authenticated profile '{profile_name}' using {}",
|
||||
managed.flow()
|
||||
);
|
||||
if !status.refreshable {
|
||||
eprintln!("The provider returned no refresh token; later commands will need another login. Check the provider's offline-access scope and consent settings.");
|
||||
}
|
||||
}
|
||||
if [
|
||||
"LANCEDB_API_KEY",
|
||||
"LANCEDB_OIDC_ACCESS_TOKEN",
|
||||
"LANCEDB_OIDC_ACCESS_TOKEN_FILE",
|
||||
]
|
||||
.iter()
|
||||
.any(|key| std::env::var_os(key).is_some())
|
||||
{
|
||||
eprintln!("Authentication environment overrides are set; ordinary commands will use those credentials instead of this session.");
|
||||
}
|
||||
}
|
||||
AuthCommand::Status(_) => {
|
||||
if managed.config.flow == OidcFlow::ClientCredentials {
|
||||
println!("Profile '{profile_name}' uses client credentials; no persistent session is stored.");
|
||||
} else {
|
||||
let status = managed.session().await?.status().await.map_err(sdk_error)?;
|
||||
println!(
|
||||
"Profile '{profile_name}' uses {}; cached refresh token: {}",
|
||||
managed.flow(),
|
||||
if status.refreshable {
|
||||
"present"
|
||||
} else {
|
||||
"not present"
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
AuthCommand::Logout(_) => {
|
||||
if managed.config.flow == OidcFlow::ClientCredentials {
|
||||
println!("Profile '{profile_name}' uses client credentials; no persistent session to clear. Configured credentials remain usable.");
|
||||
} else {
|
||||
managed.session().await?.logout().await.map_err(sdk_error)?;
|
||||
println!("Cleared the local session for profile '{profile_name}'. Provider credentials remain valid.");
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,417 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use super::*;
|
||||
use axum::{
|
||||
extract::{Form, Query},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::Redirect,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use base64::Engine;
|
||||
use std::fs;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
|
||||
struct Provider {
|
||||
config: OidcConfig,
|
||||
requests: Arc<Mutex<Vec<HashMap<String, String>>>>,
|
||||
server: tokio::task::JoinHandle<()>,
|
||||
dir: TempDir,
|
||||
}
|
||||
|
||||
impl Drop for Provider {
|
||||
fn drop(&mut self) {
|
||||
self.server.abort();
|
||||
}
|
||||
}
|
||||
|
||||
impl Provider {
|
||||
async fn start(flow: OidcFlow, lifetime: u64) -> Self {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let issuer = format!("http://{}", listener.local_addr().unwrap());
|
||||
let secret_file = dir.path().join("secret");
|
||||
fs::write(&secret_file, "test-secret\n").unwrap();
|
||||
let config = OidcConfig {
|
||||
client_auth_method: None,
|
||||
client_secret_file: (flow == OidcFlow::ClientCredentials).then_some(secret_file),
|
||||
flow,
|
||||
issuer_url: issuer.clone(),
|
||||
client_id: "cli-client".into(),
|
||||
scopes: vec!["openid".into(), "offline_access".into()],
|
||||
resource: Some("https://api.example.com/a?x=1&y=two".into()),
|
||||
audience: Some("lancedb + audience".into()),
|
||||
redirect_uri: None,
|
||||
callback_port: None,
|
||||
refresh_buffer_secs: Some(0),
|
||||
};
|
||||
let metadata = serde_json::json!({
|
||||
"issuer": issuer,
|
||||
"token_endpoint": format!("{issuer}/token"),
|
||||
"authorization_endpoint": format!("{issuer}/authorize"),
|
||||
"device_authorization_endpoint": format!("{issuer}/device"),
|
||||
});
|
||||
let device = serde_json::json!({
|
||||
"device_code": "device-code", "user_code": "CLI-CODE",
|
||||
"verification_uri": format!("{issuer}/verify"), "expires_in": 60, "interval": 1,
|
||||
});
|
||||
let requests = Arc::new(Mutex::new(Vec::new()));
|
||||
let device_requests = requests.clone();
|
||||
let auth_requests = requests.clone();
|
||||
let token_requests = requests.clone();
|
||||
let app = Router::new()
|
||||
.route("/.well-known/openid-configuration", get(move || {
|
||||
let metadata = metadata.clone();
|
||||
async move { Json(metadata) }
|
||||
}))
|
||||
.route("/verify", get(|| async { "verified" }))
|
||||
.route("/device", post(move |Form(form): Form<HashMap<String, String>>| {
|
||||
let device = device.clone();
|
||||
device_requests.lock().unwrap().push(form);
|
||||
async move { Json(device) }
|
||||
}))
|
||||
.route("/authorize", get(move |Query(form): Query<HashMap<String, String>>| {
|
||||
assert_eq!(form["code_challenge_method"], "S256");
|
||||
let mut redirect = url::Url::parse(&form["redirect_uri"]).unwrap();
|
||||
redirect.query_pairs_mut().append_pair("code", "auth-code").append_pair("state", &form["state"]);
|
||||
auth_requests.lock().unwrap().push(form);
|
||||
async move { Redirect::to(redirect.as_str()) }
|
||||
}))
|
||||
.route("/token", post(move |headers: HeaderMap, Form(form): Form<HashMap<String, String>>| {
|
||||
let mut requests = token_requests.lock().unwrap();
|
||||
requests.push(form.clone());
|
||||
let number = requests.iter().filter(|request| request.contains_key("grant_type")).count();
|
||||
async move {
|
||||
if form.get("grant_type").is_some_and(|grant| grant == "client_credentials")
|
||||
&& headers.get("authorization").and_then(|value| value.to_str().ok())
|
||||
!= Some(format!("Basic {}", base64::engine::general_purpose::STANDARD.encode("cli-client:test-secret")).as_str()) {
|
||||
return (StatusCode::UNAUTHORIZED, Json(serde_json::json!({
|
||||
"error": "invalid_client", "error_description": "private-token test-secret invalid-secret"
|
||||
})));
|
||||
}
|
||||
(StatusCode::OK, Json(serde_json::json!({
|
||||
"access_token": format!("token-{number}"), "token_type": "Bearer", "expires_in": lifetime,
|
||||
"refresh_token": "private-refresh-token",
|
||||
})))
|
||||
}
|
||||
}));
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
Self {
|
||||
config,
|
||||
requests,
|
||||
server,
|
||||
dir,
|
||||
}
|
||||
}
|
||||
|
||||
fn managed(&self) -> ManagedOidc {
|
||||
ManagedOidc::new(self.config.clone(), self.dir.path().join("cache")).unwrap()
|
||||
}
|
||||
|
||||
fn assert_targets(&self) {
|
||||
for request in self.requests.lock().unwrap().iter() {
|
||||
assert_eq!(request.get("resource"), self.config.resource.as_ref());
|
||||
assert_eq!(request.get("audience"), self.config.audience.as_ref());
|
||||
}
|
||||
}
|
||||
|
||||
async fn child(&self, action: &str) -> std::process::Output {
|
||||
let executable = std::env::current_exe().unwrap();
|
||||
let mut command = tokio::process::Command::new(&executable);
|
||||
command
|
||||
.args([
|
||||
"--exact",
|
||||
"auth::tests::session_process_helper",
|
||||
"--nocapture",
|
||||
])
|
||||
.env("LANCEDB_TEST_AUTH_ACTION", action)
|
||||
.env(
|
||||
"LANCEDB_TEST_AUTH_CONFIG",
|
||||
serde_json::to_string(&self.config).unwrap(),
|
||||
)
|
||||
.env("LANCEDB_TEST_AUTH_CACHE", self.dir.path().join("cache"))
|
||||
.env_remove(CLIENT_SECRET_ENV);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let browser = self.dir.path().join("browser");
|
||||
fs::write(&browser, "#!/bin/sh\nLANCEDB_TEST_BROWSER_URL=\"$1\" exec \"$LANCEDB_TEST_AUTH_EXE\" --exact auth::tests::browser_process_helper --nocapture\n").unwrap();
|
||||
fs::set_permissions(&browser, fs::Permissions::from_mode(0o700)).unwrap();
|
||||
command
|
||||
.env("LANCEDB_OAUTH_BROWSER", browser)
|
||||
.env("LANCEDB_TEST_AUTH_EXE", executable);
|
||||
}
|
||||
let output = command.output().await.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"stdout: {}\nstderr: {}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let output_text = format!(
|
||||
"{}{}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert!(!output_text.contains("private-refresh-token"));
|
||||
output
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_provider_is_shared_and_renews_service_tokens_in_memory() {
|
||||
let fixture = Provider::start(OidcFlow::ClientCredentials, 1).await;
|
||||
let authentication = Authentication::ManagedOidc(Arc::new(fixture.managed()));
|
||||
let cloned = authentication.clone();
|
||||
let (first, second) = tokio::join!(authentication.headers(), cloned.headers());
|
||||
assert_eq!(first.unwrap()["authorization"], "Bearer token-1");
|
||||
assert_eq!(second.unwrap()["authorization"], "Bearer token-1");
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 1);
|
||||
tokio::time::sleep(Duration::from_millis(1100)).await;
|
||||
let headers = authentication.headers().await.unwrap();
|
||||
assert_eq!(headers["authorization"], "Bearer token-2");
|
||||
assert_eq!(headers[OIDC_CREDENTIAL_HEADER], "oidc");
|
||||
assert!(headers["authorization"].is_sensitive());
|
||||
assert!(!format!("{headers:?}").contains("token-2"));
|
||||
assert!(!fixture.dir.path().join("cache").exists());
|
||||
fixture.assert_targets();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_provider_errors_do_not_expose_response_bodies() {
|
||||
let fixture = Provider::start(OidcFlow::ClientCredentials, 3600).await;
|
||||
fs::write(
|
||||
fixture.config.client_secret_file.as_ref().unwrap(),
|
||||
"invalid-secret",
|
||||
)
|
||||
.unwrap();
|
||||
let error = fixture.managed().headers().await.unwrap_err();
|
||||
let rendered = format!("{error:#} {error:?}");
|
||||
assert!(rendered.contains("OIDC authentication failed"));
|
||||
for secret in ["private-token", "test-secret", "invalid-secret"] {
|
||||
assert!(!rendered.contains(secret));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_secret_does_not_fall_back() {
|
||||
let fixture = Provider::start(OidcFlow::ClientCredentials, 3600).await;
|
||||
fs::remove_file(fixture.config.client_secret_file.as_ref().unwrap()).unwrap();
|
||||
assert!(fixture.managed().headers().await.is_err());
|
||||
assert!(fixture.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_provider_reads_updated_external_token_files() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("token");
|
||||
fs::write(&path, "first-token\n").unwrap();
|
||||
let provider = Authentication::OidcAccessTokenFile(path.clone());
|
||||
assert_eq!(
|
||||
provider.get_headers().await.unwrap()["authorization"],
|
||||
"Bearer first-token"
|
||||
);
|
||||
fs::write(&path, "second-token\n").unwrap();
|
||||
assert_eq!(
|
||||
provider.get_headers().await.unwrap()["authorization"],
|
||||
"Bearer second-token"
|
||||
);
|
||||
fs::write(&path, " \n").unwrap();
|
||||
assert!(provider.get_headers().await.is_err());
|
||||
fs::write(&path, "token\r\ninjected: header").unwrap();
|
||||
assert!(provider.get_headers().await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn profile_client_auth_method_round_trips_and_reaches_sdk() {
|
||||
let profile = r#"
|
||||
flow = "authorization_code"
|
||||
issuer_url = "http://localhost:18888/realms/lancedb"
|
||||
client_id = "cli-client"
|
||||
scopes = ["openid"]
|
||||
"#;
|
||||
for (value, expected) in [
|
||||
(None, None),
|
||||
(Some("none"), Some(ClientAuthMethod::None)),
|
||||
(
|
||||
Some("client_secret_basic"),
|
||||
Some(ClientAuthMethod::ClientSecretBasic),
|
||||
),
|
||||
(
|
||||
Some("client_secret_post"),
|
||||
Some(ClientAuthMethod::ClientSecretPost),
|
||||
),
|
||||
] {
|
||||
let text = match value {
|
||||
Some(value) => format!("{profile}\nclient_auth_method = {value:?}\n"),
|
||||
None => profile.to_owned(),
|
||||
};
|
||||
let config: OidcConfig = toml::from_str(&text).unwrap();
|
||||
let serialized = toml::to_string(&config).unwrap();
|
||||
assert_eq!(serialized.contains("client_auth_method"), value.is_some());
|
||||
assert_eq!(toml::from_str::<OidcConfig>(&serialized).unwrap(), config);
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let managed = ManagedOidc::new(config, dir.path().join("cache")).unwrap();
|
||||
assert_eq!(
|
||||
managed.sdk_config().await.unwrap().client_auth_method,
|
||||
expected
|
||||
);
|
||||
}
|
||||
assert!(toml::from_str::<OidcConfig>(&format!(
|
||||
"{profile}\nclient_auth_method = \"unsupported\"\n"
|
||||
))
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_rejects_public_client_auth_with_a_secret_before_requesting_tokens() {
|
||||
let mut fixture = Provider::start(OidcFlow::ClientCredentials, 3600).await;
|
||||
fixture.config.client_auth_method = Some(OidcClientAuthMethod::None);
|
||||
assert!(fixture.managed().provider().await.is_err());
|
||||
assert!(fixture.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sdk_config_preserves_profile_options() {
|
||||
let mut fixture = Provider::start(OidcFlow::AuthorizationCode, 3600).await;
|
||||
fixture.config.callback_port = Some(9123);
|
||||
fixture.config.redirect_uri = Some("http://127.0.0.1:9123/cli".into());
|
||||
let config = fixture.managed().sdk_config().await.unwrap();
|
||||
assert_eq!(config.resource, fixture.config.resource);
|
||||
assert_eq!(config.audience, fixture.config.audience);
|
||||
assert_eq!(config.scopes, fixture.config.scopes);
|
||||
assert_eq!(config.refresh_buffer_secs, Some(0));
|
||||
assert_eq!(
|
||||
config.token_cache.unwrap().cache_dir,
|
||||
Some(fixture.dir.path().join("cache"))
|
||||
);
|
||||
let OAuthFlow::AuthorizationCode(options) = config.flow else {
|
||||
panic!("expected browser flow")
|
||||
};
|
||||
assert!(options.use_pkce);
|
||||
assert_eq!(options.callback_port, Some(9123));
|
||||
assert_eq!(options.redirect_uri, fixture.config.redirect_uri);
|
||||
fixture.config.flow = OidcFlow::DeviceCode;
|
||||
assert!(ManagedOidc::new(fixture.config.clone(), fixture.dir.path().join("cache")).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn profile_caches_are_isolated_without_creating_directories() {
|
||||
let fixture = Provider::start(OidcFlow::DeviceCode, 3600).await;
|
||||
let first = ManagedOidc::for_profile("a/../b", fixture.config.clone()).unwrap();
|
||||
let second = ManagedOidc::for_profile("b", fixture.config.clone()).unwrap();
|
||||
assert_ne!(first.cache_dir, second.cache_dir);
|
||||
assert_eq!(first.cache_dir.parent(), second.cache_dir.parent());
|
||||
assert_eq!(first.cache_dir.file_name().unwrap().len(), 64);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn sdk_interactive_sessions_survive_processes_and_logout() {
|
||||
for flow in [OidcFlow::AuthorizationCode, OidcFlow::DeviceCode] {
|
||||
let mut fixture = Provider::start(flow, 3600).await;
|
||||
if fixture.config.flow == OidcFlow::AuthorizationCode {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
fixture.config.callback_port = Some(listener.local_addr().unwrap().port());
|
||||
}
|
||||
// Ordinary requests trigger acquisition, without an explicit login first.
|
||||
fixture.child("fetch").await;
|
||||
let original_requests = fixture.requests.lock().unwrap().len();
|
||||
fixture.child("fetch").await;
|
||||
{
|
||||
let requests = fixture.requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), original_requests + 1);
|
||||
assert_eq!(requests.last().unwrap()["grant_type"], "refresh_token");
|
||||
}
|
||||
fixture.child("status").await;
|
||||
fixture.child("logout").await;
|
||||
assert!(
|
||||
!fixture
|
||||
.managed()
|
||||
.session()
|
||||
.await
|
||||
.unwrap()
|
||||
.status()
|
||||
.await
|
||||
.unwrap()
|
||||
.refreshable
|
||||
);
|
||||
fixture.child("login").await;
|
||||
fixture.child("status").await;
|
||||
fixture.assert_targets();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_process_helper() {
|
||||
let Ok(action) = std::env::var("LANCEDB_TEST_AUTH_ACTION") else {
|
||||
return;
|
||||
};
|
||||
let config = serde_json::from_str(&std::env::var("LANCEDB_TEST_AUTH_CONFIG").unwrap()).unwrap();
|
||||
let managed = ManagedOidc::new(
|
||||
config,
|
||||
std::env::var_os("LANCEDB_TEST_AUTH_CACHE").unwrap().into(),
|
||||
)
|
||||
.unwrap();
|
||||
match action.as_str() {
|
||||
"fetch" => {
|
||||
managed.headers().await.unwrap();
|
||||
}
|
||||
"login" => {
|
||||
assert!(
|
||||
managed
|
||||
.session()
|
||||
.await
|
||||
.unwrap()
|
||||
.login()
|
||||
.await
|
||||
.unwrap()
|
||||
.refreshable
|
||||
);
|
||||
}
|
||||
"status" => {
|
||||
assert!(
|
||||
managed
|
||||
.session()
|
||||
.await
|
||||
.unwrap()
|
||||
.status()
|
||||
.await
|
||||
.unwrap()
|
||||
.refreshable
|
||||
);
|
||||
}
|
||||
"logout" => {
|
||||
assert!(
|
||||
managed
|
||||
.session()
|
||||
.await
|
||||
.unwrap()
|
||||
.logout()
|
||||
.await
|
||||
.unwrap()
|
||||
.removed
|
||||
);
|
||||
}
|
||||
_ => panic!("unknown action"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn browser_process_helper() {
|
||||
let Ok(url) = std::env::var("LANCEDB_TEST_BROWSER_URL") else {
|
||||
return;
|
||||
};
|
||||
reqwest::Client::new()
|
||||
.get(url)
|
||||
.send()
|
||||
.await
|
||||
.unwrap()
|
||||
.error_for_status()
|
||||
.unwrap();
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use crate::server_config::{ConnectionArgs, ResolvedConnection};
|
||||
use anyhow::{bail, Context, Result};
|
||||
use clap::{CommandFactory, Parser, Subcommand};
|
||||
|
||||
/// Public client commands. Runtime and administrative commands live in Enterprise.
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(
|
||||
name = "lancedb",
|
||||
version,
|
||||
about = "LanceDB command-line client",
|
||||
args_conflicts_with_subcommands = true
|
||||
)]
|
||||
pub struct Args {
|
||||
#[command(subcommand)]
|
||||
pub command: Option<Commands>,
|
||||
/// SQL options used when no subcommand is given. These are local to the
|
||||
/// default SQL invocation, so other subcommands use their own options.
|
||||
#[command(flatten)]
|
||||
pub sql: SqlArgs,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum Commands {
|
||||
/// Manage OAuth login and cached interactive sessions
|
||||
Auth(crate::auth::AuthArgs),
|
||||
/// Manage connection profiles
|
||||
Profile(crate::profile::ProfileArgs),
|
||||
/// Run SQL or open an interactive shell against a remote server
|
||||
Sql(SqlArgs),
|
||||
}
|
||||
|
||||
#[derive(clap::Args, Debug)]
|
||||
pub struct SqlArgs {
|
||||
#[command(flatten)]
|
||||
pub connection: ConnectionArgs,
|
||||
/// Execute SQL and exit
|
||||
#[arg(short = 'c', long = "command", conflicts_with = "file")]
|
||||
pub command_string: Option<String>,
|
||||
/// Execute SQL commands from a file and exit
|
||||
#[arg(short = 'f', long)]
|
||||
pub file: Option<String>,
|
||||
/// Output format: table, csv, json, ndjson
|
||||
#[arg(long, default_value = "table")]
|
||||
pub format: String,
|
||||
}
|
||||
|
||||
const SQL_ENDPOINT_HELP: &str = "No Flight SQL endpoint is configured. Use --sql-server-url, \
|
||||
select a profile with sql_server_url, or use --profile local for localhost.";
|
||||
|
||||
pub async fn run_sql(args: SqlArgs) -> Result<()> {
|
||||
let connection = args.connection.resolve_sql()?.context(SQL_ENDPOINT_HELP)?;
|
||||
run_sql_with_connection(args, connection).await
|
||||
}
|
||||
|
||||
async fn run_sql_with_connection(args: SqlArgs, connection: ResolvedConnection) -> Result<()> {
|
||||
if let Some(file) = args.file {
|
||||
crate::client::run_client_script_with_writer(
|
||||
&connection,
|
||||
std::fs::read_to_string(file)?,
|
||||
args.format,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
} else if let Some(sql) = args.command_string {
|
||||
crate::client::run_client_command_with_writer(&connection, sql, args.format, None).await
|
||||
} else {
|
||||
crate::client::run_client_with_writer(&connection, None).await
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run(args: Args) -> Result<()> {
|
||||
match args.command {
|
||||
None => match args.sql.connection.resolve_sql()? {
|
||||
Some(connection) => run_sql_with_connection(args.sql, connection).await,
|
||||
None if args.sql.command_string.is_none()
|
||||
&& args.sql.file.is_none()
|
||||
&& args.sql.connection.profile.is_none() =>
|
||||
{
|
||||
println!("{SQL_ENDPOINT_HELP}\n");
|
||||
Args::command().print_help()?;
|
||||
println!();
|
||||
Ok(())
|
||||
}
|
||||
None => bail!(SQL_ENDPOINT_HELP),
|
||||
},
|
||||
Some(Commands::Auth(args)) => crate::auth::run(args, None).await,
|
||||
Some(Commands::Profile(args)) => crate::profile::run_profile_command(args).await,
|
||||
Some(Commands::Sql(args)) => run_sql(args).await,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn sql_is_the_default_command() {
|
||||
let args = Args::try_parse_from(["lancedb"]).unwrap();
|
||||
assert!(args.command.is_none());
|
||||
assert_eq!(args.sql.format, "table");
|
||||
assert!(args.sql.command_string.is_none());
|
||||
assert!(args.sql.file.is_none());
|
||||
assert!(matches!(
|
||||
Args::try_parse_from(["lancedb", "sql"]).unwrap().command,
|
||||
Some(Commands::Sql(_))
|
||||
));
|
||||
let args = Args::try_parse_from([
|
||||
"lancedb",
|
||||
"--profile",
|
||||
"production",
|
||||
"--command",
|
||||
"SELECT 1",
|
||||
"--format",
|
||||
"json",
|
||||
])
|
||||
.unwrap();
|
||||
assert!(args.command.is_none());
|
||||
assert_eq!(args.sql.connection.profile.as_deref(), Some("production"));
|
||||
assert_eq!(args.sql.command_string.as_deref(), Some("SELECT 1"));
|
||||
assert_eq!(args.sql.format, "json");
|
||||
let args = Args::try_parse_from(["lancedb", "--file", "queries.sql"]).unwrap();
|
||||
assert_eq!(args.sql.file.as_deref(), Some("queries.sql"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_sql_options_do_not_apply_to_subcommands() {
|
||||
for args in [
|
||||
vec!["lancedb", "--format", "json", "profile", "list"],
|
||||
vec!["lancedb", "--profile", "production", "profile", "list"],
|
||||
vec!["lancedb", "--command", "SELECT 1", "profile", "list"],
|
||||
vec!["lancedb", "profile", "list", "--command", "SELECT 1"],
|
||||
vec!["lancedb", "profile", "list", "--file", "queries.sql"],
|
||||
vec!["lancedb", "-c", "SELECT 1", "-f", "query.sql"],
|
||||
vec!["lancedb", "sql", "-c", "SELECT 1", "-f", "query.sql"],
|
||||
] {
|
||||
assert!(Args::try_parse_from(&args).is_err(), "{args:?}");
|
||||
}
|
||||
assert!(Args::try_parse_from(["lancedb", "profile", "list"]).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn excluded_commands_are_rejected() {
|
||||
for command in [
|
||||
"server",
|
||||
"admin",
|
||||
"bench",
|
||||
"service",
|
||||
"pe",
|
||||
"agent",
|
||||
"management",
|
||||
"preflight",
|
||||
"run",
|
||||
"function-worker",
|
||||
"job",
|
||||
"acl",
|
||||
"load",
|
||||
] {
|
||||
assert!(
|
||||
Args::try_parse_from(["lancedb", command]).is_err(),
|
||||
"{command}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,137 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Typed DML result serializable to/from RecordBatch for Arrow Flight transport.
|
||||
|
||||
use arrow_array::{Int64Array, RecordBatch, StringArray};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum DmlOperation {
|
||||
Insert,
|
||||
Update,
|
||||
Delete,
|
||||
}
|
||||
|
||||
impl DmlOperation {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Insert => "INSERT",
|
||||
Self::Update => "UPDATE",
|
||||
Self::Delete => "DELETE",
|
||||
}
|
||||
}
|
||||
|
||||
fn parse(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"INSERT" => Some(Self::Insert),
|
||||
"UPDATE" => Some(Self::Update),
|
||||
"DELETE" => Some(Self::Delete),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for DmlOperation {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DmlResult {
|
||||
pub table: String,
|
||||
pub operation: DmlOperation,
|
||||
pub rows_affected: i64,
|
||||
pub version: i64,
|
||||
}
|
||||
|
||||
fn schema() -> Schema {
|
||||
Schema::new(vec![
|
||||
Field::new("table", DataType::Utf8, false),
|
||||
Field::new("operation", DataType::Utf8, false),
|
||||
Field::new("rows_affected", DataType::Int64, false),
|
||||
Field::new("version", DataType::Int64, false),
|
||||
])
|
||||
}
|
||||
|
||||
impl DmlResult {
|
||||
pub fn new(
|
||||
table: impl Into<String>,
|
||||
operation: DmlOperation,
|
||||
rows_affected: i64,
|
||||
version: i64,
|
||||
) -> Self {
|
||||
Self {
|
||||
table: table.into(),
|
||||
operation,
|
||||
rows_affected,
|
||||
version,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_record_batch(&self) -> RecordBatch {
|
||||
RecordBatch::try_new(
|
||||
Arc::new(schema()),
|
||||
vec![
|
||||
Arc::new(StringArray::from(vec![self.table.as_str()])),
|
||||
Arc::new(StringArray::from(vec![self.operation.as_str()])),
|
||||
Arc::new(Int64Array::from(vec![self.rows_affected])),
|
||||
Arc::new(Int64Array::from(vec![self.version])),
|
||||
],
|
||||
)
|
||||
.expect("static schema")
|
||||
}
|
||||
|
||||
pub fn try_from_batch(batch: &RecordBatch) -> Option<Self> {
|
||||
if *batch.schema().as_ref() != schema() || batch.num_rows() != 1 {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
table: batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()?
|
||||
.value(0)
|
||||
.to_string(),
|
||||
operation: DmlOperation::parse(
|
||||
batch
|
||||
.column(1)
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()?
|
||||
.value(0),
|
||||
)?,
|
||||
rows_affected: batch
|
||||
.column(2)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()?
|
||||
.value(0),
|
||||
version: batch
|
||||
.column(3)
|
||||
.as_any()
|
||||
.downcast_ref::<Int64Array>()?
|
||||
.value(0),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn roundtrip() {
|
||||
let original = DmlResult::new("foo", DmlOperation::Insert, 1, 3);
|
||||
let batch = original.to_record_batch();
|
||||
assert_eq!(DmlResult::try_from_batch(&batch), Some(original));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_dml_returns_none() {
|
||||
let s = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
|
||||
let batch = RecordBatch::try_new(s, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap();
|
||||
assert_eq!(DmlResult::try_from_batch(&batch), None);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,55 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
pub mod auth;
|
||||
pub mod cli;
|
||||
pub mod client;
|
||||
pub mod dml_result;
|
||||
pub mod formatter;
|
||||
pub mod mocks;
|
||||
pub mod profile;
|
||||
pub mod server_config;
|
||||
pub mod sql_executor;
|
||||
pub mod table_writer;
|
||||
|
||||
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
pub use cli::{run, Args, Commands};
|
||||
pub(crate) fn format_host_port(host: &str, port: impl std::fmt::Display) -> String {
|
||||
match host.parse::<std::net::IpAddr>() {
|
||||
Ok(std::net::IpAddr::V6(_)) => format!("[{host}]:{port}"),
|
||||
_ => format!("{host}:{port}"),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn lancedb_client_config(
|
||||
tls: &server_config::TlsOptions,
|
||||
authentication: &server_config::Authentication,
|
||||
) -> Option<lancedb::remote::ClientConfig> {
|
||||
if tls == &server_config::TlsOptions::default()
|
||||
&& matches!(authentication, server_config::Authentication::ApiKey(_))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let mut config = lancedb::remote::ClientConfig::default();
|
||||
if tls != &server_config::TlsOptions::default() {
|
||||
config.tls_config = Some(lancedb::remote::TlsConfig {
|
||||
ssl_ca_cert: tls.ca_cert.clone(),
|
||||
cert_file: tls.client_cert.clone(),
|
||||
key_file: tls.client_key.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
if !matches!(authentication, server_config::Authentication::ApiKey(_)) {
|
||||
config.header_provider = Some(std::sync::Arc::new(authentication.clone()));
|
||||
}
|
||||
Some(config)
|
||||
}
|
||||
|
||||
static BUILD_VERSION: std::sync::OnceLock<&'static str> = std::sync::OnceLock::new();
|
||||
/// Set the version displayed by shared command handlers when embedded in another binary.
|
||||
pub fn set_build_version(version: &'static str) {
|
||||
let _ = BUILD_VERSION.set(version);
|
||||
}
|
||||
pub fn build_version() -> &'static str {
|
||||
BUILD_VERSION.get().copied().unwrap_or(VERSION)
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use clap::Parser;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
lancedb_cli::run(lancedb_cli::Args::parse()).await
|
||||
}
|
||||
@@ -0,0 +1,119 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
/// Mock implementations for testing
|
||||
#[cfg(test)]
|
||||
pub mod test_mocks {
|
||||
use crate::sql_executor::{SqlQueryExecutor, SqlResult};
|
||||
use anyhow::Result;
|
||||
use arrow_array::RecordBatch;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// Mock SQL executor for testing
|
||||
pub struct MockSqlExecutor {
|
||||
pub is_connected: bool,
|
||||
pub queries: Arc<Mutex<Vec<String>>>,
|
||||
pub query_results: HashMap<String, Vec<RecordBatch>>,
|
||||
/// URLs passed to `connect`, in call order.
|
||||
pub connected_urls: Arc<Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
impl Default for MockSqlExecutor {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl MockSqlExecutor {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
is_connected: false,
|
||||
queries: Arc::new(Mutex::new(Vec::new())),
|
||||
query_results: HashMap::new(),
|
||||
connected_urls: Arc::new(Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_query_result(&mut self, query: &str, result: Vec<RecordBatch>) {
|
||||
self.query_results.insert(query.to_string(), result);
|
||||
}
|
||||
|
||||
pub fn get_executed_queries(&self) -> Vec<String> {
|
||||
self.queries.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
pub fn get_connected_urls(&self) -> Vec<String> {
|
||||
self.connected_urls.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SqlQueryExecutor for MockSqlExecutor {
|
||||
async fn execute_query(&mut self, query: &str) -> Result<SqlResult> {
|
||||
self.queries.lock().unwrap().push(query.to_string());
|
||||
|
||||
if let Some(result) = self.query_results.get(query) {
|
||||
Ok(SqlResult::Query(result.clone()))
|
||||
} else {
|
||||
// Return empty result by default
|
||||
Ok(SqlResult::Query(vec![]))
|
||||
}
|
||||
}
|
||||
|
||||
async fn is_connected(&self) -> bool {
|
||||
self.is_connected
|
||||
}
|
||||
|
||||
fn connection_info(&self) -> Option<String> {
|
||||
if self.is_connected {
|
||||
Some("Mock SQL connection".to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect(&mut self, server_url: &str) -> Result<()> {
|
||||
self.connected_urls
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(server_url.to_string());
|
||||
self.is_connected = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_database(&mut self, _database: String) {
|
||||
// Mock implementation - no-op
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper function to create a simple test schema
|
||||
pub fn create_test_schema() -> Arc<Schema> {
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("name", DataType::Utf8, false),
|
||||
Field::new("value", DataType::Float64, true),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Helper function to create test RecordBatch
|
||||
pub fn create_test_batch() -> RecordBatch {
|
||||
use arrow_array::{Float64Array, Int64Array, StringArray};
|
||||
|
||||
let id_array = Int64Array::from(vec![1, 2, 3]);
|
||||
let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
|
||||
let value_array = Float64Array::from(vec![Some(10.5), None, Some(20.0)]);
|
||||
|
||||
RecordBatch::try_new(
|
||||
create_test_schema(),
|
||||
vec![
|
||||
Arc::new(id_array),
|
||||
Arc::new(name_array),
|
||||
Arc::new(value_array),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,494 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::borrow::Cow;
|
||||
use std::io::{self, Write};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use clap::{Args, Subcommand};
|
||||
use rustyline::highlight::Highlighter;
|
||||
use rustyline::{Editor, Helper};
|
||||
|
||||
use crate::server_config::{Config, Profile};
|
||||
|
||||
#[derive(Args, Debug)]
|
||||
pub struct ProfileArgs {
|
||||
#[command(subcommand)]
|
||||
pub command: ProfileCommand,
|
||||
}
|
||||
|
||||
/// The settable fields of a profile, shared by `create` and `edit`.
|
||||
///
|
||||
/// `create` writes them all; `edit` applies only the ones that were passed.
|
||||
#[derive(Args, Debug, Default)]
|
||||
pub struct ProfileFields {
|
||||
/// SQL server URL (Flight SQL endpoint), as host:port or grpc(s)://host:port
|
||||
#[arg(long)]
|
||||
pub sql_server_url: Option<String>,
|
||||
|
||||
/// HTTP server URL (REST endpoint), as host:port or http(s)://host:port
|
||||
#[arg(long)]
|
||||
pub http_server_url: Option<String>,
|
||||
|
||||
/// API key for authentication
|
||||
#[arg(long, conflicts_with = "oidc_access_token_file")]
|
||||
pub api_key: Option<String>,
|
||||
|
||||
/// Path to a file containing an OIDC access token
|
||||
#[arg(long, value_name = "PATH", conflicts_with = "api_key")]
|
||||
pub oidc_access_token_file: Option<String>,
|
||||
|
||||
/// Default database name
|
||||
#[arg(long, short)]
|
||||
pub database: Option<String>,
|
||||
|
||||
/// Azure storage account name for Azure-backed databases
|
||||
#[arg(long)]
|
||||
pub azure_storage_account_name: Option<String>,
|
||||
|
||||
/// Path to a PEM CA certificate trusted in addition to the platform roots
|
||||
#[arg(long, value_name = "PATH")]
|
||||
pub ca_cert: Option<String>,
|
||||
|
||||
/// Path to a PEM client certificate presented for mTLS
|
||||
#[arg(long, value_name = "PATH")]
|
||||
pub client_cert: Option<String>,
|
||||
|
||||
/// Path to the PKCS#8 PEM private key for the client certificate
|
||||
#[arg(long, value_name = "PATH")]
|
||||
pub client_key: Option<String>,
|
||||
}
|
||||
|
||||
impl ProfileFields {
|
||||
/// Whether any field was supplied, which is what makes `create`
|
||||
/// non-interactive. Destructured so a new field cannot be forgotten here.
|
||||
fn any_set(&self) -> bool {
|
||||
let Self {
|
||||
sql_server_url,
|
||||
http_server_url,
|
||||
api_key,
|
||||
oidc_access_token_file,
|
||||
database,
|
||||
azure_storage_account_name,
|
||||
ca_cert,
|
||||
client_cert,
|
||||
client_key,
|
||||
} = self;
|
||||
[
|
||||
sql_server_url,
|
||||
http_server_url,
|
||||
api_key,
|
||||
oidc_access_token_file,
|
||||
database,
|
||||
azure_storage_account_name,
|
||||
ca_cert,
|
||||
client_cert,
|
||||
client_key,
|
||||
]
|
||||
.iter()
|
||||
.any(|field| field.is_some())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&ProfileFields> for Profile {
|
||||
fn from(fields: &ProfileFields) -> Self {
|
||||
Self {
|
||||
sql_server_url: fields.sql_server_url.clone(),
|
||||
http_server_url: fields.http_server_url.clone(),
|
||||
api_key: fields.api_key.clone(),
|
||||
oidc_access_token_file: fields.oidc_access_token_file.clone(),
|
||||
oidc: None,
|
||||
database: fields.database.clone(),
|
||||
azure_storage_account_name: fields.azure_storage_account_name.clone(),
|
||||
ca_cert: fields.ca_cert.clone(),
|
||||
client_cert: fields.client_cert.clone(),
|
||||
client_key: fields.client_key.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
pub enum ProfileCommand {
|
||||
/// List all configured profiles
|
||||
List,
|
||||
|
||||
/// Show details of a specific profile
|
||||
Show {
|
||||
/// Name of the profile to show (uses default if not specified)
|
||||
name: Option<String>,
|
||||
},
|
||||
|
||||
/// Create a new profile (interactive by default)
|
||||
Create {
|
||||
/// Name for the new profile
|
||||
name: String,
|
||||
|
||||
#[command(flatten)]
|
||||
fields: ProfileFields,
|
||||
},
|
||||
|
||||
/// Set the default profile
|
||||
#[command(name = "set")]
|
||||
SetDefault {
|
||||
/// Name of the profile to set as default
|
||||
name: String,
|
||||
},
|
||||
|
||||
/// Delete a profile
|
||||
Delete {
|
||||
/// Name of the profile to delete
|
||||
name: String,
|
||||
},
|
||||
|
||||
/// Edit an existing profile
|
||||
Edit {
|
||||
/// Name of the profile to edit
|
||||
name: String,
|
||||
|
||||
#[command(flatten)]
|
||||
fields: ProfileFields,
|
||||
},
|
||||
}
|
||||
|
||||
pub async fn run_profile_command(args: ProfileArgs) -> Result<()> {
|
||||
match args.command {
|
||||
ProfileCommand::List => list_profiles(),
|
||||
ProfileCommand::Show { name } => show_profile(name),
|
||||
ProfileCommand::Create { name, fields } => {
|
||||
// If any flags are provided, use non-interactive mode
|
||||
if fields.any_set() {
|
||||
create_profile(&name, Profile::from(&fields))
|
||||
} else {
|
||||
create_profile_interactive(&name)
|
||||
}
|
||||
}
|
||||
ProfileCommand::SetDefault { name } => set_default_profile(&name),
|
||||
ProfileCommand::Delete { name } => delete_profile(&name),
|
||||
ProfileCommand::Edit { name, fields } => edit_profile(&name, Profile::from(&fields)),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_profiles() -> Result<()> {
|
||||
let config = Config::load_default()?;
|
||||
let default_profile = config.default_profile.as_deref().unwrap_or("local");
|
||||
|
||||
println!("Configured profiles:");
|
||||
|
||||
// Always show "local" as it's built-in
|
||||
let is_default = default_profile == "local";
|
||||
let default_marker = if is_default { "*" } else { " " };
|
||||
let in_config = if config.profiles.contains_key("local") {
|
||||
""
|
||||
} else {
|
||||
" [built-in]"
|
||||
};
|
||||
println!(" {} local{}", default_marker, in_config);
|
||||
|
||||
// Show configured profiles
|
||||
let mut profile_names: Vec<_> = config.profiles.keys().collect();
|
||||
profile_names.sort();
|
||||
|
||||
for name in profile_names {
|
||||
if name == "local" {
|
||||
continue; // Already shown above
|
||||
}
|
||||
let is_default = default_profile == name;
|
||||
let default_marker = if is_default { "*" } else { " " };
|
||||
println!(" {} {}", default_marker, name);
|
||||
}
|
||||
|
||||
if config.profiles.is_empty()
|
||||
|| (config.profiles.len() == 1 && config.profiles.contains_key("local"))
|
||||
{
|
||||
println!();
|
||||
println!("No custom profiles configured. Use 'lancedb profile create <name>' to add one.");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn show_profile(name: Option<String>) -> Result<()> {
|
||||
let config = Config::load_default()?;
|
||||
let profile_name = name
|
||||
.as_deref()
|
||||
.or(config.default_profile.as_deref())
|
||||
.unwrap_or("local");
|
||||
|
||||
let profile = config.get_profile(Some(profile_name));
|
||||
let is_builtin = !config.profiles.contains_key(profile_name) && profile_name == "local";
|
||||
|
||||
let builtin_marker = if is_builtin { " [built-in]" } else { "" };
|
||||
println!("Profile: {}{}", profile_name, builtin_marker);
|
||||
println!(
|
||||
" sql_server_url: {}",
|
||||
profile.sql_server_url.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" http_server_url: {}",
|
||||
profile.http_server_url.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" api_key: {}",
|
||||
profile
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|_| "********")
|
||||
.unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" oidc_access_token_file: {}",
|
||||
profile
|
||||
.oidc_access_token_file
|
||||
.as_deref()
|
||||
.unwrap_or("<not set>")
|
||||
);
|
||||
if let Some(oidc) = &profile.oidc {
|
||||
println!(" oidc: {}", toml::to_string(oidc)?);
|
||||
}
|
||||
println!(
|
||||
" database: {}",
|
||||
profile.database.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" azure_storage_account_name: {}",
|
||||
profile
|
||||
.azure_storage_account_name
|
||||
.as_deref()
|
||||
.unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" ca_cert: {}",
|
||||
profile.ca_cert.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" client_cert: {}",
|
||||
profile.client_cert.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
println!(
|
||||
" client_key: {}",
|
||||
profile.client_key.as_deref().unwrap_or("<not set>")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_profile(name: &str, profile: Profile) -> Result<()> {
|
||||
let mut config = Config::load_default()?;
|
||||
|
||||
if config.profiles.contains_key(name) {
|
||||
return Err(anyhow!(
|
||||
"Profile '{}' already exists. Use 'lancedb profile edit {}' to modify it.",
|
||||
name,
|
||||
name
|
||||
));
|
||||
}
|
||||
|
||||
config.set_profile(name, profile);
|
||||
config.save_default()?;
|
||||
|
||||
println!("Created profile '{}'", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_profile_interactive(name: &str) -> Result<()> {
|
||||
let mut config = Config::load_default()?;
|
||||
|
||||
if config.profiles.contains_key(name) {
|
||||
return Err(anyhow!(
|
||||
"Profile '{}' already exists. Use 'lancedb profile edit {}' to modify it.",
|
||||
name,
|
||||
name
|
||||
));
|
||||
}
|
||||
|
||||
println!("Creating profile '{}'", name);
|
||||
println!("Press Enter to skip optional fields.");
|
||||
println!();
|
||||
|
||||
let sql_server_url =
|
||||
prompt_optional("SQL server URL (host:port, or grpcs://host:port for TLS)")?;
|
||||
let http_server_url =
|
||||
prompt_optional("HTTP server URL (host:port, or https://host:port for TLS)")?;
|
||||
let api_key = prompt_secret("API key")?;
|
||||
let oidc_access_token_file = if api_key.is_none() {
|
||||
prompt_optional("OIDC access token file")?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let database = prompt_optional("Default database")?;
|
||||
let azure_storage_account_name = prompt_optional("Azure storage account name")?;
|
||||
let ca_cert = prompt_optional("CA certificate path (PEM, for a private CA)")?;
|
||||
let client_cert = prompt_optional("Client certificate path (PEM, for mTLS)")?;
|
||||
let client_key = prompt_optional("Client key path (PKCS#8 PEM, for mTLS)")?;
|
||||
|
||||
let profile = Profile {
|
||||
sql_server_url,
|
||||
http_server_url,
|
||||
api_key,
|
||||
oidc_access_token_file,
|
||||
oidc: None,
|
||||
database,
|
||||
azure_storage_account_name,
|
||||
ca_cert,
|
||||
client_cert,
|
||||
client_key,
|
||||
};
|
||||
config.set_profile(name, profile);
|
||||
config.save_default()?;
|
||||
|
||||
println!();
|
||||
println!("Created profile '{}'", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_default_profile(name: &str) -> Result<()> {
|
||||
let mut config = Config::load_default()?;
|
||||
|
||||
// Verify profile exists (or is built-in "local")
|
||||
if name != "local" && !config.profiles.contains_key(name) {
|
||||
return Err(anyhow!(
|
||||
"Profile '{}' does not exist. Use 'lancedb profile list' to see available profiles.",
|
||||
name
|
||||
));
|
||||
}
|
||||
|
||||
// Persist the addresses when localhost is explicitly selected, so a later
|
||||
// default SQL invocation can distinguish this choice from an unset endpoint.
|
||||
if name == "local" && !config.profiles.contains_key(name) {
|
||||
config.set_profile(name, Profile::local_default());
|
||||
}
|
||||
config.set_default_profile(name);
|
||||
config.save_default()?;
|
||||
|
||||
println!("Set '{}' as the default profile", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn delete_profile(name: &str) -> Result<()> {
|
||||
if name == "local" {
|
||||
return Err(anyhow!("Cannot delete the built-in 'local' profile"));
|
||||
}
|
||||
|
||||
let mut config = Config::load_default()?;
|
||||
|
||||
if !config.profiles.contains_key(name) {
|
||||
return Err(anyhow!("Profile '{}' does not exist", name));
|
||||
}
|
||||
|
||||
config.remove_profile(name);
|
||||
|
||||
// If we deleted the default profile, reset to "local"
|
||||
if config.default_profile.as_deref() == Some(name) {
|
||||
config.set_default_profile("local");
|
||||
println!("Default profile reset to 'local'");
|
||||
}
|
||||
|
||||
config.save_default()?;
|
||||
|
||||
println!("Deleted profile '{}'", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn edit_profile(name: &str, updates: Profile) -> Result<()> {
|
||||
let mut config = Config::load_default()?;
|
||||
|
||||
// Get existing profile or create from built-in
|
||||
let mut profile = if config.profiles.contains_key(name) {
|
||||
config.profiles.get(name).cloned().unwrap()
|
||||
} else if name == "local" {
|
||||
Profile::local_default()
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"Profile '{}' does not exist. Use 'lancedb profile create {}' to create it.",
|
||||
name,
|
||||
name
|
||||
));
|
||||
};
|
||||
|
||||
// Update only the fields that were provided
|
||||
let Profile {
|
||||
sql_server_url,
|
||||
http_server_url,
|
||||
api_key,
|
||||
oidc_access_token_file,
|
||||
oidc,
|
||||
database,
|
||||
azure_storage_account_name,
|
||||
ca_cert,
|
||||
client_cert,
|
||||
client_key,
|
||||
} = updates;
|
||||
profile.sql_server_url = sql_server_url.or(profile.sql_server_url);
|
||||
profile.http_server_url = http_server_url.or(profile.http_server_url);
|
||||
if let Some(api_key) = api_key {
|
||||
profile.oidc = None;
|
||||
profile.api_key = Some(api_key);
|
||||
profile.oidc_access_token_file = None;
|
||||
} else if let Some(oidc_access_token_file) = oidc_access_token_file {
|
||||
profile.oidc = None;
|
||||
profile.api_key = None;
|
||||
profile.oidc_access_token_file = Some(oidc_access_token_file);
|
||||
}
|
||||
if let Some(oidc) = oidc {
|
||||
profile.api_key = None;
|
||||
profile.oidc_access_token_file = None;
|
||||
profile.oidc = Some(oidc);
|
||||
}
|
||||
profile.database = database.or(profile.database);
|
||||
profile.azure_storage_account_name =
|
||||
azure_storage_account_name.or(profile.azure_storage_account_name);
|
||||
profile.ca_cert = ca_cert.or(profile.ca_cert);
|
||||
profile.client_cert = client_cert.or(profile.client_cert);
|
||||
profile.client_key = client_key.or(profile.client_key);
|
||||
|
||||
config.set_profile(name, profile);
|
||||
config.save_default()?;
|
||||
|
||||
println!("Updated profile '{}'", name);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prompt_optional(field: &str) -> Result<Option<String>> {
|
||||
print!("{}: ", field);
|
||||
io::stdout().flush()?;
|
||||
|
||||
let mut input = String::new();
|
||||
io::stdin().read_line(&mut input)?;
|
||||
|
||||
let value = input.trim();
|
||||
if value.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(value.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_secret(field: &str) -> Result<Option<String>> {
|
||||
struct MaskHelper;
|
||||
impl Highlighter for MaskHelper {
|
||||
fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
|
||||
Cow::Owned("*".repeat(line.len()))
|
||||
}
|
||||
fn highlight_char(&self, _line: &str, _pos: usize, _forced: bool) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
impl rustyline::completion::Completer for MaskHelper {
|
||||
type Candidate = String;
|
||||
}
|
||||
impl rustyline::hint::Hinter for MaskHelper {
|
||||
type Hint = String;
|
||||
}
|
||||
impl rustyline::validate::Validator for MaskHelper {}
|
||||
impl Helper for MaskHelper {}
|
||||
|
||||
let mut editor = Editor::new()?;
|
||||
editor.set_helper(Some(MaskHelper));
|
||||
let input = editor.readline(&format!("{}: ", field))?;
|
||||
let value = input.trim();
|
||||
if value.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(value.to_string()))
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,140 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use crate::dml_result::DmlResult;
|
||||
use anyhow::Result;
|
||||
use arrow_array::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[async_trait]
|
||||
pub trait SqlQueryExecutor: Send + Sync {
|
||||
/// Execute a SQL query and return a structured result.
|
||||
async fn execute_query(&mut self, query: &str) -> Result<SqlResult>;
|
||||
|
||||
/// Check if the executor is connected and ready
|
||||
async fn is_connected(&self) -> bool;
|
||||
|
||||
/// Get the connection info for display purposes
|
||||
fn connection_info(&self) -> Option<String>;
|
||||
|
||||
/// Connect to the server. `server_url` may carry an `http://` or `https://`
|
||||
/// scheme; a bare `host:port` is treated as plaintext.
|
||||
async fn connect(&mut self, server_url: &str) -> Result<()>;
|
||||
|
||||
/// Set the current database
|
||||
fn set_database(&mut self, database: String);
|
||||
}
|
||||
|
||||
/// Result of a SQL execution
|
||||
#[derive(Debug)]
|
||||
pub enum SqlResult {
|
||||
/// Query returned data
|
||||
Query(Vec<RecordBatch>),
|
||||
/// DML result (INSERT, DELETE, UPDATE)
|
||||
Dml(DmlResult),
|
||||
/// DDL statement (CREATE, DROP, etc.) completed
|
||||
Success,
|
||||
}
|
||||
|
||||
impl From<Vec<RecordBatch>> for SqlResult {
|
||||
fn from(batches: Vec<RecordBatch>) -> Self {
|
||||
match batches.first().and_then(DmlResult::try_from_batch) {
|
||||
Some(dml) => Self::Dml(dml),
|
||||
None => Self::Query(batches),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SQL execution through the public SDK's asynchronous query lifecycle.
|
||||
pub struct SdkSqlExecutor {
|
||||
connection: Option<lancedb::Connection>,
|
||||
server_url: String,
|
||||
database: String,
|
||||
authentication: crate::server_config::Authentication,
|
||||
tls: crate::server_config::TlsOptions,
|
||||
}
|
||||
|
||||
impl SdkSqlExecutor {
|
||||
pub fn new_with_authentication(
|
||||
server_url: String,
|
||||
authentication: crate::server_config::Authentication,
|
||||
) -> Self {
|
||||
Self {
|
||||
connection: None,
|
||||
server_url,
|
||||
database: "test".into(),
|
||||
authentication,
|
||||
tls: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_with_api_key(server_url: String, api_key: String) -> Self {
|
||||
Self::new_with_authentication(
|
||||
server_url,
|
||||
crate::server_config::Authentication::ApiKey(api_key),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn with_tls(mut self, tls: crate::server_config::TlsOptions) -> Self {
|
||||
self.tls = tls;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SqlQueryExecutor for SdkSqlExecutor {
|
||||
async fn execute_query(&mut self, query: &str) -> Result<SqlResult> {
|
||||
use futures::TryStreamExt;
|
||||
let connection = self
|
||||
.connection
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("Flight SQL client not connected"))?;
|
||||
let mut builder = connection.execute_query_async(query);
|
||||
if let Ok(namespace) = std::env::var("LANCEDB_NAMESPACE_PATH") {
|
||||
builder = builder.default_namespace_path(namespace.split('$'));
|
||||
}
|
||||
let query = builder.execute().await?;
|
||||
let batches: Vec<RecordBatch> = query.reader().await?.try_collect().await?;
|
||||
Ok(batches.into())
|
||||
}
|
||||
|
||||
async fn is_connected(&self) -> bool {
|
||||
self.connection.is_some()
|
||||
}
|
||||
|
||||
fn connection_info(&self) -> Option<String> {
|
||||
self.connection
|
||||
.as_ref()
|
||||
.map(|_| format!("Connected to {}", self.server_url))
|
||||
}
|
||||
|
||||
async fn connect(&mut self, server_url: &str) -> Result<()> {
|
||||
use crate::server_config::{parse_grpc_endpoint, Authentication, DEFAULT_SQL_PORT};
|
||||
self.connection = None;
|
||||
self.server_url = server_url.to_owned();
|
||||
let endpoint = parse_grpc_endpoint(server_url, DEFAULT_SQL_PORT)?.url();
|
||||
let mut builder = lancedb::connect(&format!("db://{}", self.database))
|
||||
.region("self-hosted")
|
||||
.sql_host_override(&endpoint);
|
||||
builder = match &self.authentication {
|
||||
Authentication::ApiKey(key) => builder.api_key(key),
|
||||
_ => builder.api_key(""),
|
||||
};
|
||||
let mut config =
|
||||
crate::lancedb_client_config(&self.tls, &self.authentication).unwrap_or_default();
|
||||
// Preserve the CLI's bounded startup probe, including the TLS handshake.
|
||||
config.timeout_config.connect_timeout = Some(std::time::Duration::from_secs(5));
|
||||
builder = builder.client_config(config);
|
||||
let connection = builder.execute().await?;
|
||||
connection.check_sql_connection().await?;
|
||||
self.connection = Some(connection);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_database(&mut self, database: String) {
|
||||
if self.database != database {
|
||||
self.connection = None;
|
||||
}
|
||||
self.database = database;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,380 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
/*
|
||||
MIT License
|
||||
|
||||
Copyright (c) Foxglove Technologies Inc
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
*/
|
||||
|
||||
// Ported from
|
||||
// https://github.com/foxglove/foxglove-cli/blob/main/foxglove/util/tablewriter/tablewriter.go
|
||||
|
||||
use std::io::{self, Write};
|
||||
use std::process::Command;
|
||||
use unicode_width::UnicodeWidthStr;
|
||||
|
||||
pub struct TableWriter;
|
||||
|
||||
impl TableWriter {
|
||||
pub fn print_table<W: Write>(
|
||||
writer: &mut W,
|
||||
headers: &[String],
|
||||
data: &[Vec<String>],
|
||||
term_width: Option<usize>,
|
||||
) -> io::Result<()> {
|
||||
let term_width = term_width.unwrap_or_else(get_terminal_width);
|
||||
let table_width = compute_hotdog_table_width(headers, data);
|
||||
|
||||
if term_width < table_width {
|
||||
Self::print_hamburger(writer, term_width, headers, data)
|
||||
} else {
|
||||
Self::print_hotdog(writer, headers, data)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn print_table_with_title<W: Write>(
|
||||
writer: &mut W,
|
||||
title: &str,
|
||||
headers: &[String],
|
||||
data: &[Vec<String>],
|
||||
term_width: Option<usize>,
|
||||
) -> io::Result<()> {
|
||||
let term_width = term_width.unwrap_or_else(get_terminal_width);
|
||||
let table_width = compute_hotdog_table_width(headers, data);
|
||||
|
||||
if term_width < table_width {
|
||||
// For hamburger mode, just print the title normally
|
||||
writeln!(writer, "{}", title)?;
|
||||
Self::print_hamburger(writer, term_width, headers, data)
|
||||
} else {
|
||||
// For hotdog mode, center the title over the table
|
||||
let title_padding = if table_width > title.width() {
|
||||
(table_width - title.width()) / 2
|
||||
} else {
|
||||
0
|
||||
};
|
||||
writeln!(writer, "{}{}", " ".repeat(title_padding), title)?;
|
||||
Self::print_hotdog(writer, headers, data)
|
||||
}
|
||||
}
|
||||
|
||||
/// Prints table in horizontal "hotdog" format (PostgreSQL style)
|
||||
fn print_hotdog<W: Write>(
|
||||
writer: &mut W,
|
||||
headers: &[String],
|
||||
data: &[Vec<String>],
|
||||
) -> io::Result<()> {
|
||||
let cell_widths = compute_hotdog_cell_widths(headers, data);
|
||||
let column_alignments = compute_column_alignments(data);
|
||||
|
||||
// Write headers
|
||||
write!(writer, " ")?;
|
||||
for (i, header) in headers.iter().enumerate() {
|
||||
if i > 0 {
|
||||
write!(writer, " | ")?;
|
||||
}
|
||||
let padding = cell_widths[i] - header.width();
|
||||
write!(writer, "{}{}", header, " ".repeat(padding))?;
|
||||
}
|
||||
writeln!(writer)?;
|
||||
|
||||
// Write separator with + at junctions
|
||||
write!(writer, "-")?; // Leading edge
|
||||
for (i, &width) in cell_widths.iter().enumerate() {
|
||||
write!(writer, "{}", "-".repeat(width))?;
|
||||
if i < cell_widths.len() - 1 {
|
||||
write!(writer, "-+-")?; // Junction between columns
|
||||
}
|
||||
}
|
||||
writeln!(writer, "-")?; // Trailing edge
|
||||
|
||||
// Write data rows
|
||||
for row in data {
|
||||
write!(writer, " ")?;
|
||||
for (i, col) in row.iter().enumerate() {
|
||||
if i > 0 {
|
||||
write!(writer, " | ")?;
|
||||
}
|
||||
let padding = cell_widths[i] - col.width();
|
||||
|
||||
// Right-align numeric columns, left-align others
|
||||
if i < column_alignments.len() && column_alignments[i] {
|
||||
write!(writer, "{}{}", " ".repeat(padding), col)?;
|
||||
} else {
|
||||
write!(writer, "{}{}", col, " ".repeat(padding))?;
|
||||
}
|
||||
}
|
||||
writeln!(writer)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Prints table in vertical "hamburger" format (psql expanded display style)
|
||||
fn print_hamburger<W: Write>(
|
||||
writer: &mut W,
|
||||
term_width: usize,
|
||||
headers: &[String],
|
||||
data: &[Vec<String>],
|
||||
) -> io::Result<()> {
|
||||
let mut max_header_width = 0;
|
||||
let mut max_record_width = 0;
|
||||
|
||||
// Compute max header width
|
||||
for header in headers {
|
||||
if header.width() > max_header_width {
|
||||
max_header_width = header.width();
|
||||
}
|
||||
}
|
||||
|
||||
// Compute max record width (cap at reasonable limit)
|
||||
const MAX_VALUE_DISPLAY_WIDTH: usize = 80; // Reasonable max for value display
|
||||
for row in data {
|
||||
for col in row {
|
||||
let col_width = col.width().min(MAX_VALUE_DISPLAY_WIDTH);
|
||||
if col_width > max_record_width {
|
||||
max_record_width = col_width;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure sufficient room for record headers
|
||||
let longest_record_header = format!("-[ RECORD {} ]", data.len());
|
||||
if longest_record_header.len() > max_header_width {
|
||||
max_header_width = longest_record_header.len();
|
||||
}
|
||||
|
||||
// Calculate right dash extent - aim for a reasonable width
|
||||
let desired_dash_extent = 60; // Reasonable default width for dashes
|
||||
let mut dashes_right_extent = desired_dash_extent.min(max_record_width + 15);
|
||||
let max_allowed_extent = term_width.saturating_sub(max_header_width + 3); // +3 for "+-" and space
|
||||
if dashes_right_extent > max_allowed_extent {
|
||||
dashes_right_extent = max_allowed_extent;
|
||||
}
|
||||
let right_dashes = "-".repeat(dashes_right_extent);
|
||||
|
||||
// Print records
|
||||
for (i, row) in data.iter().enumerate() {
|
||||
let record_header = format!("-[ RECORD {} ]", i + 1);
|
||||
let header_line = format!(
|
||||
"{}{}+{}",
|
||||
record_header,
|
||||
"-".repeat(max_header_width - record_header.len()),
|
||||
right_dashes
|
||||
);
|
||||
writeln!(writer, "{}", header_line)?;
|
||||
|
||||
for (j, col) in row.iter().enumerate() {
|
||||
writeln!(
|
||||
writer,
|
||||
"{}{}| {}",
|
||||
headers[j],
|
||||
" ".repeat(max_header_width - headers[j].width()),
|
||||
col
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_hotdog_cell_widths(headers: &[String], data: &[Vec<String>]) -> Vec<usize> {
|
||||
let mut cell_widths = vec![0; headers.len()];
|
||||
|
||||
// Initialize with header widths
|
||||
for (i, header) in headers.iter().enumerate() {
|
||||
cell_widths[i] = header.width();
|
||||
}
|
||||
|
||||
// Check data column widths
|
||||
for row in data {
|
||||
for (i, col) in row.iter().enumerate() {
|
||||
if i < cell_widths.len() && cell_widths[i] < col.width() {
|
||||
cell_widths[i] = col.width();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cell_widths
|
||||
}
|
||||
|
||||
fn compute_hotdog_table_width(headers: &[String], data: &[Vec<String>]) -> usize {
|
||||
let cell_widths = compute_hotdog_cell_widths(headers, data);
|
||||
// Account for " | " separators between columns (3 chars each) + leading space + trailing space
|
||||
let separator_width = if headers.len() > 1 {
|
||||
(headers.len() - 1) * 3
|
||||
} else {
|
||||
0
|
||||
};
|
||||
2 + cell_widths.iter().sum::<usize>() + separator_width
|
||||
}
|
||||
|
||||
fn compute_column_alignments(data: &[Vec<String>]) -> Vec<bool> {
|
||||
if data.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let num_cols = data[0].len();
|
||||
let mut is_numeric = vec![true; num_cols];
|
||||
|
||||
// Check each column to see if all non-empty values are numeric
|
||||
for col_idx in 0..num_cols {
|
||||
for row in data {
|
||||
if col_idx < row.len() && !row[col_idx].is_empty() {
|
||||
let value = &row[col_idx];
|
||||
// Check if the value is numeric (integer or float)
|
||||
// Also handle NULL values
|
||||
if value != "NULL" && !is_numeric_value(value) {
|
||||
is_numeric[col_idx] = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
is_numeric
|
||||
}
|
||||
|
||||
fn is_numeric_value(s: &str) -> bool {
|
||||
// Try to parse as integer or float
|
||||
s.parse::<i64>().is_ok() || s.parse::<f64>().is_ok()
|
||||
}
|
||||
|
||||
fn get_terminal_width() -> usize {
|
||||
if let Ok(output) = Command::new("stty").arg("size").output() {
|
||||
if let Ok(output_str) = String::from_utf8(output.stdout) {
|
||||
let parts: Vec<&str> = output_str.split_whitespace().collect();
|
||||
if parts.len() >= 2 {
|
||||
if let Ok(cols) = parts[1].parse::<usize>() {
|
||||
return cols;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: try termsize crate
|
||||
if let Some(size) = termsize::get() {
|
||||
return size.cols as usize;
|
||||
}
|
||||
|
||||
// Final fallback
|
||||
80
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Cursor;
|
||||
|
||||
#[test]
|
||||
fn test_print_hotdog() {
|
||||
let mut output = Cursor::new(Vec::new());
|
||||
let headers = vec!["ID".to_string(), "Name".to_string(), "Age".to_string()];
|
||||
let data = vec![
|
||||
vec!["1".to_string(), "Alice".to_string(), "30".to_string()],
|
||||
vec!["2".to_string(), "Bob".to_string(), "25".to_string()],
|
||||
vec!["3".to_string(), "Charlie".to_string(), "35".to_string()],
|
||||
];
|
||||
|
||||
TableWriter::print_hotdog(&mut output, &headers, &data).unwrap();
|
||||
let result = String::from_utf8(output.into_inner()).unwrap();
|
||||
assert_eq!(result, expected_output());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_unicode_table_padding_uses_terminal_columns() {
|
||||
let headers = vec!["µ".to_string(), "Type".to_string()];
|
||||
let data = vec![
|
||||
vec!["µ".to_string(), "TEXT".to_string()],
|
||||
vec!["漢".to_string(), "TEXT".to_string()],
|
||||
vec!["e\u{301}".to_string(), "TEXT".to_string()],
|
||||
];
|
||||
let mut output = Vec::new();
|
||||
TableWriter::print_table(&mut output, &headers, &data, Some(11)).unwrap();
|
||||
assert_eq!(
|
||||
String::from_utf8(output).unwrap(),
|
||||
" µ | Type\n----+------\n µ | TEXT\n 漢 | TEXT\n e\u{301} | TEXT\n"
|
||||
);
|
||||
}
|
||||
|
||||
fn expected_output() -> String {
|
||||
let mut output = String::new();
|
||||
output.push_str(" ID | Name | Age\n");
|
||||
output.push_str("----+---------+-----\n");
|
||||
output.push_str(" 1 | Alice | 30\n");
|
||||
output.push_str(" 2 | Bob | 25\n");
|
||||
output.push_str(" 3 | Charlie | 35\n");
|
||||
output
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_print_hamburger() {
|
||||
let mut output = Cursor::new(Vec::new());
|
||||
let headers = vec!["ID".to_string(), "Name".to_string(), "Age".to_string()];
|
||||
let data = vec![
|
||||
vec!["1".to_string(), "Alice".to_string(), "30".to_string()],
|
||||
vec!["2".to_string(), "Bob".to_string(), "25".to_string()],
|
||||
vec!["3".to_string(), "Charlie".to_string(), "35".to_string()],
|
||||
];
|
||||
|
||||
// Test with a narrow terminal width
|
||||
TableWriter::print_hamburger(&mut output, 80, &headers, &data).unwrap();
|
||||
let result = String::from_utf8(output.into_inner()).unwrap();
|
||||
assert_eq!(result, expected_hamburger_output());
|
||||
}
|
||||
|
||||
fn expected_hamburger_output() -> String {
|
||||
let mut output = String::new();
|
||||
output.push_str("-[ RECORD 1 ]+----------------------\n");
|
||||
output.push_str("ID | 1\n");
|
||||
output.push_str("Name | Alice\n");
|
||||
output.push_str("Age | 30\n");
|
||||
output.push_str("-[ RECORD 2 ]+----------------------\n");
|
||||
output.push_str("ID | 2\n");
|
||||
output.push_str("Name | Bob\n");
|
||||
output.push_str("Age | 25\n");
|
||||
output.push_str("-[ RECORD 3 ]+----------------------\n");
|
||||
output.push_str("ID | 3\n");
|
||||
output.push_str("Name | Charlie\n");
|
||||
output.push_str("Age | 35\n");
|
||||
output
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_print_table_with_title() {
|
||||
let mut output = Cursor::new(Vec::new());
|
||||
let headers = vec!["ID".to_string(), "Name".to_string()];
|
||||
let data = vec![
|
||||
vec!["1".to_string(), "Alice".to_string()],
|
||||
vec!["2".to_string(), "Bob".to_string()],
|
||||
];
|
||||
|
||||
TableWriter::print_table_with_title(&mut output, "User List", &headers, &data, Some(200))
|
||||
.unwrap();
|
||||
let result = String::from_utf8(output.into_inner()).unwrap();
|
||||
|
||||
// Should center the title and print hotdog format
|
||||
let expected = " User List\n ID | Name \n----+-------\n 1 | Alice\n 2 | Bob \n";
|
||||
assert_eq!(result, expected);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,339 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::process::Command;
|
||||
|
||||
#[test]
|
||||
fn an_unconfigured_default_sql_invocation_prints_usage() {
|
||||
use lancedb_cli::server_config::{Config, Profile};
|
||||
|
||||
for has_http_profile in [false, true] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
if has_http_profile {
|
||||
let mut config = Config::default();
|
||||
config.set_profile(
|
||||
"cloud",
|
||||
Profile {
|
||||
http_server_url: Some("http://api.example.com".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
config.set_default_profile("cloud");
|
||||
config
|
||||
.save_to_path(&home.path().join(".lancedb/config.toml"))
|
||||
.unwrap();
|
||||
}
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.env("HOME", home.path())
|
||||
.env("LANCEDB_API_KEY", "secret-that-must-not-appear")
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
assert!(stdout.contains("Usage:"), "{stdout}");
|
||||
assert!(
|
||||
stdout.contains("No Flight SQL endpoint is configured"),
|
||||
"{stdout}"
|
||||
);
|
||||
assert!(!stdout.contains("secret-that-must-not-appear"));
|
||||
assert!(
|
||||
output.stderr.is_empty(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
if !has_http_profile {
|
||||
assert!(!home.path().join(".lancedb").exists());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_config_diagnostics_do_not_print_credentials() {
|
||||
for (contents, location) in [
|
||||
(
|
||||
"[profiles.cloud]\napi_key = \"sk_test_private\",\n",
|
||||
"line 2, column 28",
|
||||
),
|
||||
(
|
||||
"[profiles.cloud.oidc]\nflow = \"sk_test_private\"\n",
|
||||
"line 2, column 8",
|
||||
),
|
||||
] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let config_dir = home.path().join(".lancedb");
|
||||
std::fs::create_dir(&config_dir).unwrap();
|
||||
std::fs::write(config_dir.join("config.toml"), contents).unwrap();
|
||||
for args in [vec!["profile", "list"], vec!["sql"], vec![]] {
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(args)
|
||||
.env("HOME", home.path())
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(!output.status.success());
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(!stderr.contains("sk_test_private"));
|
||||
assert!(!stderr.contains("api_key ="));
|
||||
assert!(stderr.contains("Failed to parse config file"), "{stderr}");
|
||||
assert!(stderr.contains("invalid TOML or configuration value"));
|
||||
assert!(stderr.contains(location), "{stderr}");
|
||||
assert!(stderr.contains("config.toml"));
|
||||
assert!(output.stdout.is_empty());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_sql_requests_require_a_configured_endpoint() {
|
||||
for args in [
|
||||
vec!["sql"],
|
||||
vec!["--command", "SELECT 1"],
|
||||
vec!["--file", "queries.sql"],
|
||||
vec!["--profile", "missing"],
|
||||
] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(&args)
|
||||
.env("HOME", home.path())
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(!output.status.success(), "{args:?}");
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
assert!(
|
||||
stderr.contains("No Flight SQL endpoint is configured"),
|
||||
"{args:?}: {stderr}"
|
||||
);
|
||||
assert!(stderr.contains("--sql-server-url"), "{stderr}");
|
||||
assert!(stderr.contains("--profile local"), "{stderr}");
|
||||
assert!(!home.path().join(".lancedb").exists());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selecting_local_as_default_persists_its_endpoint() {
|
||||
use lancedb_cli::server_config::{Config, Profile};
|
||||
|
||||
for existing_endpoint in [None, Some("localhost:12345")] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let path = home.path().join(".lancedb/config.toml");
|
||||
if let Some(endpoint) = existing_endpoint {
|
||||
let mut config = Config::default();
|
||||
config.set_profile(
|
||||
"local",
|
||||
Profile {
|
||||
sql_server_url: Some(endpoint.into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
config.save_to_path(&path).unwrap();
|
||||
}
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(["profile", "set", "local"])
|
||||
.env("HOME", home.path())
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let config = Config::load_from_path(&path).unwrap();
|
||||
assert_eq!(config.default_profile.as_deref(), Some("local"));
|
||||
assert_eq!(
|
||||
config
|
||||
.profiles
|
||||
.get("local")
|
||||
.unwrap()
|
||||
.sql_server_url
|
||||
.as_deref(),
|
||||
Some(existing_endpoint.unwrap_or("localhost:10025"))
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn help_lists_commands_without_creating_configuration() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.arg("--help")
|
||||
.env("HOME", home.path())
|
||||
.env("LANCEDB_API_KEY", "secret-that-must-not-appear")
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(output.status.success());
|
||||
let stdout = String::from_utf8(output.stdout).unwrap();
|
||||
assert!(stdout.contains("Usage:"));
|
||||
for command in ["auth", "profile", "sql"] {
|
||||
assert!(stdout.contains(command));
|
||||
}
|
||||
for command in ["acl", "load", "job"] {
|
||||
assert!(!stdout.contains(&format!(" {command} ")), "{stdout}");
|
||||
}
|
||||
assert!(!stdout.contains("secret-that-must-not-appear"));
|
||||
assert!(!home.path().join(".lancedb").exists());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn excluded_commands_and_mixed_sql_options_are_rejected() {
|
||||
for args in [
|
||||
vec!["server"],
|
||||
vec!["management"],
|
||||
vec!["job", "--help"],
|
||||
vec!["job", "list", "--help"],
|
||||
vec!["job", "usage"],
|
||||
vec!["acl"],
|
||||
vec!["acl", "--help"],
|
||||
vec!["acl", "list"],
|
||||
vec!["acl", "add"],
|
||||
vec!["acl", "delete"],
|
||||
vec!["load"],
|
||||
vec!["load", "--help"],
|
||||
vec!["load", "--source", "lance", "--dataset", "data.lance"],
|
||||
vec!["--command", "SELECT 1", "profile", "list"],
|
||||
vec!["--file", "queries.sql", "profile", "list"],
|
||||
vec!["profile", "list", "--command", "SELECT 1"],
|
||||
vec!["profile", "list", "--file", "queries.sql"],
|
||||
] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(&args)
|
||||
.env("HOME", home.path())
|
||||
.output()
|
||||
.unwrap();
|
||||
assert_eq!(output.status.code(), Some(2), "{args:?}");
|
||||
assert!(output.stdout.is_empty(), "{args:?}");
|
||||
assert!(!home.path().join(".lancedb").exists(), "{args:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sql_options_do_not_appear_in_other_subcommand_help() {
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(["profile", "--help"])
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(output.status.success());
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
for option in ["--command", "--file", "--format", "--sql-server-url"] {
|
||||
assert!(!stdout.contains(option), "{stdout}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sql_environment_options_do_not_change_profile_commands() {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let output = Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(["profile", "list"])
|
||||
.env("HOME", home.path())
|
||||
.env("LANCEDB_CA_CERT", "/not-used-by-profile/ca.pem")
|
||||
.env("LANCEDB_CLIENT_CERT", "/not-used-by-profile/client.pem")
|
||||
.env("LANCEDB_CLIENT_KEY", "/not-used-by-profile/key.pem")
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
assert!(stdout.contains("Configured profiles:"), "{stdout}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn profile_updates_preserve_enterprise_data_directory() {
|
||||
use lancedb_cli::server_config::{Config, Profile};
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("config.toml");
|
||||
std::fs::write(
|
||||
&path,
|
||||
"data_dir = '/enterprise/data'\ndefault_profile = 'local'\n",
|
||||
)
|
||||
.unwrap();
|
||||
let mut config = Config::load_from_path(&path).unwrap();
|
||||
config.set_profile(
|
||||
"cloud",
|
||||
Profile {
|
||||
database: Some("analytics".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
config.save_to_path(&path).unwrap();
|
||||
let saved = Config::load_from_path(&path).unwrap();
|
||||
assert_eq!(saved.data_dir.as_deref(), Some("/enterprise/data"));
|
||||
assert_eq!(
|
||||
saved.get_profile(Some("cloud")).database.as_deref(),
|
||||
Some("analytics")
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn api_key_profiles_are_private_under_permissive_umasks_and_after_edits() {
|
||||
use lancedb_cli::server_config::Config;
|
||||
use std::fs;
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
for mask in ["022", "000"] {
|
||||
let home = tempfile::tempdir().unwrap();
|
||||
let config_dir = home.path().join(".lancedb");
|
||||
let path = config_dir.join("config.toml");
|
||||
let run = |command: &str, key: &str| {
|
||||
// Set umask only in the child, leaving other test threads unaffected.
|
||||
let output = Command::new("/bin/sh")
|
||||
.args([
|
||||
"-c",
|
||||
r#"umask "$1"; shift; exec "$@""#,
|
||||
"profile-permissions",
|
||||
mask,
|
||||
env!("CARGO_BIN_EXE_lancedb"),
|
||||
"profile",
|
||||
command,
|
||||
"cloud",
|
||||
"--api-key",
|
||||
key,
|
||||
"--database",
|
||||
"analytics",
|
||||
])
|
||||
.env("HOME", home.path())
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{command} under umask {mask}: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
};
|
||||
let assert_private = || {
|
||||
assert_eq!(
|
||||
fs::metadata(&config_dir).unwrap().permissions().mode() & 0o777,
|
||||
0o700,
|
||||
"config directory under umask {mask}"
|
||||
);
|
||||
assert_eq!(
|
||||
fs::metadata(&path).unwrap().permissions().mode() & 0o777,
|
||||
0o600,
|
||||
"config file under umask {mask}"
|
||||
);
|
||||
};
|
||||
|
||||
run("create", "initial-key");
|
||||
assert_private();
|
||||
|
||||
fs::set_permissions(&config_dir, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
fs::set_permissions(&path, fs::Permissions::from_mode(0o644)).unwrap();
|
||||
run("edit", "replacement-key");
|
||||
assert_private();
|
||||
let saved = Config::load_from_path(&path).unwrap();
|
||||
assert_eq!(
|
||||
saved.get_profile(Some("cloud")).api_key.as_deref(),
|
||||
Some("replacement-key")
|
||||
);
|
||||
assert_eq!(
|
||||
saved.get_profile(Some("cloud")).database.as_deref(),
|
||||
Some("analytics")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,643 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! End-to-end TLS and SQL file-execution coverage for the Flight SQL transport.
|
||||
//!
|
||||
//! A throwaway CA signs a server certificate for `127.0.0.1`, a real Flight
|
||||
//! server serves it over TLS, and the CLI's executor connects with the CA
|
||||
//! supplied through `--ca-cert`. The negative cases prove verification is
|
||||
//! genuinely happening rather than being skipped.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::pin::Pin;
|
||||
|
||||
use arrow_array::{Int32Array, RecordBatch};
|
||||
use arrow_flight::encode::FlightDataEncoderBuilder;
|
||||
use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
|
||||
use arrow_flight::sql::{Any, CommandStatementQuery};
|
||||
use arrow_flight::{
|
||||
Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo,
|
||||
HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use futures::{Stream, TryStreamExt};
|
||||
use lancedb_cli::server_config::TlsOptions;
|
||||
use lancedb_cli::sql_executor::SdkSqlExecutor as FlightSqlExecutor;
|
||||
use lancedb_cli::sql_executor::{SqlQueryExecutor, SqlResult};
|
||||
use prost::Message;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_stream::wrappers::TcpListenerStream;
|
||||
use tonic::transport::{Identity, Server, ServerTlsConfig};
|
||||
use tonic::{Request, Response, Status, Streaming};
|
||||
|
||||
/// The single column the stub server returns, asserted by the tests.
|
||||
const ANSWER: i32 = 42;
|
||||
|
||||
/// tonic's server-side TLS needs a process-level rustls crypto provider; unlike
|
||||
/// the client path it has no built-in fallback.
|
||||
fn init_crypto_provider() {
|
||||
static INIT: std::sync::Once = std::sync::Once::new();
|
||||
INIT.call_once(|| {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
});
|
||||
}
|
||||
|
||||
type BoxStream<T> = Pin<Box<dyn Stream<Item = Result<T, Status>> + Send + 'static>>;
|
||||
|
||||
/// Minimal Flight service: enough for `execute_query` to complete a round trip.
|
||||
/// Everything the CLI does not call is left unimplemented on purpose.
|
||||
#[derive(Clone, Default)]
|
||||
struct StubFlightService {
|
||||
queries: Option<Arc<Mutex<Vec<String>>>>,
|
||||
}
|
||||
|
||||
/// The SDK wraps a Flight SQL command in a ticket's field 2.
|
||||
#[derive(Clone, PartialEq, Message)]
|
||||
struct CommandTicket {
|
||||
#[prost(bytes = "vec", tag = "2")]
|
||||
command: Vec<u8>,
|
||||
}
|
||||
|
||||
impl StubFlightService {
|
||||
fn batch() -> RecordBatch {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"answer",
|
||||
DataType::Int32,
|
||||
false,
|
||||
)]));
|
||||
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![ANSWER]))]).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl FlightService for StubFlightService {
|
||||
type HandshakeStream = BoxStream<HandshakeResponse>;
|
||||
type ListFlightsStream = BoxStream<FlightInfo>;
|
||||
type DoGetStream = BoxStream<FlightData>;
|
||||
type DoPutStream = BoxStream<PutResult>;
|
||||
type DoExchangeStream = BoxStream<FlightData>;
|
||||
type DoActionStream = BoxStream<arrow_flight::Result>;
|
||||
type ListActionsStream = BoxStream<ActionType>;
|
||||
|
||||
async fn get_flight_info(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<FlightInfo>, Status> {
|
||||
let endpoint = FlightEndpoint::new().with_ticket(Ticket::new("stub"));
|
||||
Ok(Response::new(FlightInfo::new().with_endpoint(endpoint)))
|
||||
}
|
||||
|
||||
async fn do_get(
|
||||
&self,
|
||||
request: Request<Ticket>,
|
||||
) -> Result<Response<Self::DoGetStream>, Status> {
|
||||
if let Some(queries) = &self.queries {
|
||||
let ticket = CommandTicket::decode(request.into_inner().ticket).unwrap();
|
||||
let command = Any::decode(ticket.command.as_slice())
|
||||
.unwrap()
|
||||
.unpack::<CommandStatementQuery>()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
queries.lock().unwrap().push(command.query);
|
||||
}
|
||||
let batch = Self::batch();
|
||||
let stream = FlightDataEncoderBuilder::new()
|
||||
.build(futures::stream::once(async move { Ok(batch) }))
|
||||
.map_err(|err| Status::internal(err.to_string()));
|
||||
Ok(Response::new(Box::pin(stream)))
|
||||
}
|
||||
|
||||
async fn handshake(
|
||||
&self,
|
||||
_request: Request<Streaming<HandshakeRequest>>,
|
||||
) -> Result<Response<Self::HandshakeStream>, Status> {
|
||||
Err(Status::unimplemented("handshake"))
|
||||
}
|
||||
|
||||
async fn list_flights(
|
||||
&self,
|
||||
_request: Request<Criteria>,
|
||||
) -> Result<Response<Self::ListFlightsStream>, Status> {
|
||||
Err(Status::unimplemented("list_flights"))
|
||||
}
|
||||
|
||||
async fn poll_flight_info(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<PollInfo>, Status> {
|
||||
let endpoint = FlightEndpoint::new().with_ticket(Ticket::new("stub"));
|
||||
let info = FlightInfo::new()
|
||||
.with_endpoint(endpoint)
|
||||
.try_with_schema(&Self::batch().schema())
|
||||
.unwrap();
|
||||
Ok(Response::new(PollInfo {
|
||||
info: Some(info),
|
||||
flight_descriptor: None,
|
||||
progress: Some(1.0),
|
||||
expiration_time: None,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn get_schema(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<SchemaResult>, Status> {
|
||||
Err(Status::unimplemented("get_schema"))
|
||||
}
|
||||
|
||||
async fn do_put(
|
||||
&self,
|
||||
_request: Request<Streaming<FlightData>>,
|
||||
) -> Result<Response<Self::DoPutStream>, Status> {
|
||||
Err(Status::unimplemented("do_put"))
|
||||
}
|
||||
|
||||
async fn do_exchange(
|
||||
&self,
|
||||
_request: Request<Streaming<FlightData>>,
|
||||
) -> Result<Response<Self::DoExchangeStream>, Status> {
|
||||
Err(Status::unimplemented("do_exchange"))
|
||||
}
|
||||
|
||||
async fn do_action(
|
||||
&self,
|
||||
_request: Request<Action>,
|
||||
) -> Result<Response<Self::DoActionStream>, Status> {
|
||||
Err(Status::unimplemented("do_action"))
|
||||
}
|
||||
|
||||
async fn list_actions(
|
||||
&self,
|
||||
_request: Request<Empty>,
|
||||
) -> Result<Response<Self::ListActionsStream>, Status> {
|
||||
Err(Status::unimplemented("list_actions"))
|
||||
}
|
||||
}
|
||||
|
||||
/// A CA plus the server and client certificates it signed. The server
|
||||
/// certificate is valid for `127.0.0.1`.
|
||||
struct TestCerts {
|
||||
ca_pem: String,
|
||||
server_cert_pem: String,
|
||||
server_key_pem: String,
|
||||
client_cert_pem: String,
|
||||
client_key_pem: String,
|
||||
}
|
||||
|
||||
fn generate_certs() -> TestCerts {
|
||||
use rcgen::{BasicConstraints, CertificateParams, IsCa, KeyPair, PKCS_ECDSA_P256_SHA256};
|
||||
|
||||
let mut ca_params = CertificateParams::new(vec!["Test CA".to_string()]);
|
||||
ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
|
||||
ca_params.key_pair = Some(KeyPair::generate(&PKCS_ECDSA_P256_SHA256).unwrap());
|
||||
let ca_cert = rcgen::Certificate::from_params(ca_params).unwrap();
|
||||
|
||||
// The SANs must cover the address the client dials, or the handshake fails
|
||||
// hostname verification even with the right CA.
|
||||
let mut server_params =
|
||||
CertificateParams::new(vec!["localhost".to_string(), "127.0.0.1".to_string()]);
|
||||
server_params.key_pair = Some(KeyPair::generate(&PKCS_ECDSA_P256_SHA256).unwrap());
|
||||
let server_cert = rcgen::Certificate::from_params(server_params).unwrap();
|
||||
|
||||
let mut client_params = CertificateParams::new(vec!["test-client".to_string()]);
|
||||
client_params.key_pair = Some(KeyPair::generate(&PKCS_ECDSA_P256_SHA256).unwrap());
|
||||
let client_cert = rcgen::Certificate::from_params(client_params).unwrap();
|
||||
|
||||
TestCerts {
|
||||
ca_pem: ca_cert.serialize_pem().unwrap(),
|
||||
server_cert_pem: server_cert.serialize_pem_with_signer(&ca_cert).unwrap(),
|
||||
server_key_pem: server_cert.serialize_private_key_pem(),
|
||||
client_cert_pem: client_cert.serialize_pem_with_signer(&ca_cert).unwrap(),
|
||||
client_key_pem: client_cert.serialize_private_key_pem(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the stub Flight server over TLS on an ephemeral port. The server stops
|
||||
/// when the returned shutdown sender is dropped.
|
||||
async fn start_tls_server(certs: &TestCerts) -> (SocketAddr, tokio::sync::oneshot::Sender<()>) {
|
||||
start_server(certs, false).await
|
||||
}
|
||||
|
||||
/// Same, but the server demands a client certificate signed by the test CA.
|
||||
async fn start_mtls_server(certs: &TestCerts) -> (SocketAddr, tokio::sync::oneshot::Sender<()>) {
|
||||
start_server(certs, true).await
|
||||
}
|
||||
|
||||
async fn start_server(
|
||||
certs: &TestCerts,
|
||||
require_client_cert: bool,
|
||||
) -> (SocketAddr, tokio::sync::oneshot::Sender<()>) {
|
||||
init_crypto_provider();
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
|
||||
let identity = Identity::from_pem(&certs.server_cert_pem, &certs.server_key_pem);
|
||||
let mut tls = ServerTlsConfig::new().identity(identity);
|
||||
if require_client_cert {
|
||||
tls = tls.client_ca_root(tonic::transport::Certificate::from_pem(&certs.ca_pem));
|
||||
}
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.tls_config(tls)
|
||||
.unwrap()
|
||||
.add_service(FlightServiceServer::new(StubFlightService::default()))
|
||||
.serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
(addr, shutdown_tx)
|
||||
}
|
||||
|
||||
fn write_pem(dir: &tempfile::TempDir, name: &str, contents: &str) -> String {
|
||||
let path = dir.path().join(name);
|
||||
std::fs::write(&path, contents).unwrap();
|
||||
path.to_str().unwrap().to_string()
|
||||
}
|
||||
|
||||
/// Exercise argument parsing, file reading, script execution, and SQL transport.
|
||||
/// The stub records SQL without executing it against a database.
|
||||
async fn run_sql_cli(args: &[&str], script: Option<&str>) -> (std::process::Output, Vec<String>) {
|
||||
let queries = Arc::new(Mutex::new(Vec::new()));
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let sql_endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let service = StubFlightService {
|
||||
queries: Some(queries.clone()),
|
||||
};
|
||||
let sql_server = tokio::spawn(async move {
|
||||
Server::builder()
|
||||
.add_service(FlightServiceServer::new(service))
|
||||
.serve_with_incoming(TcpListenerStream::new(listener))
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||||
let http_endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let app =
|
||||
axum::Router::new().fallback(|| async { axum::Json(serde_json::json!({"tables": []})) });
|
||||
let http_server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("script.sql");
|
||||
if let Some(script) = script {
|
||||
std::fs::write(&path, script).unwrap();
|
||||
}
|
||||
let mut config = lancedb_cli::server_config::Config::default();
|
||||
config.set_profile(
|
||||
"test",
|
||||
lancedb_cli::server_config::Profile {
|
||||
sql_server_url: Some(sql_endpoint),
|
||||
http_server_url: Some(http_endpoint),
|
||||
database: Some("test".into()),
|
||||
api_key: Some("test-key".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
config.set_default_profile("test");
|
||||
config
|
||||
.save_to_path(&dir.path().join(".lancedb/config.toml"))
|
||||
.unwrap();
|
||||
let output = tokio::process::Command::new(env!("CARGO_BIN_EXE_lancedb"))
|
||||
.args(args)
|
||||
.current_dir(dir.path())
|
||||
.stdin(std::process::Stdio::null())
|
||||
.env("HOME", dir.path())
|
||||
.env("XDG_DATA_HOME", dir.path().join("data"))
|
||||
.env("XDG_CONFIG_HOME", dir.path().join("config"))
|
||||
.env_remove("LANCEDB_API_KEY")
|
||||
.env_remove("LANCEDB_OIDC_ACCESS_TOKEN")
|
||||
.env_remove("LANCEDB_OIDC_ACCESS_TOKEN_FILE")
|
||||
.env_remove("LANCEDB_CA_CERT")
|
||||
.env_remove("LANCEDB_CLIENT_CERT")
|
||||
.env_remove("LANCEDB_CLIENT_KEY")
|
||||
.output()
|
||||
.await
|
||||
.unwrap();
|
||||
sql_server.abort();
|
||||
http_server.abort();
|
||||
let recorded = queries.lock().unwrap().clone();
|
||||
(output, recorded)
|
||||
}
|
||||
|
||||
async fn run_sql_file(script: &str) -> (std::process::Output, Vec<String>) {
|
||||
run_sql_cli(
|
||||
&["sql", "--file", "script.sql", "--format", "json"],
|
||||
Some(script),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_subcommand_opens_the_sql_shell() {
|
||||
for args in [vec![], vec!["sql"]] {
|
||||
let (output, queries) = run_sql_cli(&args, None).await;
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{args:?}: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
assert!(stdout.contains("LanceDB CLI v"), "{stdout}");
|
||||
assert!(
|
||||
stdout.contains("Type \\? for help or \\q to quit"),
|
||||
"{stdout}"
|
||||
);
|
||||
assert!(queries.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sql_options_work_with_and_without_the_subcommand() {
|
||||
for prefix in [vec![], vec!["sql"]] {
|
||||
let mut args = prefix.clone();
|
||||
args.extend(["--command", "SELECT 42", "--format", "json"]);
|
||||
let (output, queries) = run_sql_cli(&args, None).await;
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{args:?}: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert_eq!(queries, ["SELECT 42"]);
|
||||
let value: serde_json::Value = serde_json::from_slice(&output.stdout).unwrap();
|
||||
assert_eq!(value, serde_json::json!([{"answer": ANSWER}]));
|
||||
|
||||
let mut args = prefix;
|
||||
args.extend(["--file", "script.sql", "--format", "json"]);
|
||||
let (output, queries) = run_sql_cli(&args, Some("SELECT 'semi;colon';")).await;
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{args:?}: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert_eq!(queries, ["SELECT 'semi;colon'"]);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sql_files_preserve_comments_and_quoted_semicolons() {
|
||||
for (script, expected) in [
|
||||
(
|
||||
"SHOW JOBS; CANCEL JOB 'job-123';",
|
||||
vec!["SHOW JOBS", "CANCEL JOB 'job-123'"],
|
||||
),
|
||||
(
|
||||
"SELECT 1 -- documentation; DELETE FROM customers\n;",
|
||||
vec!["SELECT 1 -- documentation; DELETE FROM customers"],
|
||||
),
|
||||
(
|
||||
"SELECT 1 /* documentation; DELETE FROM customers */; SELECT 2;",
|
||||
vec![
|
||||
"SELECT 1 /* documentation; DELETE FROM customers */",
|
||||
"SELECT 2",
|
||||
],
|
||||
),
|
||||
(
|
||||
"/* outer; /* nested; DELETE FROM customers */ still a comment; */ SELECT 3;",
|
||||
vec!["/* outer; /* nested; DELETE FROM customers */ still a comment; */ SELECT 3"],
|
||||
),
|
||||
(
|
||||
"SELECT 1 /*! ; DELETE FROM customers */;",
|
||||
vec!["SELECT 1 /*! ; DELETE FROM customers */"],
|
||||
),
|
||||
(
|
||||
r#"SELECT 'it''s; unchanged', "col;""name", `back;``tick`; SELECT 'next'"#,
|
||||
vec![
|
||||
r#"SELECT 'it''s; unchanged', "col;""name", `back;``tick`"#,
|
||||
"SELECT 'next'",
|
||||
],
|
||||
),
|
||||
(
|
||||
"SELECT $$text; -- still text$$, $tag$more; /* text */$tag$; SELECT 3;",
|
||||
vec![
|
||||
"SELECT $$text; -- still text$$, $tag$more; /* text */$tag$",
|
||||
"SELECT 3",
|
||||
],
|
||||
),
|
||||
(
|
||||
r"SELECT E'it\'s; unchanged'; SELECT 2;",
|
||||
vec![r"SELECT E'it\'s; unchanged'", "SELECT 2"],
|
||||
),
|
||||
(
|
||||
"SELECT 'text; -- still a string /* not a comment */'; SELECT 2;",
|
||||
vec![
|
||||
"SELECT 'text; -- still a string /* not a comment */'",
|
||||
"SELECT 2",
|
||||
],
|
||||
),
|
||||
(
|
||||
"SELECT 'café; 🦀';\r\n-- comment; DELETE FROM customers\r\nSELECT 2; ;",
|
||||
vec![
|
||||
"SELECT 'café; 🦀'",
|
||||
"-- comment; DELETE FROM customers\r\nSELECT 2",
|
||||
],
|
||||
),
|
||||
(
|
||||
"; -- documentation; DELETE FROM customers\n; /* comment; SELECT 2 */; ;",
|
||||
vec![],
|
||||
),
|
||||
] {
|
||||
let (output, queries) = run_sql_file(script).await;
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"script {script:?} failed: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert_eq!(queries, expected, "script: {script:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sql_files_with_unterminated_tokens_submit_no_statements() {
|
||||
for script in [
|
||||
"SELECT 1; SELECT 'unterminated; DELETE FROM customers",
|
||||
"SELECT 1; SELECT \"unterminated; DELETE FROM customers",
|
||||
"SELECT 1; SELECT `unterminated; DELETE FROM customers",
|
||||
"SELECT 1; SELECT $tag$unterminated; DELETE FROM customers",
|
||||
"SELECT 1; /* unterminated; DELETE FROM customers",
|
||||
] {
|
||||
let (output, queries) = run_sql_file(script).await;
|
||||
assert!(!output.status.success(), "script: {script:?}");
|
||||
assert!(
|
||||
String::from_utf8_lossy(&output.stderr).contains("Failed to tokenize SQL script"),
|
||||
"unexpected error: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
assert!(
|
||||
queries.is_empty(),
|
||||
"submitted SQL for {script:?}: {queries:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ca_cert_enables_a_tls_flight_sql_query() {
|
||||
let certs = generate_certs();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let ca_path = write_pem(&dir, "ca.pem", &certs.ca_pem);
|
||||
let (addr, _shutdown) = start_tls_server(&certs).await;
|
||||
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string())
|
||||
.with_tls(TlsOptions {
|
||||
ca_cert: Some(ca_path),
|
||||
..TlsOptions::default()
|
||||
});
|
||||
|
||||
executor
|
||||
.connect(&format!("https://{addr}"))
|
||||
.await
|
||||
.expect("TLS connect with the signing CA should succeed");
|
||||
assert!(executor.is_connected().await);
|
||||
|
||||
let result = executor
|
||||
.execute_query("SELECT 1")
|
||||
.await
|
||||
.expect("query over TLS should succeed");
|
||||
match result {
|
||||
SqlResult::Query(batches) => {
|
||||
let rows: Vec<i32> = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column(0)
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.values()
|
||||
.to_vec()
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(rows, vec![ANSWER]);
|
||||
}
|
||||
other => panic!("expected a query result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tls_connect_fails_without_the_signing_ca() {
|
||||
let certs = generate_certs();
|
||||
let (addr, _shutdown) = start_tls_server(&certs).await;
|
||||
|
||||
// No CA configured: the server's certificate chains to a CA that is not in
|
||||
// the platform trust store, so verification must reject it.
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string());
|
||||
|
||||
assert!(
|
||||
executor.connect(&format!("https://{addr}")).await.is_err(),
|
||||
"an untrusted server certificate must not be accepted"
|
||||
);
|
||||
assert!(!executor.is_connected().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tls_connect_fails_with_an_unrelated_ca() {
|
||||
let certs = generate_certs();
|
||||
// A second, unrelated CA: right file shape, wrong signer.
|
||||
let other = generate_certs();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let unrelated_ca = write_pem(&dir, "other-ca.pem", &other.ca_pem);
|
||||
let (addr, _shutdown) = start_tls_server(&certs).await;
|
||||
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string())
|
||||
.with_tls(TlsOptions {
|
||||
ca_cert: Some(unrelated_ca),
|
||||
..TlsOptions::default()
|
||||
});
|
||||
|
||||
assert!(
|
||||
executor.connect(&format!("https://{addr}")).await.is_err(),
|
||||
"a certificate signed by a different CA must not be accepted"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_missing_ca_cert_file_is_reported() {
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string())
|
||||
.with_tls(TlsOptions {
|
||||
ca_cert: Some("/nonexistent/ca.pem".to_string()),
|
||||
..TlsOptions::default()
|
||||
});
|
||||
|
||||
let err = executor
|
||||
.connect("https://127.0.0.1:1")
|
||||
.await
|
||||
.expect_err("a missing CA file must fail loudly");
|
||||
assert!(
|
||||
err.to_string().contains("/nonexistent/ca.pem"),
|
||||
"error should name the missing file, got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_client_certificate_satisfies_an_mtls_server() {
|
||||
let certs = generate_certs();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let ca_path = write_pem(&dir, "ca.pem", &certs.ca_pem);
|
||||
let cert_path = write_pem(&dir, "client.pem", &certs.client_cert_pem);
|
||||
let key_path = write_pem(&dir, "client-key.pem", &certs.client_key_pem);
|
||||
let (addr, _shutdown) = start_mtls_server(&certs).await;
|
||||
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string())
|
||||
.with_tls(TlsOptions {
|
||||
ca_cert: Some(ca_path),
|
||||
client_cert: Some(cert_path),
|
||||
client_key: Some(key_path),
|
||||
});
|
||||
|
||||
executor
|
||||
.connect(&format!("https://{addr}"))
|
||||
.await
|
||||
.expect("mTLS connect with a CA-signed client certificate should succeed");
|
||||
executor
|
||||
.execute_query("SELECT 1")
|
||||
.await
|
||||
.expect("query over mTLS should succeed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_mtls_server_rejects_a_client_with_no_certificate() {
|
||||
let certs = generate_certs();
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let ca_path = write_pem(&dir, "ca.pem", &certs.ca_pem);
|
||||
let (addr, _shutdown) = start_mtls_server(&certs).await;
|
||||
|
||||
// Trusts the server, but presents no identity of its own.
|
||||
let mut executor = FlightSqlExecutor::new_with_api_key(String::new(), "sk_test".to_string())
|
||||
.with_tls(TlsOptions {
|
||||
ca_cert: Some(ca_path),
|
||||
..TlsOptions::default()
|
||||
});
|
||||
|
||||
// TLS 1.3 sends the client certificate after the client considers the
|
||||
// handshake done, so the rejection can surface at connect or at the first
|
||||
// request. Either way the exchange must not succeed.
|
||||
let rejected = match executor.connect(&format!("https://{addr}")).await {
|
||||
Err(_) => true,
|
||||
Ok(()) => executor.execute_query("SELECT 1").await.is_err(),
|
||||
};
|
||||
assert!(
|
||||
rejected,
|
||||
"an mTLS server must reject a client with no certificate"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_client_certificate_without_its_key_is_rejected_at_resolve() {
|
||||
use lancedb_cli::server_config::ConnectionArgs;
|
||||
|
||||
let args = ConnectionArgs {
|
||||
client_cert: Some("/tmp/client.pem".to_string()),
|
||||
..ConnectionArgs::default()
|
||||
};
|
||||
let err = args
|
||||
.resolve(None)
|
||||
.expect_err("a client certificate without a key must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("--client-key"),
|
||||
"error should name the missing flag, got: {err}"
|
||||
);
|
||||
}
|
||||
@@ -528,6 +528,12 @@ impl Connection {
|
||||
ExecuteQueryAsyncBuilder::new(self.internal.clone(), query.into())
|
||||
}
|
||||
|
||||
/// Establish the remote SQL transport without executing a statement.
|
||||
/// This checks connectivity and TLS; authentication is checked by subsequent requests.
|
||||
pub async fn check_sql_connection(&self) -> Result<()> {
|
||||
self.internal.check_sql_connection().await
|
||||
}
|
||||
|
||||
/// Describe a submitted SQL query by its connection-scoped id.
|
||||
///
|
||||
/// This performs one bounded status poll using state retained by this
|
||||
|
||||
@@ -579,6 +579,12 @@ pub trait Database:
|
||||
async fn resume_job(&self, _job_id: &str) -> Result<ResumeJobStatus> {
|
||||
job_op_not_supported("resume_job")
|
||||
}
|
||||
/// Establish the SQL transport without executing a statement.
|
||||
async fn check_sql_connection(&self) -> Result<()> {
|
||||
Err(crate::Error::NotSupported {
|
||||
message: "SQL is not supported by this database".into(),
|
||||
})
|
||||
}
|
||||
/// Start executing a SQL statement on a LanceDB Distributed database.
|
||||
async fn execute_query_async(
|
||||
&self,
|
||||
|
||||
@@ -92,6 +92,7 @@ impl TerminalResult {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) fn remote(value: Option<Value>, request_id: String) -> Self {
|
||||
Self {
|
||||
value,
|
||||
@@ -242,6 +243,7 @@ impl Job<()> {
|
||||
|
||||
/// A handle whose record the caller has already fetched, so the detail
|
||||
/// accessors answer without a second round trip.
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) fn opened(handle: Box<dyn JobHandle>, description: JobDescription) -> Self {
|
||||
let job = Self::new(handle);
|
||||
{
|
||||
|
||||
@@ -1357,6 +1357,16 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
})
|
||||
}
|
||||
|
||||
async fn check_sql_connection(&self) -> Result<()> {
|
||||
self.sql_client
|
||||
.as_ref()
|
||||
.ok_or_else(|| Error::NotSupported {
|
||||
message: "SQL is unavailable for this remote database client".into(),
|
||||
})?
|
||||
.check_connection()
|
||||
.await
|
||||
}
|
||||
|
||||
async fn execute_query_async(&self, request: ExecuteQueryRequest) -> Result<crate::sql::Query> {
|
||||
let client = self
|
||||
.sql_client
|
||||
|
||||
@@ -170,6 +170,13 @@ impl std::fmt::Debug for SqlClient {
|
||||
}
|
||||
|
||||
impl SqlClient {
|
||||
pub(super) async fn check_connection(&self) -> Result<()> {
|
||||
self.inner
|
||||
.connection(&Uuid::new_v4().to_string())
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
pub(super) fn new(
|
||||
database: String,
|
||||
database_prefix: Option<String>,
|
||||
@@ -1428,7 +1435,18 @@ async fn connect_channel(
|
||||
tokio::time::timeout(connect_timeout, endpoint.connect())
|
||||
.await
|
||||
.map_err(|_| sql_error(request_id, "SQL connection timed out"))?
|
||||
.map_err(|err| sql_error(request_id, err))
|
||||
.map_err(|err| {
|
||||
// tonic's Display only says "transport error"; preserve the source chain
|
||||
// so callers can distinguish DNS, refused connections, and TLS failures.
|
||||
let mut message = err.to_string();
|
||||
let mut source = std::error::Error::source(&err);
|
||||
while let Some(error) = source {
|
||||
message.push_str(": ");
|
||||
message.push_str(&error.to_string());
|
||||
source = error.source();
|
||||
}
|
||||
sql_error(request_id, message)
|
||||
})
|
||||
}
|
||||
|
||||
fn tls_config(config: Option<&TlsConfig>) -> Result<ClientTlsConfig> {
|
||||
|
||||
Reference in New Issue
Block a user