diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index ad24644b6..fa4e0748a 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -25,6 +25,27 @@ the underlying connection has been closed. ## Methods +### cancelJob() + +```ts +abstract cancelJob(jobId): Promise +``` + +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`<`boolean`> + +*** + ### cloneTable() ```ts @@ -365,6 +386,26 @@ Drop an existing table. *** +### getJob() + +```ts +abstract getJob(jobId): Promise +``` + +Describe a single server-side job by id. + +Resolves to `null` when the server has no such job. + +#### Parameters + +* **jobId**: `string` + +#### Returns + +`Promise`<`null` \| [`JobDescription`](../interfaces/JobDescription.md)> + +*** + ### 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> +``` + +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`<`Table`<`any`>> + +*** + +### listJobs() + +```ts +abstract listJobs(): Promise +``` + +List server-side jobs across the database's tables. + +#### Returns + +`Promise`<[`JobInfo`](../interfaces/JobInfo.md)[]> + +*** + ### listNamespaces() ```ts diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index c46b9651f..9f723e9e8 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -51,6 +51,25 @@ Request cancellation. Cancelling a finished operation is a no-op. *** +### status() + +```ts +status(): Promise +``` + +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`<`string`> + +*** + ### wait() ```ts diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 323a24a49..7455a81ce 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -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) diff --git a/docs/src/js/interfaces/JobDescription.md b/docs/src/js/interfaces/JobDescription.md new file mode 100644 index 000000000..5118bf5d6 --- /dev/null +++ b/docs/src/js/interfaces/JobDescription.md @@ -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". diff --git a/docs/src/js/interfaces/JobFailureInfo.md b/docs/src/js/interfaces/JobFailureInfo.md new file mode 100644 index 000000000..5683ee5f0 --- /dev/null +++ b/docs/src/js/interfaces/JobFailureInfo.md @@ -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; +``` diff --git a/docs/src/js/interfaces/JobInfo.md b/docs/src/js/interfaces/JobInfo.md new file mode 100644 index 000000000..3596fc968 --- /dev/null +++ b/docs/src/js/interfaces/JobInfo.md @@ -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. diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 75c5540eb..89a9e992c 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -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"); + }, + ); + }); +}); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index ef03e276b..e63a7ae65 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -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; + + /** + * 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; + + /** + * Describe a single server-side job by id. + * + * Resolves to `null` when the server has no such job. + */ + abstract getJob(jobId: string): Promise; + + /** + * 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; + + /** + * 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; } /** @hideconstructor */ @@ -722,6 +760,30 @@ export class LocalConnection extends Connection { options?.newNamespacePath, ); } + + job(jobId: string): Job { + return this.inner.job(jobId); + } + + async listJobs(): Promise { + return this.inner.listJobs(); + } + + async getJob(jobId: string): Promise { + return this.inner.getJob(jobId); + } + + async cancelJob(jobId: string): Promise { + return this.inner.cancelJob(jobId); + } + + async jobHistory(jobId?: string): Promise { + const buf = await this.inner.jobHistory(jobId); + if (buf.length === 0) { + return new ArrowTable(); + } + return tableFromIPC(buf); + } } /** diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 2ce031458..319222421 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -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, diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 1a18651f0..c45321aba 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -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 { + 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> { + 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> { + 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 { + 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) -> napi::Result { + 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::::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( diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs index 52d082beb..6aaeee174 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -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 { + 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 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, + pub message: Option, + pub retryable: Option, +} + +/// 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, + /// Why the job failed, when the job is failed and the server reports a + /// reason. + pub failure: Option, +} + +impl From 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, + }), + } + } +} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index f5c59c152..47e727f99 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 90b8dc75b..b48e84cc5 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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. diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index 82911119b..d33b62cbf 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -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. diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 3171c3638..332886590 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -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. diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 9bb162b15..757d8a758 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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)) diff --git a/python/src/connection.rs b/python/src/connection.rs index e22813640..b97d48ad8 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -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 { + 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> { + 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::>()) + }) + } + + pub fn get_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> { + 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> { + 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, + ) -> PyResult> { + 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] diff --git a/python/src/job.rs b/python/src/job.rs index 83303ddea..56ee211f4 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -28,6 +28,14 @@ impl Job { self.inner.id().map(str::to_string) } + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py( + self_.py(), + async move { inner.status().await.infer_error() }, + ) + } + pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { 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 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, + message: Option, + retryable: Option, +} + +#[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, + failure: Option, +} + +#[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 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, + }), + } + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index 6c9f62df9..6b0c0cf97 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -46,6 +46,9 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index f4ac0018b..89e59e12e 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -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) -> Result { + self.internal.job(job_id.as_ref()) + } + + /// List server-side jobs across the database's tables. + pub async fn list_jobs(&self) -> Result> { + 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) -> Result> { + 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) -> Result { + 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> { + self.internal.job_history(job_id).await + } + /// Drop a table in the database. /// /// # Arguments diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 6ce5bdbc9..f99f6e12a 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -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, +} + +fn job_op_not_supported(what: &str) -> Result { + 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>; + /// 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 { + job_op_not_supported("job") + } + /// List server-side jobs across the database's tables. + async fn list_jobs(&self) -> Result> { + 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> { + 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 { + 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> { + job_op_not_supported("job_history") + } /// Open a table in the database async fn open_table(&self, request: OpenTableRequest) -> Result>; /// Rename a table in the database diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 4c2c76d95..789ce8312 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -18,6 +18,7 @@ pub(crate) trait JobHandle: Send + Sync { fn id(&self) -> Option<&str> { None } + async fn status(&self) -> Result; 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 { + 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 { + 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 diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index a67f9dbb3..839cb3797 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -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, + #[serde(default)] + page_token: Option, +} + +/// 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, + #[serde(default)] + message: Option, + #[serde(default)] + retryable: Option, +} + +#[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, +} + +/// 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 Database for RemoteDatabase { fn uri(&self) -> &str { @@ -445,6 +512,108 @@ impl Database for RemoteDatabase { }) } + fn job(&self, job_id: &str) -> Result { + Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new( + self.client.clone(), + job_id.to_string(), + )))) + } + + async fn list_jobs(&self) -> Result> { + let mut out = Vec::new(); + let mut page_token: Option = 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> { + 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 { + 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> { + 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::, _>>() + .map_err(Into::into) + } + async fn table_names(&self, request: TableNamesRequest) -> Result> { 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); + } } diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 26bbfbbf9..2fc99da59 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -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, + #[serde(default)] + message: Option, + #[serde(default)] + retryable: Option, } #[derive(Deserialize)] @@ -98,6 +117,10 @@ impl JobHandle for RemoteJob { Some(&self.job_id) } + async fn status(&self) -> Result { + 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 JobHandle for RemoteJob { 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(),