diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index 1c5abd89f..18c1deb3b 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -448,26 +448,6 @@ on the returned job to know when cleanup has finished. *** -### 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 @@ -482,48 +462,6 @@ 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 @@ -648,6 +586,30 @@ A page of table names and an *** +### openJob() + +```ts +abstract openJob(jobId): Promise +``` + +Open a server-side job by id, returning a handle with its record already +populated. Rejects when the server has no such job, the way +[Connection.openTable](Connection.md#opentable) does for a missing table. + +The returned [Job](Job.md) answers for its own state, specification, +result, failure and event history, so there is no separate +connection-level call for any of them. + +#### Parameters + +* **jobId**: `string` + +#### Returns + +`Promise`<[`Job`](Job.md)> + +*** + ### openMaterializedView() ```ts diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index 9f723e9e8..cfbe454ec 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -8,19 +8,46 @@ A handle to an operation that may still be running. -## Constructors +The operation may already be complete when the handle is created. -### new Job() +The detail getters read what the handle last observed. Submitting an +operation returns only a job id, so populating them eagerly would cost an +extra round trip on every call: + +- [Job.refresh](Job.md#refresh) and [Job.status](Job.md#status) fetch the whole record. +- [Job.wait](Job.md#wait) records the terminal state it establishes, but not the + rest of the record. +- Everything is null until one of those runs. + +## Accessors + +### creationMs ```ts -new Job(): Job +get creationMs(): null | number ``` +When the job was created, in milliseconds since the epoch. + #### Returns -[`Job`](Job.md) +`null` \| `number` -## Accessors +*** + +### failure + +```ts +get failure(): null | JobFailureInfo +``` + +Why the job failed, when it failed and the server reports a reason. + +#### Returns + +`null` \| [`JobFailureInfo`](../interfaces/JobFailureInfo.md) + +*** ### id @@ -28,8 +55,69 @@ new Job(): Job get id(): null | string ``` -Identifies the operation on the server that is running it. Operations -that run in this process have no server id. The value is opaque. +Identifies the operation on the server that is running it. + +Operations that run in this process have no server id. The value is +opaque: parsing it or storing it to resume the job later is not supported. + +#### Returns + +`null` \| `string` + +*** + +### jobType + +```ts +get jobType(): null | string +``` + +The job's type, as the server names it. Null for an in-process job, which +has no server-side record. + +#### Returns + +`null` \| `string` + +*** + +### result + +```ts +get result(): any +``` + +The job-type-specific terminal result. Null until the job succeeds, so a +job that never terminates reports its progress through [Job.events](Job.md#events) +instead. + +#### Returns + +`any` + +*** + +### spec + +```ts +get spec(): any +``` + +The job-type-specific specification it was submitted with. + +#### Returns + +`any` + +*** + +### state + +```ts +get state(): null | string +``` + +The last observed lifecycle state, without contacting the backend. #### Returns @@ -51,18 +139,61 @@ Request cancellation. Cancelling a finished operation is a no-op. *** +### events() + +```ts +events(options?): Promise> +``` + +This job's recorded lifecycle events. + +Where the getters above report a terminal result only once the job reaches +one, events are written as the job runs and outlive the workers that +produced them. A distributed job records a `claim`/`claim_complete` pair +per unit of work, each carrying `rows_processed`, so a job that never +finishes still accounts for what it did. + +The server caps results at 1000 rows by default and 10,000 at most, and +truncates without saying so, so pass `limit` for a job that emits an event +per fragment. `filter` is a SQL-like expression over the `state`, +`updated_by`, `emitted_from`, `emitted_by`, and `claim_entity` columns. + +#### Parameters + +* **options?**: [`JobEventsOptions`](../interfaces/JobEventsOptions.md) + +#### Returns + +`Promise`<`Table`<`any`>> + +*** + +### refresh() + +```ts +refresh(): Promise +``` + +Ask the backend for this job's current state, and for a server-side job +its full record, then cache it for the getters above. + +#### Returns + +`Promise`<`void`> + +*** + ### status() ```ts status(): Promise ``` -The operation's current lifecycle state: "running", "finished", -"failed", or "cancelled". +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. +A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject on a +terminal failure state. Also refreshes the getters above. #### Returns @@ -70,6 +201,22 @@ client version does not know pass through as-is. *** +### toString() + +```ts +toString(): string +``` + +Every field the handle currently knows, one per line, with the JSON +payloads indented -- a refresh job's spec and result are the point of +printing it. + +#### Returns + +`string` + +*** + ### wait() ```ts diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index beb9cbeff..eb0fc7d5a 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -96,7 +96,7 @@ - [IvfFlatOptions](interfaces/IvfFlatOptions.md) - [IvfPqOptions](interfaces/IvfPqOptions.md) - [IvfRqOptions](interfaces/IvfRqOptions.md) -- [JobDescription](interfaces/JobDescription.md) +- [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) - [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) diff --git a/docs/src/js/interfaces/JobDescription.md b/docs/src/js/interfaces/JobDescription.md deleted file mode 100644 index 5118bf5d6..000000000 --- a/docs/src/js/interfaces/JobDescription.md +++ /dev/null @@ -1,66 +0,0 @@ -[**@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/JobEventsOptions.md b/docs/src/js/interfaces/JobEventsOptions.md new file mode 100644 index 000000000..24831f4f1 --- /dev/null +++ b/docs/src/js/interfaces/JobEventsOptions.md @@ -0,0 +1,29 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / JobEventsOptions + +# Interface: JobEventsOptions + +Which of a job's events [Job.events](../classes/Job.md#events) returns. + +## Properties + +### filter? + +```ts +optional filter: string; +``` + +SQL-like filter over the event columns. + +*** + +### limit? + +```ts +optional limit: number; +``` + +Maximum event rows to return, up to the server maximum of 10,000. diff --git a/docs/src/js/interfaces/JobInfo.md b/docs/src/js/interfaces/JobInfo.md index 3596fc968..01a7ceefc 100644 --- a/docs/src/js/interfaces/JobInfo.md +++ b/docs/src/js/interfaces/JobInfo.md @@ -26,7 +26,7 @@ When the job was created, in milliseconds since the epoch. jobId: string; ``` -The job id -- what `Connection.getJob` and `Connection.cancelJob` +The job id -- what `Connection.openJob` and `Connection.cancelJob` accept. *** diff --git a/docs/src/python/python.md b/docs/src/python/python.md index b0d1bb426..28e774473 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -157,6 +157,12 @@ listing a storage directory. ::: lancedb.job.AsyncJob +::: lancedb.job.JobInfo + +::: lancedb.job.JobDescription + +::: lancedb.job.JobFailureInfo + ::: lancedb.sql.Query ::: lancedb.sql.AsyncQuery @@ -310,6 +316,12 @@ still work. Queries return descriptors. Call ::: lancedb.exceptions.MissingColumnError +::: lancedb.exceptions.JobNotFoundError + +::: lancedb.exceptions.JobFailedError + +::: lancedb.exceptions.JobCancelledError + ## Integrations ## Pydantic diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 708559b7b..519f0eb5f 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -939,6 +939,7 @@ describe("remote connection jobs surface", () => { const { tableFromArrays, tableToIPC } = await import("apache-arrow"); const eventsTable = tableFromArrays({ state: ["created", "succeeded"] }); const eventsBody = Buffer.from(tableToIPC(eventsTable, "stream")); + const queryEventsPayloads: Record[] = []; await withMockDatabase( (req, res) => { @@ -967,6 +968,16 @@ describe("remote connection jobs surface", () => { ); } } else if (req.url === "/v1/jobs/describe") { + if (payload["job_id"] === "job-2") { + res + .writeHead(200, { "Content-Type": "application/json" }) + .end( + '{"job_id": "job-2", "job_type": "refresh_column", ' + + '"job_state": "DONE", "creation_ms": 2000, ' + + '"result": {"rows_assigned": 1000000}}', + ); + return; + } if (payload["job_id"] !== "job-1") { res.writeHead(404).end("no such job"); return; @@ -988,6 +999,7 @@ describe("remote connection jobs surface", () => { .writeHead(200, { "Content-Type": "application/json" }) .end('{"job_id": "job-1"}'); } else if (req.url === "/v1/jobs/query_events") { + queryEventsPayloads.push(payload); res .writeHead(200, { "Content-Type": "application/vnd.apache.arrow.stream", @@ -1004,22 +1016,65 @@ describe("remote connection jobs surface", () => { 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); + // Opening a job hands back a populated handle; a missing one rejects. + await expect(db.openJob("missing")).rejects.toThrow("not found"); + const finished = await db.openJob("job-2"); + expect(finished.state).toEqual("finished"); + expect(finished.result).toEqual({ + // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format + rows_assigned: 1000000, + }); - const job = db.job("job-1"); + const job = await db.openJob("job-1"); expect(job.id).toEqual("job-1"); + + // openJob already populated the handle; refresh() re-reads it. + expect(job.state).toEqual("failed"); + await job.refresh(); + expect(job.state).toEqual("failed"); + expect(job.jobType).toEqual("create_index"); + expect(job.creationMs).toEqual(1000); + expect(job.spec).toEqual({ column: "vec" }); + expect(job.result).toBeNull(); + expect(job.failure?.message).toEqual("worker died"); + + // The handle reaches its own events, supplying its job id. + const jobEvents = await job.events({ + limit: 500, + filter: "state = 'claim_complete'", + }); + expect(jobEvents.numRows).toEqual(2); + expect(queryEventsPayloads.pop()).toEqual({ + // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format + job_id: "job-1", + limit: 500, + filter: "state = 'claim_complete'", + }); + + // Printing lays every known field out on its own line, with the JSON + // payloads indented rather than crammed onto one line. + expect(`${job}`).toEqual( + [ + "Job(", + ' id="job-1",', + ' state="failed",', + ' jobType="create_index",', + " creationMs=1000,", + " spec={", + ' "column": "vec"', + " },", + " failure={", + ' "phase": "execute",', + ' "message": "worker died",', + ' "retryable": true', + " },", + ")", + ].join("\n"), + ); + 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 263a338ab..094819ec1 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { tableFromIPC } from "apache-arrow"; import { Data, SchemaLike, @@ -16,6 +15,7 @@ import { makeEmptyTable, } from "./arrow"; import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry"; +import { Job } from "./job"; import { MaterializedView, MaterializedViewSelect, @@ -27,8 +27,6 @@ import type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - Job, - JobDescription, JobInfo, ListNamespacesResponse, ListTablesResponse, @@ -557,24 +555,19 @@ export abstract class Connection { ): Promise; /** - * A {@link Job} handle for a server-side job by id. + * Open a server-side job by id, returning a handle with its record already + * populated. Rejects when the server has no such job, the way + * {@link Connection.openTable} does for a missing table. * - * 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. + * The returned {@link Job} answers for its own state, specification, + * result, failure and event history, so there is no separate + * connection-level call for any of them. */ - abstract job(jobId: string): Job; + abstract openJob(jobId: string): Promise; /** 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. * @@ -582,13 +575,6 @@ export abstract class Connection { * 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 */ @@ -869,7 +855,7 @@ export class LocalConnection extends Connection { } async dropTableAsync(name: string, namespacePath?: string[]): Promise { - return this.inner.dropTableAsync(name, namespacePath ?? []); + return new Job(await this.inner.dropTableAsync(name, namespacePath ?? [])); } async dropAllTables(namespacePath?: string[]): Promise { @@ -928,29 +914,17 @@ export class LocalConnection extends Connection { ); } - job(jobId: string): Job { - return this.inner.job(jobId); + async openJob(jobId: string): Promise { + return new Job(await this.inner.openJob(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 34d7ce4d9..4f8ff77e5 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -94,13 +94,9 @@ export { RenameTableOptions, } from "./connection"; -export { - Job, - JobDescription, - JobFailureInfo, - JobInfo, - Session, -} from "./native.js"; +export { JobFailureInfo, JobInfo, Session } from "./native.js"; + +export { Job, JobEventsOptions } from "./job"; export { AutoQuery, diff --git a/nodejs/lancedb/job.ts b/nodejs/lancedb/job.ts new file mode 100644 index 000000000..0baa0f571 --- /dev/null +++ b/nodejs/lancedb/job.ts @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import { Table as ArrowTable, tableFromIPC } from "apache-arrow"; +import { JobFailureInfo, Job as NativeJob } from "./native"; + +/** Which of a job's events {@link Job.events} returns. */ +export interface JobEventsOptions { + /** Maximum event rows to return, up to the server maximum of 10,000. */ + limit?: number; + /** SQL-like filter over the event columns. */ + filter?: string; +} + +/** + * A handle to an operation that may still be running. + * + * The operation may already be complete when the handle is created. + * + * The detail getters read what the handle last observed. Submitting an + * operation returns only a job id, so populating them eagerly would cost an + * extra round trip on every call: + * + * - {@link Job.refresh} and {@link Job.status} fetch the whole record. + * - {@link Job.wait} records the terminal state it establishes, but not the + * rest of the record. + * - Everything is null until one of those runs. + * + * @hideconstructor + */ +export class Job { + private readonly inner: NativeJob; + + constructor(inner: NativeJob) { + this.inner = inner; + } + + /** + * Identifies the operation on the server that is running it. + * + * Operations that run in this process have no server id. The value is + * opaque: parsing it or storing it to resume the job later is not supported. + */ + get id(): string | null { + return this.inner.id ?? null; + } + + /** The last observed lifecycle state, without contacting the backend. */ + get state(): string | null { + return this.inner.state ?? null; + } + + /** + * The job's type, as the server names it. Null for an in-process job, which + * has no server-side record. + */ + get jobType(): string | null { + return this.inner.jobType ?? null; + } + + /** When the job was created, in milliseconds since the epoch. */ + get creationMs(): number | null { + return this.inner.creationMs ?? null; + } + + /** The job-type-specific specification it was submitted with. */ + // biome-ignore lint/suspicious/noExplicitAny: shape varies by job type + get spec(): any | null { + return parseJson(this.inner.specJson); + } + + /** + * The job-type-specific terminal result. Null until the job succeeds, so a + * job that never terminates reports its progress through {@link Job.events} + * instead. + */ + // biome-ignore lint/suspicious/noExplicitAny: shape varies by job type + get result(): any | null { + return parseJson(this.inner.resultJson); + } + + /** Why the job failed, when it failed and the server reports a reason. */ + get failure(): JobFailureInfo | null { + return this.inner.failure ?? null; + } + + /** + * 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. Also refreshes the getters above. + */ + async status(): Promise { + return this.inner.status(); + } + + /** Wait until the operation reaches a terminal state. */ + async wait(): Promise { + return this.inner.wait(); + } + + /** Request cancellation. Cancelling a finished operation is a no-op. */ + async cancel(): Promise { + return this.inner.cancel(); + } + + /** + * Ask the backend for this job's current state, and for a server-side job + * its full record, then cache it for the getters above. + */ + async refresh(): Promise { + return this.inner.refresh(); + } + + /** + * This job's recorded lifecycle events. + * + * Where the getters above report a terminal result only once the job reaches + * one, events are written as the job runs and outlive the workers that + * produced them. A distributed job records a `claim`/`claim_complete` pair + * per unit of work, each carrying `rows_processed`, so a job that never + * finishes still accounts for what it did. + * + * The server caps results at 1000 rows by default and 10,000 at most, and + * truncates without saying so, so pass `limit` for a job that emits an event + * per fragment. `filter` is a SQL-like expression over the `state`, + * `updated_by`, `emitted_from`, `emitted_by`, and `claim_entity` columns. + */ + async events(options?: JobEventsOptions): Promise { + const buf = await this.inner.events(options?.limit, options?.filter); + if (buf.length === 0) { + return new ArrowTable(); + } + return tableFromIPC(buf); + } + + /** + * Every field the handle currently knows, one per line, with the JSON + * payloads indented -- a refresh job's spec and result are the point of + * printing it. + */ + toString(): string { + if (this.state === null) { + const known = this.id === null ? "" : `id=${JSON.stringify(this.id)}, `; + return `Job(${known}not refreshed)`; + } + const fields: string[] = []; + if (this.id !== null) { + fields.push(`id=${JSON.stringify(this.id)}`); + } + fields.push(`state=${JSON.stringify(this.state)}`); + if (this.jobType !== null) { + fields.push(`jobType=${JSON.stringify(this.jobType)}`); + } + if (this.creationMs !== null) { + fields.push(`creationMs=${this.creationMs}`); + } + for (const [name, value] of [ + ["spec", this.spec], + ["result", this.result], + ] as const) { + if (value !== null) { + fields.push(`${name}=${indentJson(value)}`); + } + } + if (this.failure !== null) { + fields.push(`failure=${indentJson(this.failure)}`); + } + return `Job(${fields.map((field) => `\n${REPR_INDENT}${field},`).join("")}\n)`; + } + + [Symbol.for("nodejs.util.inspect.custom")](): string { + return this.toString(); + } +} + +const REPR_INDENT = " "; + +// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type +function indentJson(value: any): string { + return JSON.stringify(value, null, 4).replace(/\n/g, `\n${REPR_INDENT}`); +} + +// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type +function parseJson(raw: string | null | undefined): any | null { + return raw === null || raw === undefined ? null : JSON.parse(raw); +} diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 28591f51c..06f8cd991 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -19,6 +19,7 @@ import { import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry"; import { IndexOptions } from "./indices"; +import { Job } from "./job"; import { MergeInsertBuilder } from "./merge"; import { AddColumnsResult, @@ -30,7 +31,6 @@ import { DropColumnsResult, IndexConfig, IndexStatistics, - Job, LsmStats, Branches as NativeBranches, OptimizeStats, @@ -1124,13 +1124,15 @@ export class LocalTable extends Table { ): Promise { // biome-ignore lint/suspicious/noExplicitAny: skip const nativeIndex = (options?.config as any)?.inner; - return await this.inner.createIndexAsync( - nativeIndex, - column, - options?.replace, - options?.waitTimeoutSeconds, - options?.name, - options?.train, + return new Job( + await this.inner.createIndexAsync( + nativeIndex, + column, + options?.replace, + options?.waitTimeoutSeconds, + options?.name, + options?.train, + ), ); } @@ -1313,7 +1315,7 @@ export class LocalTable extends Table { } async refreshColumnAsync(column: string): Promise { - return await this.inner.refreshColumnAsync(column); + return new Job(await this.inner.refreshColumnAsync(column)); } async refreshMaterializedView( diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 5cf676256..586238359 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -442,13 +442,15 @@ impl Connection { self.get_inner()?.drop_all_tables(&ns).await.default_error() } - /// A `Job` handle for a server-side job by id. + /// Open a server-side job by id, returning a handle with its record + /// already populated. Rejects when the server has no such job. /// - /// 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()?; + /// The returned handle answers for its own state, specification, result, + /// failure and event history, so there is no separate connection-level + /// call for any of them. + #[napi(catch_unwind)] + pub async fn open_job(&self, job_id: String) -> napi::Result { + let job = self.get_inner()?.open_job(&job_id).await.default_error()?; Ok(crate::job::Job::new(job)) } @@ -459,17 +461,6 @@ impl Connection { 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)] @@ -477,34 +468,6 @@ impl Connection { 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 14013fd27..9c6559dfd 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -3,6 +3,9 @@ use std::sync::Arc; +use arrow_array::RecordBatch; +use lancedb::job::JobEventsRequest; +use napi::bindgen_prelude::Buffer; use napi_derive::napi; use crate::error::NapiErrorExt; @@ -55,12 +58,98 @@ impl Job { pub async fn cancel(&self) -> napi::Result<()> { self.inner.cancel().await.default_error() } + + /// Ask the backend for this job's current state, and for a server-side job + /// its full record, then cache it for the getters below. + /// + /// They are all null until this runs, because submitting an operation + /// returns only a job id. {@link Job.status} fetches the whole record too; + /// {@link Job.wait} records only the terminal state it establishes. + #[napi(catch_unwind)] + pub async fn refresh(&self) -> napi::Result<()> { + self.inner.refresh().await.default_error() + } + + /// The last observed lifecycle state, without contacting the backend. + #[napi(getter)] + pub fn state(&self) -> Option { + self.inner.state() + } + + /// The job's type, as the server names it. Null for an in-process job, + /// which has no server-side record. + #[napi(getter)] + pub fn job_type(&self) -> Option { + self.inner.job_type() + } + + /// When the job was created, in milliseconds since the epoch. + #[napi(getter)] + pub fn creation_ms(&self) -> Option { + self.inner.creation_ms() + } + + /// The job-type-specific specification as a JSON string, when present. + #[napi(getter)] + pub fn spec_json(&self) -> Option { + self.inner.spec().map(|spec| spec.to_string()) + } + + /// The job-type-specific terminal result as a JSON string. Null until the + /// job succeeds, so a job that never terminates reports its progress + /// through {@link Job.events} instead. + #[napi(getter)] + pub fn result_json(&self) -> Option { + self.inner.result().map(|result| result.to_string()) + } + + /// Why the job failed, when it failed and the server reports a reason. + #[napi(getter)] + pub fn failure(&self) -> Option { + self.inner.failure().map(|failure| JobFailureInfo { + phase: failure.phase, + message: failure.message, + retryable: failure.retryable, + }) + } + + /// This job's recorded lifecycle events, as an Arrow IPC stream buffer. + /// The TypeScript wrapper turns it into an Arrow table. + #[napi(catch_unwind)] + pub async fn events(&self, limit: Option, filter: Option) -> napi::Result { + let batches = self + .inner + .events(JobEventsRequest { limit, filter }) + .await + .default_error()?; + batches_to_ipc_buffer(&batches) + } +} + +/// Serialise Arrow batches as a single IPC stream for the TypeScript layer. +pub(crate) fn batches_to_ipc_buffer(batches: &[RecordBatch]) -> napi::Result { + 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)) } /// A row from `Connection.listJobs`: one server-side job. #[napi(object)] pub struct JobInfo { - /// The job id -- what `Connection.getJob` and `Connection.cancelJob` + /// The job id -- what `Connection.openJob` and `Connection.cancelJob` /// accept. pub job_id: String, /// The table the job runs against, without URI or namespace. @@ -91,36 +180,3 @@ pub struct JobFailureInfo { 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 3ea8d15c7..0f7b110ac 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -148,17 +148,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 open_job(self, job_id: str) -> Job: ... async def create_function_async(self, request_json: str) -> Job: ... async def get_function(self, name: str, version: str) -> str: ... async def list_functions(self) -> List[str]: ... async def drop_function(self, name: str, version: str) -> bool: ... 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 execute_query_async( self, query: str, @@ -244,9 +240,20 @@ class BlobFile: class Job: @property def id(self) -> Optional[str]: ... + @property + def _state(self) -> Optional[str]: ... + @property + def _description(self) -> Optional[JobDescription]: ... async def status(self) -> str: ... async def wait(self) -> Optional[str]: ... async def cancel(self) -> None: ... + async def refresh(self) -> None: ... + async def events( + self, + *, + limit: Optional[int] = None, + filter: Optional[str] = None, + ) -> pa.Table: ... class JobInfo: @property @@ -278,7 +285,13 @@ class JobDescription: @property def creation_ms(self) -> int: ... @property - def spec_json(self) -> Optional[str]: ... + def _spec_json(self) -> Optional[str]: ... + @property + def _result_json(self) -> Optional[str]: ... + @property + def spec(self) -> Optional[Any]: ... + @property + def result(self) -> Optional[Any]: ... @property def failure(self) -> Optional[JobFailureInfo]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 82dfa22f1..88718e968 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -76,7 +76,7 @@ if TYPE_CHECKING: from .pydantic import LanceModel from ._lancedb import Connection as LanceDbConnection - from ._lancedb import JobDescription, JobInfo + from ._lancedb import JobInfo from .common import DATA, URI from .embeddings import EmbeddingFunctionConfig from ._lancedb import Session @@ -745,26 +745,23 @@ class DBConnection(EnforceOverrides): "Function catalog operations are 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. + def open_job(self, job_id: str) -> Job: + """Open a server-side job by id, returning a handle with its record + already populated. - 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. + The returned [Job][lancedb.job.Job] answers for its own state, + specification, result, failure and event history, so there is no + separate connection-level call for any of them. + + Raises `JobNotFoundError` when the server has no such job, the way + `open_table` does for a missing table. """ - raise NotImplementedError("job is not supported for this connection type") + raise NotImplementedError("open_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. @@ -776,15 +773,6 @@ class DBConnection(EnforceOverrides): "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" - ) - def execute_query( self, query: str, @@ -1462,14 +1450,11 @@ 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. + def open_job(self, job_id: str) -> Job: + """Open a server-side job by id. See + [DBConnection.open_job][lancedb.db.DBConnection.open_job]. """ - return Job(self._conn.job(job_id)) + return Job(LOOP.run(self._conn.open_job(job_id))) @override def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]: @@ -1493,14 +1478,6 @@ class LanceDBConnection(DBConnection): """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. @@ -1511,14 +1488,6 @@ class LanceDBConnection(DBConnection): """ 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. @@ -2289,15 +2258,11 @@ 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. + async def open_job(self, job_id: str) -> AsyncJob: + """Open a server-side job by id. See + [DBConnection.open_job][lancedb.db.DBConnection.open_job]. """ - return AsyncJob(self._inner.job(job_id)) + return AsyncJob(await self._inner.open_job(job_id)) async def create_function_async( self, definition: UdfDefinition @@ -2337,13 +2302,6 @@ class AsyncConnection(object): """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. @@ -2353,13 +2311,6 @@ class AsyncConnection(object): """ 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 execute_query( self, query: str, diff --git a/python/python/lancedb/exceptions.py b/python/python/lancedb/exceptions.py index daa98ee6e..67f15cabe 100644 --- a/python/python/lancedb/exceptions.py +++ b/python/python/lancedb/exceptions.py @@ -35,3 +35,9 @@ class JobCancelledError(RuntimeError): """Exception raised when an asynchronous job was cancelled.""" pass + + +class JobNotFoundError(ValueError): + """Exception raised when opening a job the server does not have.""" + + pass diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index f688768cb..e57fde0ea 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -4,15 +4,27 @@ """Handles to operations a server may run asynchronously.""" import asyncio +import json from datetime import timedelta from typing import Any, Callable, Generic, Optional, TypeVar, cast +import pyarrow as pa + from lancedb.background_loop import LOOP from . import _lancedb +from ._lancedb import JobDescription, JobFailureInfo, JobInfo T = TypeVar("T") +__all__ = [ + "AsyncJob", + "Job", + "JobDescription", + "JobFailureInfo", + "JobInfo", +] + class AsyncJob(Generic[T]): """A handle to an operation that may still be running. @@ -78,6 +90,149 @@ class AsyncJob(Generic[T]): return await self._inner.cancel() + async def refresh(self) -> None: + """Ask the backend for this job's current state, and for a server-side + job its full record, then cache it for the properties below. + + The properties are all `None` until this runs, because submitting an + operation returns only a job id. `status` fetches the whole record too; + `wait` records only the terminal state it establishes. + """ + if self._inner is None: + return + await self._inner.refresh() + + @property + def state(self) -> Optional[str]: + """The last observed lifecycle state, without contacting the backend. + + `None` until the handle has talked to it. See :meth:`AsyncJob.refresh`. + """ + if self._inner is None: + return "finished" + return self._inner._state + + @property + def job_type(self) -> Optional[str]: + """The job's type, as the server names it. + + `None` for an in-process job, which has no server-side record. + """ + return self._field("job_type") + + @property + def creation_ms(self) -> Optional[int]: + """When the job was created, in milliseconds since the epoch.""" + return self._field("creation_ms") + + @property + def spec(self) -> Optional[Any]: + """The job-type-specific specification it was submitted with.""" + return self._field("spec") + + @property + def result(self) -> Optional[Any]: + """The job-type-specific terminal result, as reported data rather than + the typed model :meth:`AsyncJob.wait` returns. + + `None` until the job succeeds, so a job that never terminates reports + its progress through :meth:`AsyncJob.events` instead. + """ + return self._field("result") + + @property + def failure(self) -> Optional[JobFailureInfo]: + """Why the job failed, when it failed and the server reports a reason.""" + return self._field("failure") + + @property + def _spec_json(self) -> Optional[str]: + return self._field("_spec_json") + + @property + def _result_json(self) -> Optional[str]: + return self._field("_result_json") + + def _field(self, name: str) -> Optional[Any]: + description = self._inner._description if self._inner is not None else None + return getattr(description, name) if description is not None else None + + async def events( + self, + *, + limit: Optional[int] = None, + filter: Optional[str] = None, + ) -> "pa.Table": + """This job's recorded lifecycle events. + + Where the properties above report a terminal result only once the job + reaches one, events are written as the job runs and outlive the workers + that produced them. A distributed job records a `claim`/`claim_complete` + pair per unit of work, each carrying `rows_processed`, so a job that + never finishes still accounts for what it did. + + Parameters + ---------- + limit: int, optional + Maximum event rows to return. The server caps results at 1000 by + default and 10,000 at most, and truncates without saying so, so + pass this for a job that emits an event per fragment. + filter: str, optional + SQL-like expression over the `state`, `updated_by`, `emitted_from`, + `emitted_by`, and `claim_entity` columns, such as + ``state = 'claim_complete'``. + """ + if self._inner is None: + raise NotImplementedError( + "job event history is only available for server-side jobs" + ) + return await self._inner.events(limit=limit, filter=filter) + + def __repr__(self) -> str: + return _job_repr("AsyncJob", self) + + +_REPR_INDENT = " " * 4 + + +def _repr_payload(value: Any) -> str: + """Render a job payload as indented JSON, aligned under its field.""" + try: + rendered = json.dumps(value, indent=4) + except TypeError: + return repr(value) + return rendered.replace("\n", "\n" + _REPR_INDENT) + + +def _job_repr(kind: str, job: Any) -> str: + """Render every field the handle currently knows, omitting the rest. + + One field per line, with the JSON payloads indented, because a refresh + job's spec and result are the point of printing it. + """ + state = job.state + if state is None: + # Nothing has been fetched yet, so there is nothing to lay out. + known = f"id={job.id!r}, " if job.id is not None else "" + return f"{kind}({known}not refreshed)" + + fields = [] + if job.id is not None: + fields.append(f"id={job.id!r}") + fields.append(f"state={state!r}") + for name in ("job_type", "creation_ms"): + value = getattr(job, name) + if value is not None: + fields.append(f"{name}={value!r}") + for name in ("spec", "result"): + value = getattr(job, name) + if value is not None: + fields.append(f"{name}={_repr_payload(value)}") + if job.failure is not None: + fields.append(f"failure={job.failure!r}") + body = "".join(f"\n{_REPR_INDENT}{field}," for field in fields) + return f"{kind}({body}\n)" + class Job(Generic[T]): """Synchronous counterpart of `AsyncJob` with the same result type.""" @@ -122,6 +277,75 @@ class Job(Generic[T]): return LOOP.run(self._inner.cancel()) + def refresh(self) -> None: + """Ask the backend for this job's current state and record. + + See :meth:`AsyncJob.refresh`. + """ + if self._inner is None: + return + LOOP.run(self._inner.refresh()) + + @property + def state(self) -> Optional[str]: + """The last observed lifecycle state. See :attr:`AsyncJob.state`.""" + return self._inner.state if self._inner is not None else "finished" + + @property + def job_type(self) -> Optional[str]: + """The job's type. See :attr:`AsyncJob.job_type`.""" + return self._field("job_type") + + @property + def creation_ms(self) -> Optional[int]: + """When the job was created. See :attr:`AsyncJob.creation_ms`.""" + return self._field("creation_ms") + + @property + def spec(self) -> Optional[Any]: + """The job's specification. See :attr:`AsyncJob.spec`.""" + return self._field("spec") + + @property + def result(self) -> Optional[Any]: + """The job's terminal result. See :attr:`AsyncJob.result`.""" + return self._field("result") + + @property + def failure(self) -> Optional[JobFailureInfo]: + """Why the job failed. See :attr:`AsyncJob.failure`.""" + return self._field("failure") + + @property + def _spec_json(self) -> Optional[str]: + return self._field("_spec_json") + + @property + def _result_json(self) -> Optional[str]: + return self._field("_result_json") + + def _field(self, name: str) -> Optional[Any]: + return getattr(self._inner, name) if self._inner is not None else None + + def events( + self, + *, + limit: Optional[int] = None, + filter: Optional[str] = None, + ) -> "pa.Table": + """This job's recorded lifecycle events. + + See :meth:`AsyncJob.events`. + """ + if self._inner is None: + raise NotImplementedError( + "job event history is only available for server-side jobs" + ) + return LOOP.run(self._inner.events(limit=limit, filter=filter)) + + def __repr__(self) -> str: + return _job_repr("Job", self) + def _typed_job( inner: "_lancedb.Job", result_decoder: Callable[[str], T] diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 1c0dd3afa..5b6828b04 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -31,7 +31,7 @@ from ..sql import QueryDescription from ..materialized_view import MaterializedView, SelectArg if TYPE_CHECKING: - from .._lancedb import JobDescription, JobInfo + from .._lancedb import JobInfo from ..embeddings import EmbeddingFunctionConfig from lance_namespace import ( LanceNamespace, @@ -739,14 +739,11 @@ 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. + def open_job(self, job_id: str) -> Job: + """Open a server-side job by id. See + [DBConnection.open_job][lancedb.db.DBConnection.open_job]. """ - return Job(self._conn.job(job_id)) + return Job(LOOP.run(self._conn.open_job(job_id))) @override def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]: @@ -769,14 +766,6 @@ class RemoteDBConnection(DBConnection): """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. @@ -787,14 +776,6 @@ class RemoteDBConnection(DBConnection): """ 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 execute_query_async( self, diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index add995de1..1e5a71e9a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2467,7 +2467,7 @@ def test_remote_blob_byte_apis_not_supported_on_old_server(): def test_remote_connection_jobs_surface(): - from lancedb.exceptions import JobFailedError + from lancedb.exceptions import JobFailedError, JobNotFoundError schema = pa.schema([("state", pa.string())]) batch = pa.record_batch([pa.array(["created", "done"])], schema=schema) @@ -2475,6 +2475,7 @@ def test_remote_connection_jobs_surface(): with pa.ipc.new_stream(sink, schema) as writer: writer.write_batch(batch) events_body = sink.getvalue().to_pybytes() + query_events_payloads = [] def handler(request): content_len = int(request.headers.get("Content-Length", 0)) @@ -2512,6 +2513,22 @@ def test_remote_connection_jobs_surface(): request.end_headers() request.wfile.write(json.dumps(rsp).encode()) elif request.path == "/v1/jobs/describe": + if payload["job_id"] == "job-2": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + json.dumps( + dict( + job_id="job-2", + job_type="refresh_column", + job_state="DONE", + creation_ms=2000, + result=dict(rows_assigned=1000000, rows_failed=0), + ) + ).encode() + ) + return if payload["job_id"] != "job-1": request.send_response(404) request.end_headers() @@ -2543,7 +2560,7 @@ def test_remote_connection_jobs_surface(): request.end_headers() request.wfile.write(b'{"job_id": "job-1"}') elif request.path == "/v1/jobs/query_events": - assert payload["job_id"] == "job-1" + query_events_payloads.append(payload) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.stream") request.end_headers() @@ -2559,24 +2576,109 @@ def test_remote_connection_jobs_surface(): 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"] + # Opening a job hands back a populated handle; a missing one fails. + with pytest.raises(JobNotFoundError, match="missing"): + db.open_job("missing") + finished = db.open_job("job-2") + assert finished.state == "finished" + assert finished.result == {"rows_assigned": 1000000, "rows_failed": 0} - job = db.job("job-1") + job = db.open_job("job-1") assert job.id == "job-1" + # Opening already populated the handle. + assert job.state == "failed" + assert job.spec == {"column": "vec"} + assert job.failure.message == "worker died" assert job.status() == "failed" with pytest.raises(JobFailedError, match="worker died"): job.wait(timeout=timedelta(seconds=5)) + + +def test_remote_job_handle_reports_its_own_detail(): + schema = pa.schema([("state", pa.string())]) + batch = pa.record_batch([pa.array(["claim_complete"])], schema=schema) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, schema) as writer: + writer.write_batch(batch) + events_body = sink.getvalue().to_pybytes() + event_payloads = [] + + 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/describe": + 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="refresh_column", + job_state="DONE", + creation_ms=2000, + spec=dict(column="vec"), + result=dict(rows_assigned=1000000), + ) + ).encode() + ) + elif request.path == "/v1/jobs/query_events": + event_payloads.append(payload) + 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: + job = db.open_job("job-1") + + # Opening populates the handle in the same round trip. + assert job.state == "finished" + job.refresh() + assert job.job_type == "refresh_column" + assert job.creation_ms == 2000 + assert job.spec == {"column": "vec"} + assert job.result == {"rows_assigned": 1000000} + assert job.failure is None + # The JSON payloads stay reachable, but as internal APIs. + assert json.loads(job._spec_json) == {"column": "vec"} + assert json.loads(job._result_json) == {"rows_assigned": 1000000} + + # print() shows everything the handle knows and nothing it does not. + # print() lays every known field out on its own line, with the JSON + # payloads indented rather than crammed onto one line. + assert repr(job) == "\n".join( + [ + "Job(", + " id='job-1',", + " state='finished',", + " job_type='refresh_column',", + " creation_ms=2000,", + " spec={", + ' "column": "vec"', + " },", + " result={", + ' "rows_assigned": 1000000', + " },", + ")", + ] + ) + # Nothing it does not know shows up. + assert "failure" not in repr(job) + + events = job.events(filter="state = 'claim_complete'", limit=500) + assert isinstance(events, pa.Table) + assert events.column("state").to_pylist() == ["claim_complete"] + # The handle supplies job_id; the caller only narrows the query. + assert event_payloads[-1] == { + "job_id": "job-1", + "limit": 500, + "filter": "state = 'claim_complete'", + } diff --git a/python/src/connection.rs b/python/src/connection.rs index 882fdfd29..2d613966a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -13,11 +13,7 @@ use crate::{ runtime::future_into_py, table::Table, }; -use arrow::{ - datatypes::Schema, - ffi_stream::ArrowArrayStreamReader, - pyarrow::{FromPyArrow, ToPyArrow}, -}; +use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow}; use lancedb::{ connection::Connection as LanceConnection, connection::NamespaceClientPushdownOperation, @@ -28,7 +24,7 @@ use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, exceptions::{PyRuntimeError, PyValueError}, pyclass, pyfunction, pymethods, - types::{PyAnyMethods, PyDict, PyDictMethods, PyList, PyListMethods}, + types::{PyAnyMethods, PyDict, PyDictMethods, PyList}, }; #[pyclass] @@ -644,9 +640,12 @@ 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 open_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + let job = inner.open_job(&job_id).await.infer_error()?; + Ok(crate::job::Job::new(job)) + }) } pub fn create_function_async( @@ -716,38 +715,12 @@ impl Connection { }) } - 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/error.rs b/python/src/error.rs index b66afe47b..aa13a8e87 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -114,6 +114,12 @@ impl PythonErrorExt for std::result::Result { .getattr(intern!(py, "JobCancelledError"))?; Err(PyErr::from_value(cls.call1((err.to_string(),))?)) }), + LanceError::JobNotFound { .. } => Python::attach(|py| { + let cls = py + .import(intern!(py, "lancedb.exceptions"))? + .getattr(intern!(py, "JobNotFoundError"))?; + Err(PyErr::from_value(cls.call1((err.to_string(),))?)) + }), _ => self.runtime_error(), }, } diff --git a/python/src/job.rs b/python/src/job.rs index 688cba7f9..4922c701a 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -4,11 +4,50 @@ use std::sync::Arc; use crate::runtime::future_into_py; -use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods}; +use arrow::{ + datatypes::Schema, + pyarrow::{IntoPyArrow, Table as PyArrowTable}, +}; +use lancedb::job::JobEventsRequest; +use pyo3::{ + Bound, PyAny, PyRef, PyResult, Python, + exceptions::PyValueError, + pyclass, pymethods, + types::{PyAnyMethods, PyDict, PyDictMethods}, +}; use serde::Serialize; use crate::error::PythonErrorExt; +const REPR_INDENT: &str = " "; + +/// Parse a stored JSON payload into Python data. The bindings carry these as +/// strings because that is what crosses the boundary cheaply; the public +/// Python surface is the parsed form. +fn parse_json_payload<'py>( + py: Python<'py>, + raw: Option<&str>, +) -> PyResult>> { + match raw { + None => Ok(None), + Some(raw) => Ok(Some(py.import("json")?.call_method1("loads", (raw,))?)), + } +} + +/// A payload rendered as indented JSON, aligned under the field that holds it. +fn pretty_json_payload(py: Python<'_>, raw: Option<&str>) -> PyResult> { + let Some(parsed) = parse_json_payload(py, raw)? else { + return Ok(None); + }; + let kwargs = PyDict::new(py); + kwargs.set_item("indent", 4)?; + let rendered: String = py + .import("json")? + .call_method("dumps", (parsed,), Some(&kwargs))? + .extract()?; + Ok(Some(rendered.replace('\n', &format!("\n{REPR_INDENT}")))) +} + #[pyclass] pub struct Job { inner: Arc, String>>>, @@ -67,6 +106,48 @@ impl Job { Ok(()) }) } + + pub fn refresh(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.refresh().await.infer_error()?; + Ok(()) + }) + } + + /// The last observed lifecycle state, without contacting the backend. + #[getter] + pub fn _state(&self) -> Option { + self.inner.state() + } + + /// The last observed server-side record. `None` for an in-process job. + #[getter] + pub fn _description(&self) -> Option { + self.inner.description().map(JobDescription::from) + } + + #[pyo3(signature = (*, limit=None, filter=None))] + pub fn events( + self_: PyRef<'_, Self>, + limit: Option, + filter: Option, + ) -> PyResult> { + let inner = self_.inner.clone(); + let request = JobEventsRequest { limit, filter }; + future_into_py(self_.py(), async move { + let batches = inner.events(request).await.infer_error()?; + Python::attach(|py| { + let schema = batches + .first() + .map(|batch| batch.schema()) + .unwrap_or_else(|| Arc::new(Schema::empty())); + let table = PyArrowTable::try_new(batches, schema) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + table.into_pyarrow(py).map(|table| table.unbind()) + }) + }) + } } /// A row from `Connection.list_jobs`: one server-side job. @@ -121,7 +202,7 @@ impl JobFailureInfo { } } -/// A described job from `Connection.get_job`. +/// The server-side record behind a `Job` handle. #[pyclass(get_all, skip_from_py_object)] #[derive(Clone)] pub struct JobDescription { @@ -129,17 +210,49 @@ pub struct JobDescription { job_type: String, state: String, creation_ms: i64, - spec_json: Option, + /// Internal: the wire form behind the `spec` property. + _spec_json: Option, + /// Internal: the wire form behind the `result` property. + _result_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 - ) + /// The job-type-specific specification it was submitted with. + #[getter] + fn spec<'py>(&self, py: Python<'py>) -> PyResult>> { + parse_json_payload(py, self._spec_json.as_deref()) + } + + /// The job-type-specific terminal result. `None` until the job succeeds. + #[getter] + fn result<'py>(&self, py: Python<'py>) -> PyResult>> { + parse_json_payload(py, self._result_json.as_deref()) + } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + let mut fields = vec![ + format!("job_id={:?}", self.job_id), + format!("job_type={:?}", self.job_type), + format!("state={:?}", self.state), + format!("creation_ms={}", self.creation_ms), + ]; + // Lay the payloads out as indented JSON, the same way the `Job` repr + // does, so the two agree on how the same data looks. + for (name, payload) in [("spec", &self._spec_json), ("result", &self._result_json)] { + if let Some(rendered) = pretty_json_payload(py, payload.as_deref())? { + fields.push(format!("{name}={rendered}")); + } + } + if let Some(failure) = &self.failure { + fields.push(format!("failure={}", failure.__repr__())); + } + let body = fields + .iter() + .map(|field| format!("\n{REPR_INDENT}{field},")) + .collect::(); + Ok(format!("JobDescription({body}\n)")) } } @@ -150,7 +263,11 @@ impl From for JobDescription { job_type: description.job_type, state: description.state, creation_ms: description.creation_ms, - spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()), + _spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()), + _result_json: description + .result + .filter(|result| !result.is_null()) + .map(|result| result.to_string()), failure: description.failure.map(|failure| JobFailureInfo { phase: failure.phase, message: failure.message, diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index f04f04ce2..6ec4a6ec1 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, JobDescription, JobInfo, OpenTableRequest, - ReadConsistency, TableNamesRequest, + CloneTableRequest, Database, DatabaseOptions, JobInfo, OpenTableRequest, ReadConsistency, + TableNamesRequest, }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -670,14 +670,34 @@ 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. + /// Open a server-side job by id, returning a handle with its record + /// already populated. Fails with [`crate::Error::JobNotFound`] when the + /// server has no such job, the way [`Connection::open_table`] does for a + /// missing table. /// - /// 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()) + /// This is the one way in: the returned [`crate::job::Job`] answers for + /// its own state, specification, result, failure and event history, so + /// there is no separate connection-level call for any of them. + /// + /// # Example + /// + /// ```no_run + /// # use lancedb::job::JobEventsRequest; + /// # async fn open_job( + /// # connection: &lancedb::Connection, + /// # job_id: &str, + /// # ) -> Result<(), Box> { + /// let job = connection.open_job(job_id).await?; + /// println!("{:?} {:?}", job.state(), job.result()); + /// let done = job + /// .events(JobEventsRequest::default().filter("state = 'claim_complete'")) + /// .await?; + /// println!("{} completions", done.iter().map(|b| b.num_rows()).sum::()); + /// # Ok(()) + /// # } + /// ``` + pub async fn open_job(&self, job_id: impl AsRef) -> Result { + self.internal.open_job(job_id.as_ref()).await } /// List server-side jobs across the database's tables. @@ -685,24 +705,12 @@ impl Connection { 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 532ea3658..843170030 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -18,8 +18,6 @@ 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::{ @@ -206,8 +204,8 @@ pub enum ReadConsistency { /// compaction, column refresh, ...). #[derive(Debug, Clone)] pub struct JobInfo { - /// The job id -- what [`Database::get_job`] and [`Database::cancel_job`] - /// accept. + /// The job id -- what [`Database::open_job`] and + /// [`Database::cancel_job`] accept. pub job_id: String, /// The table the job runs against, without URI or namespace. pub table: String, @@ -218,8 +216,8 @@ pub struct JobInfo { pub created_at_millis: i64, } -/// A described job from [`Database::get_job`]: lifecycle state plus the -/// job-type-specific specification. +/// The server-side record behind a [`crate::job::Job`] handle: lifecycle +/// state plus the job-type-specific specification and result. #[derive(Debug, Clone)] pub struct JobDescription { pub job_id: String, @@ -230,6 +228,10 @@ pub struct JobDescription { pub creation_ms: i64, /// The job-type-specific specification. Null when the server omits it. pub spec: serde_json::Value, + /// The job-type-specific terminal result, for job types that define one. + /// `None` until the job succeeds, so a job that never terminates reports + /// its progress through [`crate::job::Job::events`] instead. + pub result: Option, /// Why the job failed, when the job is failed and the server reports a /// reason. pub failure: Option, @@ -315,31 +317,22 @@ pub trait Database: async fn drop_function(&self, _name: &str, _version: &str) -> Result { function_catalog_not_supported() } - /// 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") + /// Open a job by id, returning a handle with its record already + /// populated. Fails with [`crate::Error::JobNotFound`] when the server has + /// no such job. + async fn open_job(&self, _job_id: &str) -> Result { + job_op_not_supported("open_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") - } /// Start executing a SQL statement on a remote database. async fn execute_query_async( &self, diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index be4641388..f6a2df9a6 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -102,6 +102,8 @@ pub enum Error { }, #[snafu(display("Job{} was cancelled", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))] JobCancelled { job_id: Option }, + #[snafu(display("Job '{job_id}' was not found"))] + JobNotFound { job_id: String }, // 3rd party / external errors #[snafu(display("object_store error: {source}"))] diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 22f1a0450..54b46fc11 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -3,16 +3,53 @@ //! Handles to operations a server may run asynchronously. -use std::sync::Arc; +use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use arrow_array::RecordBatch; use async_trait::async_trait; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; +use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; +/// Which of a job's events [`Job::events`] returns. +/// +/// The handle already knows which job to ask about, so this narrows the +/// query rather than naming one. +#[derive(Debug, Clone, Default)] +pub struct JobEventsRequest { + /// Maximum event rows to return. The server applies its own default + /// (1000 rows) and maximum (10,000 rows) when this is `None`, and + /// truncates without saying so, which matters for a job with one event + /// per fragment. + pub limit: Option, + /// SQL-like filter over the event columns `state`, `updated_by`, + /// `emitted_from`, `emitted_by`, and `claim_entity`. For example + /// `state = 'claim_complete'` selects only per-claim completions. + pub filter: Option, +} + +impl JobEventsRequest { + pub fn limit(mut self, limit: u32) -> Self { + self.limit = Some(limit); + self + } + + pub fn filter(mut self, filter: impl Into) -> Self { + self.filter = Some(filter.into()); + self + } +} + +fn job_detail_not_supported(what: &str) -> Result { + Err(Error::NotSupported { + message: format!("{what} is only available for server-side jobs"), + }) +} + /// Backend-specific tracking for an asynchronous operation. #[async_trait] pub(crate) trait JobHandle: Send + Sync { @@ -23,6 +60,15 @@ pub(crate) trait JobHandle: Send + Sync { async fn status(&self) -> Result; async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; + /// The job's full server-side record. Backends that run the operation in + /// this process have none and keep the default. + async fn describe(&self) -> Result { + job_detail_not_supported("describing a job") + } + /// The job's recorded lifecycle events. + async fn events(&self, _request: JobEventsRequest) -> Result> { + job_detail_not_supported("job event history") + } } /// A backend-neutral successful terminal result. @@ -85,16 +131,34 @@ enum JobInner { Completed(T), } +/// What a handle last learned about its job. `state` is separate because an +/// in-process job can report one but has no server-side record behind it. +#[derive(Default)] +struct JobCache { + state: Option, + description: Option, +} + /// A handle to an operation that may still be running. /// /// The operation may already be complete when the handle is created. `T` is /// the endpoint's successful terminal result; unit-result operations use the /// default `Job<()>`. +/// +/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the +/// handle last observed. Submitting an operation returns only a job id, so +/// populating them eagerly would cost an extra round trip on every call: +/// +/// - [`Job::refresh`] and [`Job::status`] fetch the whole record. +/// - [`Job::wait`] records the terminal state it establishes, but not the rest +/// of the record; call [`Job::refresh`] for that. +/// - Everything is `None` until one of those runs. pub struct Job where T: Clone + Send + Sync + 'static, { inner: JobInner, + cache: RwLock, } impl std::fmt::Debug for Job @@ -102,18 +166,40 @@ where T: Clone + Send + Sync + 'static, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("Job") - .field("id", &self.id()) - .field("done", &matches!(self.inner, JobInner::Completed(_))) - .finish() + let cache = self.cache_read(); + let mut out = f.debug_struct("Job"); + out.field("id", &self.id()) + .field("done", &matches!(self.inner, JobInner::Completed(_))); + if let Some(state) = &cache.state { + out.field("state", state); + } + if let Some(description) = &cache.description { + out.field("job_type", &description.job_type) + .field("creation_ms", &description.creation_ms); + if !description.spec.is_null() { + out.field("spec", &description.spec); + } + if let Some(result) = &description.result { + out.field("result", result); + } + if let Some(failure) = &description.failure { + out.field("failure", failure); + } + } + out.finish() } } impl Job<()> { - /// A job whose operation finished before the handle was created. + /// A job whose operation finished before the handle was created. Its + /// state is known without asking anyone, so the cache starts populated. pub(crate) fn new_done() -> Self { Self { inner: JobInner::Completed(()), + cache: RwLock::new(JobCache { + state: Some("finished".to_string()), + description: None, + }), } } @@ -123,8 +209,21 @@ impl Job<()> { handle, decode: Arc::new(|_| Ok(())), }, + cache: RwLock::default(), } } + + /// A handle whose record the caller has already fetched, so the detail + /// accessors answer without a second round trip. + pub(crate) fn opened(handle: Box, description: JobDescription) -> Self { + let job = Self::new(handle); + { + let mut cache = job.cache_write(); + cache.state = Some(description.state.clone()); + cache.description = Some(description); + } + job + } } impl Job @@ -138,6 +237,7 @@ where handle, decode: Arc::new(TerminalResult::decode::), }, + cache: RwLock::default(), } } } @@ -169,16 +269,124 @@ where } } + fn cache_read(&self) -> RwLockReadGuard<'_, JobCache> { + self.cache.read().unwrap_or_else(|err| err.into_inner()) + } + + fn cache_write(&self) -> RwLockWriteGuard<'_, JobCache> { + self.cache.write().unwrap_or_else(|err| err.into_inner()) + } + + /// Asks the backend for this job's current state, and for a server-side + /// job its full record, then caches the answer for the detail accessors. + /// + /// In-process operations have no server-side record, so only + /// [`Job::state`] is populated for them. + pub async fn refresh(&self) -> Result<()> { + self.refresh_state().await.map(|_| ()) + } + + /// Refreshes and reports the state, which every backend can answer. + async fn refresh_state(&self) -> Result { + let JobInner::Handle { handle, .. } = &self.inner else { + let state = "finished".to_string(); + self.cache_write().state = Some(state.clone()); + return Ok(state); + }; + match handle.describe().await { + Ok(description) => { + let state = description.state.clone(); + let mut cache = self.cache_write(); + cache.state = Some(state.clone()); + cache.description = Some(description); + Ok(state) + } + // An in-process job knows its own state and nothing more. + Err(Error::NotSupported { .. }) => { + let state = handle.status().await?; + self.cache_write().state = Some(state.clone()); + Ok(state) + } + Err(err) => Err(err), + } + } + /// 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. + /// this client version does not know pass through as-is. Also refreshes + /// the detail accessors. pub async fn status(&self) -> Result { + self.refresh_state().await + } + + /// The last lifecycle state this handle observed, without contacting the + /// backend. `None` until the handle has. + pub fn state(&self) -> Option { + self.cache_read().state.clone() + } + + /// The whole server-side record this handle last observed. The accessors + /// below read individual fields out of it. `None` for an in-process job, + /// which has no such record. + pub fn description(&self) -> Option { + self.cache_read().description.clone() + } + + /// The job's type, as the server names it. `None` for an in-process job. + pub fn job_type(&self) -> Option { + self.with_description(|description| description.job_type.clone()) + } + + /// When the job was created, in milliseconds since the epoch. `None` for + /// an in-process job. + pub fn creation_ms(&self) -> Option { + self.with_description(|description| description.creation_ms) + } + + /// The job-type-specific specification it was submitted with. + pub fn spec(&self) -> Option { + self.with_description(|description| description.spec.clone()) + .filter(|spec| !spec.is_null()) + } + + /// The job-type-specific terminal result, as reported data rather than the + /// typed model [`Job::wait`] returns. `None` until the job succeeds. + pub fn result(&self) -> Option { + self.with_description(|description| description.result.clone()) + .flatten() + } + + /// Why the job failed, when it failed and the server reports a reason. + pub fn failure(&self) -> Option { + self.with_description(|description| description.failure.clone()) + .flatten() + } + + fn with_description(&self, read: impl FnOnce(&JobDescription) -> R) -> Option { + self.cache_read().description.as_ref().map(read) + } + + /// This job's recorded lifecycle events. + /// + /// Unlike the detail accessors, which report a terminal result only once + /// the job reaches one, events are written as the job runs and outlive the + /// workers that produced them. A distributed job records a + /// `claim`/`claim_complete` pair per unit of work, each carrying + /// `rows_processed`, so a job that never finishes still accounts for what + /// it did. In-process operations keep no event history. + pub async fn events(&self, request: JobEventsRequest) -> Result> { match &self.inner { - JobInner::Handle { handle, .. } => handle.status().await, - JobInner::Completed(_) => Ok("finished".to_string()), + JobInner::Handle { handle, .. } => handle.events(request).await, + // The operation finished before the handle existed, so there is no + // id to query with even when a server ran it. + JobInner::Completed(_) => Err(Error::NotSupported { + message: "this operation completed before its handle was created, so it \ + carries no job id to query events with" + .to_string(), + }), } } @@ -190,8 +398,19 @@ where /// [`crate::Error::JobCancelled`] if it was cancelled. pub async fn wait(&self) -> Result { match &self.inner { - JobInner::Handle { handle, decode } => (decode)(handle.wait().await?), - JobInner::Completed(result) => Ok(result.clone()), + JobInner::Handle { handle, decode } => { + let settled = handle.wait().await; + // Waiting already established a terminal state; record it so + // the detail accessors do not need another round trip for it. + if let Some(state) = terminal_state(&settled) { + self.cache_write().state = Some(state.to_string()); + } + (decode)(settled?) + } + JobInner::Completed(result) => { + self.cache_write().state = Some("finished".to_string()); + Ok(result.clone()) + } } } @@ -224,20 +443,36 @@ where U: Clone + Send + Sync + 'static, F: Fn(T) -> U + Send + Sync + 'static, { - match self.inner { + // The mapped handle tracks the same job, so it inherits what this one + // has already learned about it. + let Self { inner, cache } = self; + match inner { JobInner::Handle { handle, decode } => Job { inner: JobInner::Handle { handle, decode: Arc::new(move |result| Ok(map((decode)(result)?))), }, + cache, }, JobInner::Completed(result) => Job { inner: JobInner::Completed(map(result)), + cache, }, } } } +/// The lifecycle state a settled [`JobHandle::wait`] implies. +fn terminal_state(settled: &Result) -> Option<&'static str> { + match settled { + Ok(_) => Some("finished"), + Err(Error::JobFailed { .. }) => Some("failed"), + Err(Error::JobCancelled { .. }) => Some("cancelled"), + // Anything else is a transport failure, not a verdict on the job. + Err(_) => None, + } +} + /// How an in-process operation ended. Cloneable so every waiter can be given /// the outcome; [`Error`] is not, so failures share one behind an [`Arc`]. #[derive(Clone)] diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 87486a522..32ace368e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -20,13 +20,13 @@ use lance_namespace::models::{ use crate::Error; use crate::database::{ - CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, - JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, + CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, JobInfo, + OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; use crate::function::{FunctionRegistrationRequest, FunctionVersion}; use crate::job::Job; -use crate::remote::job::{DescribeJobResponse, RemoteJob, job_state_to_client}; +use crate::remote::job::{RemoteJob, job_state_to_client}; use crate::remote::util::stream_as_body; use crate::table::BaseTable; @@ -671,11 +671,18 @@ impl Database for RemoteDatabase { Ok(response.dropped) } - 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 open_job(&self, job_id: &str) -> Result { + let handle = super::job::RemoteJob::new(self.client.clone(), job_id.to_string()); + match crate::job::JobHandle::describe(&handle).await { + Ok(description) => Ok(Job::opened(Box::new(handle), description)), + Err(Error::Http { + status_code: Some(StatusCode::NOT_FOUND), + .. + }) => Err(Error::JobNotFound { + job_id: job_id.to_string(), + }), + Err(err) => Err(err), + } } async fn list_jobs(&self) -> Result> { @@ -712,31 +719,6 @@ impl Database for RemoteDatabase { 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: DescribeJobResponse = 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| reported.into_job_failure()), - })) - } - async fn cancel_job(&self, job_id: &str) -> Result { let req = self .client @@ -753,21 +735,6 @@ impl Database for RemoteDatabase { } } - 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 execute_query_async( &self, query: &str, @@ -1307,6 +1274,7 @@ mod tests { use crate::{ Connection, Error, database::CreateTableMode, + job::JobEventsRequest, remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE}, }; @@ -2655,7 +2623,7 @@ mod tests { } #[tokio::test] - async fn test_get_job() { + async fn test_open_job() { let conn = Connection::new_with_handler(|request| { assert_eq!(request.method(), &reqwest::Method::POST); assert_eq!(request.url().path(), "/v1/jobs/describe"); @@ -2669,51 +2637,55 @@ mod tests { ) .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(); + // Opening populates the handle, so the accessors answer without a + // second round trip. + let job = conn.open_job("job-1").await.unwrap(); + assert_eq!(job.id(), Some("job-1")); + assert_eq!(job.job_type().as_deref(), Some("create_index")); + assert_eq!(job.state().as_deref(), Some("failed")); + assert_eq!(job.creation_ms(), Some(1000)); + assert_eq!(job.spec().unwrap()["column"], "vec"); + assert!(job.result().is_none()); + 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() { + async fn test_open_job_reports_the_terminal_result() { 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"}"#) + .body( + r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "DONE", "creation_ms": 1000, "result": {"rows_assigned": 1000000, "rows_failed": 0}}"#, + ) .unwrap() }); - assert!(conn.cancel_job("job-1").await.unwrap()); + let job = conn.open_job("job-1").await.unwrap(); + assert_eq!(job.state().as_deref(), Some("finished")); + let result = job.result().unwrap(); + assert_eq!(result["rows_assigned"], 1_000_000); + assert_eq!(result["rows_failed"], 0); + } + #[tokio::test] + async fn test_open_job_missing_fails() { let conn = Connection::new_with_handler(|_| { http::Response::builder() .status(404) .body("no such job") .unwrap() }); - assert!(!conn.cancel_job("nope").await.unwrap()); + let err = conn.open_job("nope").await.unwrap_err(); + assert!( + matches!(&err, Error::JobNotFound { job_id } if job_id == "nope"), + "{err:?}" + ); } #[tokio::test] - async fn test_job_history_parses_arrow_stream() { + async fn test_job_events_scope_to_that_job() { let schema = Arc::new(Schema::new(vec![Field::new( "state", DataType::Utf8, @@ -2722,29 +2694,91 @@ mod tests { let batch = RecordBatch::try_new( schema.clone(), vec![Arc::new(arrow_array::StringArray::from(vec![ - "created", "done", + "claim_complete", ]))], ) .unwrap(); - let mut body = Vec::new(); + let mut events = Vec::new(); { - let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap(); + let mut writer = + arrow_ipc::writer::StreamWriter::try_new(&mut events, &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 = + let body: serde_json::Value = serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(req_body["job_id"], "job-1"); + if request.url().path() == "/v1/jobs/describe" { + return http::Response::builder() + .status(200) + .body( + r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"# + .as_bytes() + .to_vec(), + ) + .unwrap(); + } + assert_eq!(request.url().path(), "/v1/jobs/query_events"); + // The handle supplies job_id; the caller only narrows the query. + assert_eq!(body["job_id"], "job-1"); + assert_eq!(body["limit"], 500); + assert_eq!(body["filter"], "state = 'claim_complete'"); http::Response::builder() .status(200) - .body(body.clone()) + .body(events.clone()) .unwrap() }); - let batches = conn.job_history(Some("job-1")).await.unwrap(); + let job = conn.open_job("job-1").await.unwrap(); + let batches = job + .events( + JobEventsRequest::default() + .limit(500) + .filter("state = 'claim_complete'"), + ) + .await + .unwrap(); assert_eq!(batches.len(), 1); - assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[0].num_rows(), 1); + } + + #[tokio::test] + async fn test_job_events_keep_the_schema_when_nothing_matches() { + let schema = Arc::new(Schema::new(vec![Field::new( + "state", + DataType::Utf8, + false, + )])); + let mut events = Vec::new(); + { + let mut writer = + arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap(); + writer.finish().unwrap(); + } + let conn = Connection::new_with_handler(move |request| { + if request.url().path() == "/v1/jobs/describe" { + return http::Response::builder() + .status(200) + .body( + r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"# + .as_bytes() + .to_vec(), + ) + .unwrap(); + } + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + // Only the job id when the caller narrows nothing. + assert_eq!(body, serde_json::json!({ "job_id": "job-1" })); + http::Response::builder() + .status(200) + .body(events.clone()) + .unwrap() + }); + let job = conn.open_job("job-1").await.unwrap(); + let batches = job.events(JobEventsRequest::default()).await.unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 0); + assert_eq!(batches[0].schema(), schema); } #[tokio::test] @@ -2939,7 +2973,9 @@ mod tests { 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 { + // Two in-progress answers: one for the load, one for the first + // status poll. + let state = if polls_ref.fetch_add(1, Ordering::SeqCst) < 2 { "IN_PROGRESS" } else { "DONE" @@ -2952,11 +2988,13 @@ mod tests { )) .unwrap() }); - let job = conn.job("job-1").unwrap(); + let job = conn.open_job("job-1").await.unwrap(); assert_eq!(job.id(), Some("job-1")); + // Opening already answered the state; no extra call needed for it. + assert_eq!(job.state().as_deref(), Some("running")); assert_eq!(job.status().await.unwrap(), "running"); job.wait().await.unwrap(); assert_eq!(job.status().await.unwrap(), "finished"); - assert!(polls.load(Ordering::SeqCst) >= 3); + assert!(polls.load(Ordering::SeqCst) >= 4); } } diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 0d41dbb35..c7acc3c05 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -5,13 +5,15 @@ use std::time::Duration; +use arrow_array::RecordBatch; use async_trait::async_trait; use tokio::time::sleep; use serde::Deserialize; +use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; -use crate::job::{JobHandle, TerminalResult}; +use crate::job::{JobEventsRequest, JobHandle, TerminalResult}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. @@ -86,7 +88,7 @@ pub(super) struct DescribeJobResponse { #[serde(default)] pub(super) spec: serde_json::Value, #[serde(default)] - result: Option, + pub(super) result: Option, #[serde(default)] pub(super) failure: Option, } @@ -110,6 +112,39 @@ impl DescribeJobResponse { fn into_terminal_result(self, request_id: String) -> TerminalResult { TerminalResult::remote(self.result, request_id) } + + /// The public description this wire envelope stands for. + pub(super) fn into_description(self) -> JobDescription { + JobDescription { + job_id: self.job_id, + job_type: self.job_type, + state: JobState::from(self.job_state.as_str()).client_label(), + creation_ms: self.creation_ms, + spec: self.spec, + result: self.result, + failure: self.failure.map(ReportedFailure::into_job_failure), + } + } +} + +/// One `/v1/jobs/query_events` round trip. +pub(super) async fn fetch_job_events( + client: &RestfulLanceDbClient, + body: serde_json::Value, +) -> Result> { + let request = client.post("/v1/jobs/query_events").json(&body); + let (request_id, response) = client.send(request).await?; + let response = client.check_response(&request_id, response).await?; + let bytes = response.bytes().await.err_to_http(request_id)?; + let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)?; + let schema = reader.schema(); + let mut batches = reader.collect::, _>>()?; + // A query that matched nothing still describes the event columns. + // Keep that schema so callers can build a typed empty result. + if batches.is_empty() { + batches.push(RecordBatch::new_empty(schema)); + } + Ok(batches) } pub struct RemoteJob { @@ -123,7 +158,7 @@ impl RemoteJob { } /// One `/v1/jobs/describe` round trip. - async fn describe(&self) -> Result<(String, DescribeJobResponse)> { + async fn fetch_description(&self) -> Result<(String, DescribeJobResponse)> { let request = self .client .post("/v1/jobs/describe") @@ -148,13 +183,28 @@ impl JobHandle for RemoteJob { } async fn status(&self) -> Result { - Ok(self.describe().await?.1.state().client_label()) + Ok(self.fetch_description().await?.1.state().client_label()) + } + + async fn describe(&self) -> Result { + Ok(self.fetch_description().await?.1.into_description()) + } + + async fn events(&self, request: JobEventsRequest) -> Result> { + let mut body = serde_json::json!({ "job_id": self.job_id }); + if let Some(limit) = request.limit { + body["limit"] = serde_json::Value::from(limit); + } + if let Some(filter) = request.filter { + body["filter"] = serde_json::Value::String(filter); + } + fetch_job_events(&self.client, body).await } async fn wait(&self) -> Result { let mut interval = INITIAL_POLL_INTERVAL; loop { - let (request_id, description) = self.describe().await?; + let (request_id, description) = self.fetch_description().await?; match description.state() { JobState::Done => return Ok(description.into_terminal_result(request_id)), JobState::Failed => { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 5faffa8c7..89885061e 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -264,6 +264,17 @@ impl crate::job::JobHandle for FreshnessJob { crate::job::JobHandle::status(&self.inner).await } + async fn describe(&self) -> Result { + crate::job::JobHandle::describe(&self.inner).await + } + + async fn events( + &self, + request: crate::job::JobEventsRequest, + ) -> Result> { + crate::job::JobHandle::events(&self.inner, request).await + } + async fn wait(&self) -> Result { let result = crate::job::JobHandle::wait(&self.inner).await?; let version = self.version.read().await; @@ -7982,6 +7993,72 @@ mod tests { ); } + /// The refresh handle is wrapped for read-freshness tracking, so it has to + /// forward the detail APIs too -- this is the job an operator is holding + /// when a backfill goes quiet. + #[tokio::test] + async fn test_refresh_job_handle_reports_detail_and_events() { + 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![ + "claim_complete", + ]))], + ) + .unwrap(); + let mut events = Vec::new(); + { + let mut writer = + arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + let table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-42"}"#.as_bytes().to_vec()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body( + r#"{"job_id": "j-42", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 7, "spec": {"column": "doubled"}}"# + .as_bytes() + .to_vec(), + ) + .unwrap(), + "/v1/jobs/query_events" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["job_id"], "j-42"); + http::Response::builder() + .status(200) + .body(events.clone()) + .unwrap() + } + other => panic!("unexpected path {other}"), + } + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + job.refresh().await.unwrap(); + assert_eq!(job.state().as_deref(), Some("running")); + assert_eq!(job.job_type().as_deref(), Some("refresh_column")); + assert_eq!(job.creation_ms(), Some(7)); + assert_eq!(job.spec().unwrap()["column"], "doubled"); + + let batches = job + .events(crate::job::JobEventsRequest::default()) + .await + .unwrap(); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); + } + #[tokio::test] async fn test_refresh_submission_uses_add_columns_version_fence() { let table = Table::new_with_handler("my_table", |request| match request.url().path() {