feat: connection-level job operations (#3755)

Adds job operations to the connection surface, building on the Job
handle from #3742: job(id), list_jobs, get_job, cancel_job, and
job_history, plus a non-blocking Job.status(). Implemented on the
Database trait (defaulting to NotSupported), the remote backend
(/v1/jobs), and the Python and Node bindings; job_history returns Arrow
batches.

errors() and progress() are not included.

Tested with mocked endpoints in all three languages.

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Wyatt Alt
2026-07-31 12:51:43 -07:00
committed by GitHub
parent a6418b6cb9
commit e3b472c212
24 changed files with 1569 additions and 14 deletions
+97
View File
@@ -25,6 +25,27 @@ the underlying connection has been closed.
## Methods
### cancelJob()
```ts
abstract cancelJob(jobId): Promise<boolean>
```
Request cancellation of a server-side job by id.
Resolves to true if the server accepted the cancellation, false if no
such job exists. Cancelling an already-terminal job is a no-op success.
#### Parameters
* **jobId**: `string`
#### Returns
`Promise`&lt;`boolean`&gt;
***
### cloneTable()
```ts
@@ -365,6 +386,26 @@ Drop an existing table.
***
### getJob()
```ts
abstract getJob(jobId): Promise<null | JobDescription>
```
Describe a single server-side job by id.
Resolves to `null` when the server has no such job.
#### Parameters
* **jobId**: `string`
#### Returns
`Promise`&lt;`null` \| [`JobDescription`](../interfaces/JobDescription.md)&gt;
***
### isOpen()
```ts
@@ -379,6 +420,62 @@ Return true if the connection has not been closed
***
### job()
```ts
abstract job(jobId): Job
```
A [Job](Job.md) handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect on
the job itself.
#### Parameters
* **jobId**: `string`
#### Returns
[`Job`](Job.md)
***
### jobHistory()
```ts
abstract jobHistory(jobId?): Promise<Table<any>>
```
The lifecycle event history of a server-side job, as an Arrow table.
Lists history across all jobs when `jobId` is omitted.
#### Parameters
* **jobId?**: `string`
#### Returns
`Promise`&lt;`Table`&lt;`any`&gt;&gt;
***
### listJobs()
```ts
abstract listJobs(): Promise<JobInfo[]>
```
List server-side jobs across the database's tables.
#### Returns
`Promise`&lt;[`JobInfo`](../interfaces/JobInfo.md)[]&gt;
***
### listNamespaces()
```ts
+19
View File
@@ -51,6 +51,25 @@ Request cancellation. Cancelling a finished operation is a no-op.
***
### status()
```ts
status(): Promise<string>
```
The operation's current lifecycle state: "running", "finished",
"failed", or "cancelled".
A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject
on a terminal failure state. States a newer server reports that this
client version does not know pass through as-is.
#### Returns
`Promise`&lt;`string`&gt;
***
### wait()
```ts
+3
View File
@@ -89,6 +89,9 @@
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
- [IvfPqOptions](interfaces/IvfPqOptions.md)
- [IvfRqOptions](interfaces/IvfRqOptions.md)
- [JobDescription](interfaces/JobDescription.md)
- [JobFailureInfo](interfaces/JobFailureInfo.md)
- [JobInfo](interfaces/JobInfo.md)
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
+66
View File
@@ -0,0 +1,66 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / JobDescription
# Interface: JobDescription
A described job from `Connection.getJob`.
## Properties
### creationMs
```ts
creationMs: number;
```
When the job was created, in milliseconds since the epoch.
***
### failure?
```ts
optional failure: JobFailureInfo;
```
Why the job failed, when the job is failed and the server reports a
reason.
***
### jobId
```ts
jobId: string;
```
***
### jobType
```ts
jobType: string;
```
***
### specJson?
```ts
optional specJson: string;
```
The job-type-specific specification as a JSON string, when present.
***
### state
```ts
state: string;
```
Lifecycle state: "running", "finished", "failed", or "cancelled".
+33
View File
@@ -0,0 +1,33 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / JobFailureInfo
# Interface: JobFailureInfo
The server's account of why a job failed.
## Properties
### message?
```ts
optional message: string;
```
***
### phase?
```ts
optional phase: string;
```
***
### retryable?
```ts
optional retryable: boolean;
```
+58
View File
@@ -0,0 +1,58 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / JobInfo
# Interface: JobInfo
A row from `Connection.listJobs`: one server-side job.
## Properties
### createdAtMillis
```ts
createdAtMillis: number;
```
When the job was created, in milliseconds since the epoch.
***
### jobId
```ts
jobId: string;
```
The job id -- what `Connection.getJob` and `Connection.cancelJob`
accept.
***
### jobType
```ts
jobType: string;
```
***
### state
```ts
state: string;
```
Lifecycle state: "running", "finished", "failed", or "cancelled".
***
### table
```ts
table: string;
```
The table the job runs against, without URI or namespace.
+93
View File
@@ -877,3 +877,96 @@ describe("remote connection", () => {
});
});
});
describe("remote connection jobs surface", () => {
it("lists, describes, cancels, and reads history", async () => {
const { tableFromArrays, tableToIPC } = await import("apache-arrow");
const eventsTable = tableFromArrays({ state: ["created", "succeeded"] });
const eventsBody = Buffer.from(tableToIPC(eventsTable, "stream"));
await withMockDatabase(
(req, res) => {
let body = "";
req.on("data", (chunk) => {
body += chunk;
});
req.on("end", () => {
const payload = body.length > 0 ? JSON.parse(body) : {};
if (req.url === "/v1/jobs/list") {
if (payload["page_token"] === undefined) {
res
.writeHead(200, { "Content-Type": "application/json" })
.end(
'{"jobs": [{"job_id": "job-1", "table": "t1", ' +
'"job_type": "create_index", "state": "in_progress", ' +
'"created_at_millis": 1000}], "page_token": "next"}',
);
} else {
res
.writeHead(200, { "Content-Type": "application/json" })
.end(
'{"jobs": [{"job_id": "job-2", "table": "t2", ' +
'"job_type": "create_index", "state": "succeeded", ' +
'"created_at_millis": 2000}]}',
);
}
} else if (req.url === "/v1/jobs/describe") {
if (payload["job_id"] !== "job-1") {
res.writeHead(404).end("no such job");
return;
}
res
.writeHead(200, { "Content-Type": "application/json" })
.end(
'{"job_id": "job-1", "job_type": "create_index", ' +
'"job_state": "FAILED", "creation_ms": 1000, ' +
'"spec": {"column": "vec"}, "failure": {"phase": "execute", ' +
'"message": "worker died", "retryable": true}}',
);
} else if (req.url === "/v1/jobs/cancel") {
if (payload["job_id"] !== "job-1") {
res.writeHead(404).end("no such job");
return;
}
res
.writeHead(200, { "Content-Type": "application/json" })
.end('{"job_id": "job-1"}');
} else if (req.url === "/v1/jobs/query_events") {
res
.writeHead(200, {
"Content-Type": "application/vnd.apache.arrow.stream",
})
.end(eventsBody);
} else {
res.writeHead(404).end();
}
});
},
async (db) => {
const jobs = await db.listJobs();
expect(jobs.map((job) => job.jobId)).toEqual(["job-1", "job-2"]);
expect(jobs[0].state).toEqual("running");
expect(jobs[1].state).toEqual("finished");
const description = await db.getJob("job-1");
expect(description?.state).toEqual("failed");
expect(JSON.parse(description?.specJson ?? "")).toEqual({
column: "vec",
});
expect(description?.failure?.message).toEqual("worker died");
expect(await db.getJob("missing")).toBeNull();
expect(await db.cancelJob("job-1")).toBe(true);
expect(await db.cancelJob("missing")).toBe(false);
const history = await db.jobHistory("job-1");
expect(history.numRows).toEqual(2);
const job = db.job("job-1");
expect(job.id).toEqual("job-1");
expect(await job.status()).toEqual("failed");
await expect(job.wait()).rejects.toThrow("worker died");
},
);
});
});
+62
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import { tableFromIPC } from "apache-arrow";
import {
Data,
SchemaLike,
@@ -20,6 +21,9 @@ import type {
CreateNamespaceResponse,
DescribeNamespaceResponse,
DropNamespaceResponse,
Job,
JobDescription,
JobInfo,
ListNamespacesResponse,
} from "./native";
export type {
@@ -436,6 +440,40 @@ export abstract class Connection {
newName: string,
options?: RenameTableOptions,
): Promise<void>;
/**
* A {@link Job} handle for a server-side job by id.
*
* The handle is constructed without a server round trip; an unknown id
* surfaces when the handle is used. Dropping the handle has no effect on
* the job itself.
*/
abstract job(jobId: string): Job;
/** List server-side jobs across the database's tables. */
abstract listJobs(): Promise<JobInfo[]>;
/**
* Describe a single server-side job by id.
*
* Resolves to `null` when the server has no such job.
*/
abstract getJob(jobId: string): Promise<JobDescription | null>;
/**
* Request cancellation of a server-side job by id.
*
* Resolves to true if the server accepted the cancellation, false if no
* such job exists. Cancelling an already-terminal job is a no-op success.
*/
abstract cancelJob(jobId: string): Promise<boolean>;
/**
* The lifecycle event history of a server-side job, as an Arrow table.
*
* Lists history across all jobs when `jobId` is omitted.
*/
abstract jobHistory(jobId?: string): Promise<ArrowTable>;
}
/** @hideconstructor */
@@ -722,6 +760,30 @@ export class LocalConnection extends Connection {
options?.newNamespacePath,
);
}
job(jobId: string): Job {
return this.inner.job(jobId);
}
async listJobs(): Promise<JobInfo[]> {
return this.inner.listJobs();
}
async getJob(jobId: string): Promise<JobDescription | null> {
return this.inner.getJob(jobId);
}
async cancelJob(jobId: string): Promise<boolean> {
return this.inner.cancelJob(jobId);
}
async jobHistory(jobId?: string): Promise<ArrowTable> {
const buf = await this.inner.jobHistory(jobId);
if (buf.length === 0) {
return new ArrowTable();
}
return tableFromIPC(buf);
}
}
/**
+7 -1
View File
@@ -85,7 +85,13 @@ export {
RenameTableOptions,
} from "./connection";
export { Job, Session } from "./native.js";
export {
Job,
JobDescription,
JobFailureInfo,
JobInfo,
Session,
} from "./native.js";
export {
ExecutableQuery,
+63
View File
@@ -340,6 +340,69 @@ impl Connection {
self.get_inner()?.drop_all_tables(&ns).await.default_error()
}
/// A `Job` handle for a server-side job by id.
///
/// The handle is constructed without a server round trip; an unknown id
/// surfaces when the handle is used.
#[napi]
pub fn job(&self, job_id: String) -> napi::Result<crate::job::Job> {
let job = self.get_inner()?.job(job_id).default_error()?;
Ok(crate::job::Job::new(job))
}
/// List server-side jobs across the database's tables.
#[napi(catch_unwind)]
pub async fn list_jobs(&self) -> napi::Result<Vec<crate::job::JobInfo>> {
let jobs = self.get_inner()?.list_jobs().await.default_error()?;
Ok(jobs.into_iter().map(Into::into).collect())
}
/// Describe a single server-side job by id. `null` when the server has
/// no such job.
#[napi(catch_unwind)]
pub async fn get_job(
&self,
job_id: String,
) -> napi::Result<Option<crate::job::JobDescription>> {
let description = self.get_inner()?.get_job(&job_id).await.default_error()?;
Ok(description.map(Into::into))
}
/// Request cancellation of a server-side job by id. Returns true if the
/// server accepted the cancellation, false if no such job exists.
#[napi(catch_unwind)]
pub async fn cancel_job(&self, job_id: String) -> napi::Result<bool> {
self.get_inner()?.cancel_job(&job_id).await.default_error()
}
/// The lifecycle event history of a server-side job (all jobs when
/// `job_id` is null), as an Arrow IPC stream buffer. Empty when there is
/// no history.
#[napi(catch_unwind)]
pub async fn job_history(&self, job_id: Option<String>) -> napi::Result<Buffer> {
let batches = self
.get_inner()?
.job_history(job_id.as_deref())
.await
.default_error()?;
let Some(first) = batches.first() else {
return Ok(Buffer::from(Vec::<u8>::new()));
};
let mut out = Vec::new();
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut out, &first.schema())
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
for batch in &batches {
writer
.write(batch)
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
}
writer
.finish()
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
drop(writer);
Ok(Buffer::from(out))
}
#[napi(catch_unwind)]
/// Describe a namespace and return its properties.
pub async fn describe_namespace(
+79
View File
@@ -30,6 +30,17 @@ impl Job {
self.inner.id().map(str::to_string)
}
/// The operation's current lifecycle state: "running", "finished",
/// "failed", or "cancelled".
///
/// A point snapshot; unlike {@link Job.wait} it does not block or reject
/// on a terminal failure state. States a newer server reports that this
/// client version does not know pass through as-is.
#[napi(catch_unwind)]
pub async fn status(&self) -> napi::Result<String> {
self.inner.status().await.default_error()
}
/// Wait until the operation reaches a terminal state.
#[napi(catch_unwind)]
pub async fn wait(&self) -> napi::Result<()> {
@@ -42,3 +53,71 @@ impl Job {
self.inner.cancel().await.default_error()
}
}
/// A row from `Connection.listJobs`: one server-side job.
#[napi(object)]
pub struct JobInfo {
/// The job id -- what `Connection.getJob` and `Connection.cancelJob`
/// accept.
pub job_id: String,
/// The table the job runs against, without URI or namespace.
pub table: String,
pub job_type: String,
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
pub state: String,
/// When the job was created, in milliseconds since the epoch.
pub created_at_millis: i64,
}
impl From<lancedb::database::JobInfo> for JobInfo {
fn from(info: lancedb::database::JobInfo) -> Self {
Self {
job_id: info.job_id,
table: info.table,
job_type: info.job_type,
state: info.state,
created_at_millis: info.created_at_millis,
}
}
}
/// The server's account of why a job failed.
#[napi(object)]
pub struct JobFailureInfo {
pub phase: Option<String>,
pub message: Option<String>,
pub retryable: Option<bool>,
}
/// A described job from `Connection.getJob`.
#[napi(object)]
pub struct JobDescription {
pub job_id: String,
pub job_type: String,
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
pub state: String,
/// When the job was created, in milliseconds since the epoch.
pub creation_ms: i64,
/// The job-type-specific specification as a JSON string, when present.
pub spec_json: Option<String>,
/// Why the job failed, when the job is failed and the server reports a
/// reason.
pub failure: Option<JobFailureInfo>,
}
impl From<lancedb::database::JobDescription> for JobDescription {
fn from(description: lancedb::database::JobDescription) -> Self {
Self {
job_id: description.job_id,
job_type: description.job_type,
state: description.state,
creation_ms: description.creation_ms,
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
failure: description.failure.map(|failure| JobFailureInfo {
phase: failure.phase,
message: failure.message,
retryable: failure.retryable,
}),
}
}
}
+42
View File
@@ -146,6 +146,13 @@ class Connection(object):
start_after: Optional[str],
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
def job(self, job_id: str) -> Job: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
async def cancel_job(self, job_id: str) -> bool: ...
async def job_history(
self, job_id: Optional[str] = None
) -> List[pa.RecordBatch]: ...
async def create_table(
self,
name: str,
@@ -212,9 +219,44 @@ class BlobFile:
class Job:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> None: ...
async def cancel(self) -> None: ...
class JobInfo:
@property
def job_id(self) -> str: ...
@property
def table(self) -> str: ...
@property
def job_type(self) -> str: ...
@property
def state(self) -> str: ...
@property
def created_at_millis(self) -> int: ...
class JobFailureInfo:
@property
def phase(self) -> Optional[str]: ...
@property
def message(self) -> Optional[str]: ...
@property
def retryable(self) -> Optional[bool]: ...
class JobDescription:
@property
def job_id(self) -> str: ...
@property
def job_type(self) -> str: ...
@property
def state(self) -> str: ...
@property
def creation_ms(self) -> int: ...
@property
def spec_json(self) -> Optional[str]: ...
@property
def failure(self) -> Optional[JobFailureInfo]: ...
class Table:
def name(self) -> str: ...
def __repr__(self) -> str: ...
+120
View File
@@ -45,6 +45,7 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from . import __version__
from ._lancedb import connect as lancedb_connect # type: ignore
from .job import AsyncJob, Job
from .table import (
AsyncTable,
LanceTable,
@@ -63,6 +64,7 @@ if TYPE_CHECKING:
from .pydantic import LanceModel
from ._lancedb import Connection as LanceDbConnection
from ._lancedb import JobDescription, JobInfo
from .common import DATA, URI
from .embeddings import EmbeddingFunctionConfig
from ._lancedb import Session
@@ -608,6 +610,46 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError("serialize is not supported for this connection type")
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
"""
raise NotImplementedError("job is not supported for this connection type")
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
raise NotImplementedError("list_jobs is not supported for this connection type")
def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
raise NotImplementedError("get_job is not supported for this connection type")
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
Returns True if the server accepted the cancellation, False if no
such job exists. Cancelling an already-terminal job is a no-op
success.
"""
raise NotImplementedError(
"cancel_job is not supported for this connection type"
)
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
Lists history across all jobs when `job_id` is None.
"""
raise NotImplementedError(
"job_history is not supported for this connection type"
)
class LanceDBConnection(DBConnection):
"""
@@ -1170,6 +1212,47 @@ class LanceDBConnection(DBConnection):
)
)
@override
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
"""
return Job(self._conn.job(job_id))
@override
def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return LOOP.run(self._conn.list_jobs())
@override
def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return LOOP.run(self._conn.get_job(job_id))
@override
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
Returns True if the server accepted the cancellation, False if no
such job exists. Cancelling an already-terminal job is a no-op
success.
"""
return LOOP.run(self._conn.cancel_job(job_id))
@override
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
Lists history across all jobs when `job_id` is None.
"""
return LOOP.run(self._conn.job_history(job_id))
@override
def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
@@ -1879,6 +1962,43 @@ class AsyncConnection(object):
namespace_path = []
await self._inner.drop_all_tables(namespace_path=namespace_path)
def job(self, job_id: str) -> AsyncJob:
"""An [AsyncJob][lancedb.job.AsyncJob] handle for a server-side job
by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
"""
return AsyncJob(self._inner.job(job_id))
async def list_jobs(self) -> List[JobInfo]:
"""List server-side jobs across the database's tables."""
return await self._inner.list_jobs()
async def get_job(self, job_id: str) -> Optional[JobDescription]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return await self._inner.get_job(job_id)
async def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
Returns True if the server accepted the cancellation, False if no
such job exists. Cancelling an already-terminal job is a no-op
success.
"""
return await self._inner.cancel_job(job_id)
async def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
Lists history across all jobs when `job_id` is None.
"""
return await self._inner.job_history(job_id)
async def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
+22
View File
@@ -32,6 +32,18 @@ class AsyncJob:
"""
return self._inner.id if self._inner is not None else None
async def status(self) -> str:
"""The operation's current lifecycle state: "running", "finished",
"failed", or "cancelled".
A point snapshot; unlike `wait` it does not block or raise on a
terminal failure state. States a newer server reports that this
client version does not know pass through as-is.
"""
if self._inner is None:
return "finished"
return await self._inner.status()
async def wait(self, timeout: Optional[timedelta] = None):
"""Wait until the operation reaches a terminal state.
@@ -66,6 +78,16 @@ class Job:
"""
return self._inner.id if self._inner is not None else None
def status(self) -> str:
"""The operation's current lifecycle state: "running", "finished",
"failed", or "cancelled".
See :meth:`AsyncJob.status`.
"""
if self._inner is None:
return "finished"
return LOOP.run(self._inner.status())
def wait(self, timeout: Optional[timedelta] = None):
"""Block until the operation reaches a terminal state.
+46 -1
View File
@@ -7,7 +7,7 @@ import json
import logging
from concurrent.futures import ThreadPoolExecutor
import sys
from typing import Any, Dict, Iterable, List, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
from urllib.parse import urlparse
import warnings
@@ -23,6 +23,10 @@ import pyarrow as pa
from ..common import DATA
from ..db import DBConnection, LOOP
from ..job import Job
if TYPE_CHECKING:
from .._lancedb import JobDescription, JobInfo
from ..embeddings import EmbeddingFunctionConfig
from lance_namespace import (
LanceNamespace,
@@ -689,6 +693,47 @@ class RemoteDBConnection(DBConnection):
)
)
@override
def job(self, job_id: str) -> Job:
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
The handle is constructed without a server round trip; an unknown id
surfaces when the handle is used. Dropping the handle has no effect
on the job itself.
"""
return Job(self._conn.job(job_id))
@override
def list_jobs(self) -> List["JobInfo"]:
"""List server-side jobs across the database's tables."""
return LOOP.run(self._conn.list_jobs())
@override
def get_job(self, job_id: str) -> Optional["JobDescription"]:
"""Describe a single server-side job by id.
Returns None when the server has no such job.
"""
return LOOP.run(self._conn.get_job(job_id))
@override
def cancel_job(self, job_id: str) -> bool:
"""Request cancellation of a server-side job by id.
Returns True if the server accepted the cancellation, False if no
such job exists. Cancelling an already-terminal job is a no-op
success.
"""
return LOOP.run(self._conn.cancel_job(job_id))
@override
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
"""The lifecycle event history of a server-side job, as Arrow batches.
Lists history across all jobs when `job_id` is None.
"""
return LOOP.run(self._conn.job_history(job_id))
@override
def namespace_client(self) -> LanceNamespace:
"""Get the equivalent namespace client for this connection.
+116
View File
@@ -2153,3 +2153,119 @@ def test_remote_blob_byte_apis_not_supported_on_old_server():
table.fetch_blobs("image", [1])
with pytest.raises(NotImplementedError, match="not supported"):
table.fetch_blob_files("image", [1])
def test_remote_connection_jobs_surface():
from lancedb.exceptions import JobFailedError
schema = pa.schema([("state", pa.string())])
batch = pa.record_batch([pa.array(["created", "done"])], schema=schema)
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, schema) as writer:
writer.write_batch(batch)
events_body = sink.getvalue().to_pybytes()
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
payload = json.loads(body) if body else {}
if request.path == "/v1/jobs/list":
if payload.get("page_token") is None:
rsp = dict(
jobs=[
dict(
job_id="job-1",
table="t1",
job_type="create_index",
state="in_progress",
created_at_millis=1000,
)
],
page_token="next",
)
else:
assert payload["page_token"] == "next"
rsp = dict(
jobs=[
dict(
job_id="job-2",
table="t2",
job_type="create_index",
state="succeeded",
created_at_millis=2000,
)
]
)
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(rsp).encode())
elif request.path == "/v1/jobs/describe":
if payload["job_id"] != "job-1":
request.send_response(404)
request.end_headers()
return
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
dict(
job_id="job-1",
job_type="create_index",
job_state="FAILED",
creation_ms=1000,
spec=dict(column="vec"),
failure=dict(
phase="execute", message="worker died", retryable=True
),
)
).encode()
)
elif request.path == "/v1/jobs/cancel":
if payload["job_id"] != "job-1":
request.send_response(404)
request.end_headers()
return
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"job_id": "job-1"}')
elif request.path == "/v1/jobs/query_events":
assert payload["job_id"] == "job-1"
request.send_response(200)
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
request.end_headers()
request.wfile.write(events_body)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
jobs = db.list_jobs()
assert [job.job_id for job in jobs] == ["job-1", "job-2"]
assert jobs[0].state == "running"
assert jobs[0].table == "t1"
assert jobs[1].state == "finished"
description = db.get_job("job-1")
assert description.job_type == "create_index"
assert description.state == "failed"
assert json.loads(description.spec_json) == {"column": "vec"}
assert description.failure.message == "worker died"
assert description.failure.retryable is True
assert db.get_job("missing") is None
assert db.cancel_job("job-1") is True
assert db.cancel_job("missing") is False
batches = db.job_history("job-1")
assert len(batches) == 1
assert batches[0].num_rows == 2
assert batches[0].column("state").to_pylist() == ["created", "done"]
job = db.job("job-1")
assert job.id == "job-1"
assert job.status() == "failed"
with pytest.raises(JobFailedError, match="worker died"):
job.wait(timeout=timedelta(seconds=5))
+55 -2
View File
@@ -13,7 +13,11 @@ use crate::{
runtime::future_into_py,
table::Table,
};
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
use arrow::{
datatypes::Schema,
ffi_stream::ArrowArrayStreamReader,
pyarrow::{FromPyArrow, ToPyArrow},
};
use lancedb::{
connection::Connection as LanceConnection,
connection::NamespaceClientPushdownOperation,
@@ -24,7 +28,7 @@ use pyo3::{
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
exceptions::{PyRuntimeError, PyValueError},
pyclass, pyfunction, pymethods,
types::{PyDict, PyDictMethods},
types::{PyDict, PyDictMethods, PyList, PyListMethods},
};
#[pyclass]
@@ -536,6 +540,55 @@ impl Connection {
})
})
}
pub fn job(&self, job_id: String) -> PyResult<crate::job::Job> {
let inner = self.get_inner()?.clone();
Ok(crate::job::Job::new(inner.job(job_id).infer_error()?))
}
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let jobs = inner.list_jobs().await.infer_error()?;
Ok(jobs
.into_iter()
.map(crate::job::JobInfo::from)
.collect::<Vec<_>>())
})
}
pub fn get_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let description = inner.get_job(&job_id).await.infer_error()?;
Ok(description.map(crate::job::JobDescription::from))
})
}
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
inner.cancel_job(&job_id).await.infer_error()
})
}
#[pyo3(signature = (job_id=None))]
pub fn job_history(
self_: PyRef<'_, Self>,
job_id: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let batches = inner.job_history(job_id.as_deref()).await.infer_error()?;
Python::attach(|py| {
let list = PyList::empty(py);
for batch in batches {
list.append(batch.to_pyarrow(py)?)?;
}
Ok(list.unbind())
})
})
}
}
#[pyfunction]
+99
View File
@@ -28,6 +28,14 @@ impl Job {
self.inner.id().map(str::to_string)
}
pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(
self_.py(),
async move { inner.status().await.infer_error() },
)
}
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
@@ -44,3 +52,94 @@ impl Job {
})
}
}
/// A row from `Connection.list_jobs`: one server-side job.
#[pyclass(get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobInfo {
job_id: String,
table: String,
job_type: String,
state: String,
created_at_millis: i64,
}
#[pymethods]
impl JobInfo {
fn __repr__(&self) -> String {
format!(
"JobInfo(job_id={:?}, table={:?}, job_type={:?}, state={:?}, created_at_millis={})",
self.job_id, self.table, self.job_type, self.state, self.created_at_millis
)
}
}
impl From<lancedb::database::JobInfo> for JobInfo {
fn from(info: lancedb::database::JobInfo) -> Self {
Self {
job_id: info.job_id,
table: info.table,
job_type: info.job_type,
state: info.state,
created_at_millis: info.created_at_millis,
}
}
}
/// The server's account of why a job failed.
#[pyclass(get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobFailureInfo {
phase: Option<String>,
message: Option<String>,
retryable: Option<bool>,
}
#[pymethods]
impl JobFailureInfo {
fn __repr__(&self) -> String {
format!(
"JobFailureInfo(phase={:?}, message={:?}, retryable={:?})",
self.phase, self.message, self.retryable
)
}
}
/// A described job from `Connection.get_job`.
#[pyclass(get_all, skip_from_py_object)]
#[derive(Clone)]
pub struct JobDescription {
job_id: String,
job_type: String,
state: String,
creation_ms: i64,
spec_json: Option<String>,
failure: Option<JobFailureInfo>,
}
#[pymethods]
impl JobDescription {
fn __repr__(&self) -> String {
format!(
"JobDescription(job_id={:?}, job_type={:?}, state={:?}, creation_ms={})",
self.job_id, self.job_type, self.state, self.creation_ms
)
}
}
impl From<lancedb::database::JobDescription> for JobDescription {
fn from(description: lancedb::database::JobDescription) -> Self {
Self {
job_id: description.job_id,
job_type: description.job_type,
state: description.state,
creation_ms: description.creation_ms,
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
failure: description.failure.map(|failure| JobFailureInfo {
phase: failure.phase,
message: failure.message,
retryable: failure.retryable,
}),
}
}
}
+3
View File
@@ -46,6 +46,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Session>()?;
m.add_class::<Table>()?;
m.add_class::<crate::job::Job>()?;
m.add_class::<crate::job::JobInfo>()?;
m.add_class::<crate::job::JobDescription>()?;
m.add_class::<crate::job::JobFailureInfo>()?;
m.add_class::<PyBlobFile>()?;
m.add_class::<IndexConfig>()?;
m.add_class::<Query>()?;
+35 -2
View File
@@ -23,8 +23,8 @@ use crate::connection::create_table::CreateTableBuilder;
use crate::data::scannable::Scannable;
use crate::database::listing::ListingDatabase;
use crate::database::{
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
TableNamesRequest,
CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest,
ReadConsistency, TableNamesRequest,
};
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
use crate::error::{Error, Result};
@@ -517,6 +517,39 @@ impl Connection {
self.internal.read_consistency().await
}
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
/// waiting on or cancelling the job.
///
/// The handle is constructed without a server round trip; an unknown id
/// surfaces when the handle is used. Only server-backed databases support
/// job handles by id.
pub fn job(&self, job_id: impl AsRef<str>) -> Result<crate::job::Job> {
self.internal.job(job_id.as_ref())
}
/// List server-side jobs across the database's tables.
pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
self.internal.list_jobs().await
}
/// Describe a single server-side job by id. `None` when the server has no
/// such job.
pub async fn get_job(&self, job_id: impl AsRef<str>) -> Result<Option<JobDescription>> {
self.internal.get_job(job_id.as_ref()).await
}
/// Request cancellation of a server-side job by id. Returns true if the
/// server accepted the cancellation, false if no such job exists.
pub async fn cancel_job(&self, job_id: impl AsRef<str>) -> Result<bool> {
self.internal.cancel_job(job_id.as_ref()).await
}
/// The lifecycle event history of a server-side job (all jobs when
/// `job_id` is `None`), as recorded Arrow batches.
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
self.internal.job_history(job_id).await
}
/// Drop a table in the database.
///
/// # Arguments
+66
View File
@@ -18,6 +18,8 @@ use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use arrow_array::RecordBatch;
use lance::dataset::ReadParams;
use lance_namespace::LanceNamespace;
use lance_namespace::models::{
@@ -200,6 +202,45 @@ pub enum ReadConsistency {
Strong,
}
/// A row from [`Database::list_jobs`]: one server-side job (index build,
/// compaction, column refresh, ...).
#[derive(Debug, Clone)]
pub struct JobInfo {
/// The job id -- what [`Database::get_job`] and [`Database::cancel_job`]
/// accept.
pub job_id: String,
/// The table the job runs against, without URI or namespace.
pub table: String,
pub job_type: String,
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
pub state: String,
/// When the job was created, in milliseconds since the epoch.
pub created_at_millis: i64,
}
/// A described job from [`Database::get_job`]: lifecycle state plus the
/// job-type-specific specification.
#[derive(Debug, Clone)]
pub struct JobDescription {
pub job_id: String,
pub job_type: String,
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
pub state: String,
/// When the job was created, in milliseconds since the epoch.
pub creation_ms: i64,
/// The job-type-specific specification. Null when the server omits it.
pub spec: serde_json::Value,
/// Why the job failed, when the job is failed and the server reports a
/// reason.
pub failure: Option<crate::error::JobFailure>,
}
fn job_op_not_supported<T>(what: &str) -> Result<T> {
Err(crate::error::Error::NotSupported {
message: format!("{} is not supported by this database", what),
})
}
/// The `Database` trait defines the interface for database implementations.
///
/// A database is responsible for managing tables and their metadata.
@@ -245,6 +286,31 @@ pub trait Database:
///
/// See [`CloneTableRequest`] for detailed documentation and examples.
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
/// waiting on or cancelling the job. The handle is constructed without a
/// server round trip; an unknown id surfaces when the handle is used.
fn job(&self, _job_id: &str) -> Result<crate::job::Job> {
job_op_not_supported("job")
}
/// List server-side jobs across the database's tables.
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
job_op_not_supported("list_jobs")
}
/// Describe a single job by id. `None` when the server has no such job.
async fn get_job(&self, _job_id: &str) -> Result<Option<JobDescription>> {
job_op_not_supported("get_job")
}
/// Request cancellation of a job by id. Returns true if the server
/// accepted the cancellation, false if no such job exists. Cancelling an
/// already-terminal job is a no-op success.
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
job_op_not_supported("cancel_job")
}
/// The lifecycle event history of a job (all jobs when `job_id` is
/// `None`), as recorded Arrow batches.
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
job_op_not_supported("job_history")
}
/// Open a table in the database
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
/// Rename a table in the database
+24
View File
@@ -18,6 +18,7 @@ pub(crate) trait JobHandle: Send + Sync {
fn id(&self) -> Option<&str> {
None
}
async fn status(&self) -> Result<String>;
async fn wait(&self) -> Result<()>;
async fn cancel(&self) -> Result<()>;
}
@@ -65,6 +66,19 @@ impl Job {
self.handle.as_ref().and_then(|handle| handle.id())
}
/// The operation's current lifecycle state: "running", "finished",
/// "failed", or "cancelled".
///
/// A point snapshot; unlike [`Job::wait`] it does not block, raise on a
/// terminal failure state, or retry. States a newer server reports that
/// this client version does not know pass through as-is.
pub async fn status(&self) -> Result<String> {
match &self.handle {
None => Ok("finished".to_string()),
Some(handle) => handle.status().await,
}
}
/// Waits until the operation reaches a terminal state.
///
/// Returns [`crate::Error::JobFailed`] if the operation failed and
@@ -138,6 +152,16 @@ impl SpawnedJob {
#[async_trait]
impl JobHandle for SpawnedJob {
async fn status(&self) -> Result<String> {
let label = match &*self.outcome.borrow() {
None => "running",
Some(Outcome::Succeeded) => "finished",
Some(Outcome::Failed(_)) => "failed",
Some(Outcome::Cancelled) => "cancelled",
};
Ok(label.to_string())
}
async fn wait(&self) -> Result<()> {
let mut outcome = self.outcome.clone();
let settled = outcome
+331 -1
View File
@@ -20,7 +20,7 @@ use lance_namespace::models::{
use crate::Error;
use crate::database::{
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
OpenTableRequest, ReadConsistency, TableNamesRequest,
JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest,
};
use crate::error::Result;
use crate::remote::util::stream_as_body;
@@ -432,6 +432,73 @@ fn build_cache_key(name: &str, namespace: &[String]) -> String {
key.iter().map(|b| format!("{:02x}", b)).collect()
}
#[derive(serde::Deserialize)]
struct RemoteListJobRow {
job_id: String,
#[serde(default)]
table: String,
#[serde(default)]
job_type: String,
#[serde(default)]
state: String,
#[serde(default)]
created_at_millis: i64,
}
#[derive(serde::Deserialize)]
struct RemoteListJobsResponse {
#[serde(default)]
jobs: Vec<RemoteListJobRow>,
#[serde(default)]
page_token: Option<String>,
}
/// The server's account of why a job failed. Absent from older servers,
/// which report only the terminal state.
#[derive(serde::Deserialize)]
struct RemoteReportedFailure {
#[serde(default)]
phase: Option<String>,
#[serde(default)]
message: Option<String>,
#[serde(default)]
retryable: Option<bool>,
}
#[derive(serde::Deserialize)]
struct RemoteDescribeJobResponse {
job_id: String,
#[serde(default)]
job_type: String,
job_state: String,
#[serde(default)]
creation_ms: i64,
#[serde(default)]
spec: serde_json::Value,
#[serde(default)]
failure: Option<RemoteReportedFailure>,
}
/// Server job states -> the client vocabulary ("running" / "finished" /
/// "failed" / "cancelled"). Covers both the describe enum (IN_PROGRESS /
/// DONE / FAILED / CANCELLED) and the registry's lowercase list-row states
/// (in_progress / succeeded / failed / canceled / timed_out). States this
/// client version does not know (e.g. created, queued) pass through as-is.
fn job_state_to_client(state: &str) -> String {
match state {
"IN_PROGRESS" | "in_progress" => "running",
"DONE" | "done" | "succeeded" => "finished",
"FAILED" | "failed" | "TIMED_OUT" | "timed_out" => "failed",
"CANCELLED" | "cancelled" | "canceled" => "cancelled",
other => other,
}
.to_string()
}
/// Bound on `list_jobs` page walking; a warning is logged when the listing
/// is truncated at this many pages.
const MAX_LIST_JOBS_PAGES: usize = 100;
#[async_trait]
impl<S: HttpSend> Database for RemoteDatabase<S> {
fn uri(&self) -> &str {
@@ -445,6 +512,108 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
})
}
fn job(&self, job_id: &str) -> Result<crate::job::Job> {
Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new(
self.client.clone(),
job_id.to_string(),
))))
}
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
let mut out = Vec::new();
let mut page_token: Option<String> = None;
for page in 0..MAX_LIST_JOBS_PAGES {
let mut body = serde_json::json!({});
if let Some(token) = &page_token {
body["page_token"] = serde_json::Value::String(token.clone());
}
let req = self.client.post("/v1/jobs/list").json(&body);
let (request_id, rsp) = self.client.send(req).await?;
let rsp = self.client.check_response(&request_id, rsp).await?;
let body: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id)?;
out.extend(body.jobs.into_iter().map(|row| JobInfo {
job_id: row.job_id,
table: row.table,
job_type: row.job_type,
state: job_state_to_client(&row.state),
created_at_millis: row.created_at_millis,
}));
page_token = body.page_token;
if page_token.is_none() {
break;
}
if page + 1 == MAX_LIST_JOBS_PAGES {
log::warn!(
"list_jobs truncated after {} pages ({} jobs)",
MAX_LIST_JOBS_PAGES,
out.len()
);
}
}
Ok(out)
}
async fn get_job(&self, job_id: &str) -> Result<Option<JobDescription>> {
let req = self
.client
.post("/v1/jobs/describe")
.json(&serde_json::json!({ "job_id": job_id }));
let (request_id, rsp) = self.client.send(req).await?;
let rsp = match self.client.check_response(&request_id, rsp).await {
Ok(rsp) => rsp,
Err(Error::Http {
status_code: Some(StatusCode::NOT_FOUND),
..
}) => return Ok(None),
Err(err) => return Err(err),
};
let body: RemoteDescribeJobResponse = rsp.json().await.err_to_http(request_id)?;
Ok(Some(JobDescription {
job_id: body.job_id,
job_type: body.job_type,
state: job_state_to_client(&body.job_state),
creation_ms: body.creation_ms,
spec: body.spec,
failure: body.failure.map(|reported| crate::error::JobFailure {
phase: reported.phase,
message: reported.message,
retryable: reported.retryable,
source: None,
}),
}))
}
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
let req = self
.client
.post("/v1/jobs/cancel")
.json(&serde_json::json!({ "job_id": job_id }));
let (request_id, rsp) = self.client.send(req).await?;
match self.client.check_response(&request_id, rsp).await {
Ok(_) => Ok(true),
Err(Error::Http {
status_code: Some(StatusCode::NOT_FOUND),
..
}) => Ok(false),
Err(err) => Err(err),
}
}
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<arrow_array::RecordBatch>> {
let mut body = serde_json::json!({});
if let Some(job_id) = job_id {
body["job_id"] = serde_json::Value::String(job_id.to_string());
}
let req = self.client.post("/v1/jobs/query_events").json(&body);
let (request_id, rsp) = self.client.send(req).await?;
let rsp = self.client.check_response(&request_id, rsp).await?;
let bytes = rsp.bytes().await.err_to_http(request_id)?;
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
reader
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(Into::into)
}
async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
let mut req = if !request.namespace_path.is_empty() {
let namespace_id =
@@ -2094,4 +2263,165 @@ mod tests {
assert!(list_response.tables.contains(&"table3".to_string()));
}
}
#[tokio::test]
async fn test_list_jobs_paginates() {
let page = Arc::new(AtomicUsize::new(0));
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/jobs/list");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
match page.fetch_add(1, Ordering::SeqCst) {
0 => {
assert!(body.get("page_token").is_none());
http::Response::builder()
.status(200)
.body(
r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000}], "page_token": "next"}"#,
)
.unwrap()
}
_ => {
assert_eq!(body["page_token"], "next");
http::Response::builder()
.status(200)
.body(
r#"{"jobs": [{"job_id": "job-2", "table": "t2", "job_type": "create_index", "state": "succeeded", "created_at_millis": 2000}, {"job_id": "job-3", "table": "t3", "job_type": "create_index", "state": "timed_out", "created_at_millis": 3000}]}"#,
)
.unwrap()
}
}
});
let jobs = conn.list_jobs().await.unwrap();
assert_eq!(jobs.len(), 3);
assert_eq!(jobs[0].job_id, "job-1");
assert_eq!(jobs[0].table, "t1");
assert_eq!(jobs[0].state, "running");
assert_eq!(jobs[1].job_id, "job-2");
assert_eq!(jobs[1].state, "finished");
assert_eq!(jobs[1].created_at_millis, 2000);
assert_eq!(jobs[2].job_id, "job-3");
assert_eq!(jobs[2].state, "failed");
}
#[tokio::test]
async fn test_get_job() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/jobs/describe");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(body["job_id"], "job-1");
http::Response::builder()
.status(200)
.body(
r#"{"job_id": "job-1", "job_type": "create_index", "job_state": "FAILED", "creation_ms": 1000, "spec": {"column": "vec"}, "failure": {"phase": "execute", "message": "worker died", "retryable": true}}"#,
)
.unwrap()
});
let job = conn.get_job("job-1").await.unwrap().unwrap();
assert_eq!(job.job_id, "job-1");
assert_eq!(job.job_type, "create_index");
assert_eq!(job.state, "failed");
assert_eq!(job.creation_ms, 1000);
assert_eq!(job.spec["column"], "vec");
let failure = job.failure.unwrap();
assert_eq!(failure.phase.as_deref(), Some("execute"));
assert_eq!(failure.message.as_deref(), Some("worker died"));
assert_eq!(failure.retryable, Some(true));
}
#[tokio::test]
async fn test_get_job_missing_is_none() {
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(404)
.body("no such job")
.unwrap()
});
assert!(conn.get_job("nope").await.unwrap().is_none());
}
#[tokio::test]
async fn test_cancel_job() {
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.url().path(), "/v1/jobs/cancel");
http::Response::builder()
.status(200)
.body(r#"{"job_id": "job-1"}"#)
.unwrap()
});
assert!(conn.cancel_job("job-1").await.unwrap());
let conn = Connection::new_with_handler(|_| {
http::Response::builder()
.status(404)
.body("no such job")
.unwrap()
});
assert!(!conn.cancel_job("nope").await.unwrap());
}
#[tokio::test]
async fn test_job_history_parses_arrow_stream() {
let schema = Arc::new(Schema::new(vec![Field::new(
"state",
DataType::Utf8,
false,
)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(arrow_array::StringArray::from(vec![
"created", "done",
]))],
)
.unwrap();
let mut body = Vec::new();
{
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
}
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.url().path(), "/v1/jobs/query_events");
let req_body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(req_body["job_id"], "job-1");
http::Response::builder()
.status(200)
.body(body.clone())
.unwrap()
});
let batches = conn.job_history(Some("job-1")).await.unwrap();
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 2);
}
#[tokio::test]
async fn test_conn_job_waits_to_done() {
let polls = Arc::new(AtomicUsize::new(0));
let polls_ref = polls.clone();
let conn = Connection::new_with_handler(move |request| {
assert_eq!(request.url().path(), "/v1/jobs/describe");
let state = if polls_ref.fetch_add(1, Ordering::SeqCst) == 0 {
"IN_PROGRESS"
} else {
"DONE"
};
http::Response::builder()
.status(200)
.body(format!(
r#"{{"job_id": "job-1", "job_type": "create_index", "job_state": "{}", "creation_ms": 1}}"#,
state
))
.unwrap()
});
let job = conn.job("job-1").unwrap();
assert_eq!(job.id(), Some("job-1"));
assert_eq!(job.status().await.unwrap(), "running");
job.wait().await.unwrap();
assert_eq!(job.status().await.unwrap(), "finished");
assert!(polls.load(Ordering::SeqCst) >= 3);
}
}
+30 -7
View File
@@ -35,12 +35,28 @@ impl<'de> Deserialize<'de> for JobState {
}
}
impl JobState {
/// The client vocabulary label for this state.
fn client_label(&self) -> String {
match self {
Self::InProgress => "running".to_string(),
Self::Done => "finished".to_string(),
Self::Failed => "failed".to_string(),
Self::Cancelled => "cancelled".to_string(),
Self::Other(state) => state.clone(),
}
}
}
impl From<&str> for JobState {
fn from(state: &str) -> Self {
match state {
"IN_PROGRESS" => Self::InProgress,
"CANCELLED" => Self::Cancelled,
"FAILED" => Self::Failed,
// The server reports a timed-out job as FAILED on describe;
// accept the raw registry state too in case a future server
// stops folding it.
"FAILED" | "TIMED_OUT" => Self::Failed,
"DONE" => Self::Done,
other => Self::Other(other.to_string()),
}
@@ -51,9 +67,12 @@ impl From<&str> for JobState {
/// report only the terminal state.
#[derive(Deserialize)]
struct ReportedFailure {
phase: String,
message: String,
retryable: bool,
#[serde(default)]
phase: Option<String>,
#[serde(default)]
message: Option<String>,
#[serde(default)]
retryable: Option<bool>,
}
#[derive(Deserialize)]
@@ -98,6 +117,10 @@ impl<S: HttpSend> JobHandle for RemoteJob<S> {
Some(&self.job_id)
}
async fn status(&self) -> Result<String> {
Ok(self.describe().await?.job_state.client_label())
}
async fn wait(&self) -> Result<()> {
let mut interval = INITIAL_POLL_INTERVAL;
loop {
@@ -110,9 +133,9 @@ impl<S: HttpSend> JobHandle for RemoteJob<S> {
failure: description
.failure
.map(|reported| JobFailure {
phase: Some(reported.phase),
message: Some(reported.message),
retryable: Some(reported.retryable),
phase: reported.phase,
message: reported.message,
retryable: reported.retryable,
source: None,
})
.unwrap_or_default(),