diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index cfbe454ec..74e1c00e2 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -17,7 +17,8 @@ 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. +- Registry details are null until one of those runs. Query handles report + jobType "query" immediately. ## Accessors @@ -57,8 +58,10 @@ 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: parsing it or storing it to resume the job later is not supported. +Registry handles expose a registry job ID; query handles expose a +connection-scoped SQL query UUID string. Local handles have no ID. +Query handles report jobType "query". Their IDs are not registry IDs or +durable resume tokens. #### Returns @@ -72,8 +75,8 @@ opaque: parsing it or storing it to resume the job later is not supported. 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. +"query" immediately for SQL-backed handles; otherwise the cached +server-reported job type. Null for in-process handles. #### Returns diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index 0577cd120..a87663a51 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -20,19 +20,20 @@ failures. Methods such as Python's `create_index_async` and TypeScript's `createIndexAsync` return job handles. In OSS, the work runs in your application -process; in Enterprise, it runs on server workers. A returned handle can -represent an operation that has already finished, so check its status before +process; in Enterprise, it runs on server workers or through the SQL query +backend. A returned handle can represent an operation that has already finished, +so check its status before waiting or requesting cancellation. See the operation's API reference for when the method returns. ## Submit and wait for a job {#waiting} 1. **Submit the operation.** Keep the returned job handle. For an Enterprise - job, also keep its ID if another process needs to look it up later. + registry job, also keep its ID if another process needs to look it up later. 2. **Check progress.** Read the current status while the operation runs. 3. **Wait for completion.** Check that the job succeeded before using its - result. If it failed or was cancelled, inspect the failure details and any - results it committed before stopping. + result. If it failed or was cancelled, inspect the failure details and any + results it committed before stopping. In Python, `job.status()` reports the current state and `job.wait()` waits for completion, raising if the job fails or is cancelled. If `job.wait()` times out, the client stops waiting. @@ -59,13 +60,37 @@ history. ## OSS and Enterprise execution {#in-process} -| | OSS | Enterprise | +The following table describes local operations and Enterprise registry jobs. +SQL-backed job handles follow the query lifecycle described below. + +| | OSS | Enterprise registry job | | --- | --- | --- | | Execution | In your application process | On server workers | | Job identity | No server ID | A server ID for looking up the job | | Event history | No server event history | Available through the jobs API | | Reattaching after a client restart | No server job to reopen | Reopen by ID while the record is retained | +### SQL-backed job handles + +Remote materialized-view CREATE and REFRESH use the SQL query backend. Python's +`create_materialized_view_async()` and `mv.refresh_async()` still return job +handles with `status()`, `wait()` and `cancel()`. Synchronous creation and refresh +wait on the same SQL execution before returning. + +The existing `job.job_type` identifies query-backed handles as `"query"` from +submission, and `job.id` contains the connection-scoped SQL query UUID string. +Use `describe_query(job.id)` on the submitting connection to inspect the query. +Registry handles retain their server-reported job types (available after +`status()` or `refresh()`) and registry IDs. Local handles have no ID or job type. + +SQL-backed handles do not expose registry records or event history and cannot +be reopened through `open_job()` or reattached after a client restart. Their IDs +are not durable resume tokens. Cancellation follows the SQL server's +capabilities and can be unsupported. + +The server decides whether SQL executes on its serving fleet or starts SQL job +workers. That placement does not change the handle's query-backed identity. + ## Learn more {#next} - [Build and repair indexes](/indexing/reindexing): submit index work. diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 3e47a00e9..3852af717 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -129,6 +129,7 @@ describe("materialized views", () => { const job = await db.dropMaterializedViewAsync("adults"); expect(job.id).toBeNull(); + expect(job.jobType).toBeNull(); await job.wait(); expect(await collect(db.listMaterializedViews())).toEqual([]); }); diff --git a/nodejs/lancedb/job.ts b/nodejs/lancedb/job.ts index 0baa0f571..1c8d74ad2 100644 --- a/nodejs/lancedb/job.ts +++ b/nodejs/lancedb/job.ts @@ -24,7 +24,8 @@ export interface JobEventsOptions { * - {@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. + * - Registry details are null until one of those runs. Query handles report + * jobType "query" immediately. * * @hideconstructor */ @@ -38,8 +39,10 @@ export class Job { /** * 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. + * Registry handles expose a registry job ID; query handles expose a + * connection-scoped SQL query UUID string. Local handles have no ID. + * Query handles report jobType "query". Their IDs are not registry IDs or + * durable resume tokens. */ get id(): string | null { return this.inner.id ?? null; @@ -51,8 +54,8 @@ export class Job { } /** - * The job's type, as the server names it. Null for an in-process job, which - * has no server-side record. + * "query" immediately for SQL-backed handles; otherwise the cached + * server-reported job type. Null for in-process handles. */ get jobType(): string | null { return this.inner.jobType ?? null; diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs index 214e687b4..320812128 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -29,8 +29,8 @@ impl Job { #[napi] impl Job { - /// Identifies the operation on the server that is running it. Operations - /// that run in this process have no server id. The value is opaque. + /// The operation's opaque ID. Interpret registry and query IDs using job_type(). + /// Local handles have no ID. #[napi(getter)] pub fn id(&self) -> Option { self.inner.id().map(str::to_string) @@ -62,8 +62,8 @@ impl Job { /// 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; + /// Registry details are null until this runs; query handles already expose + /// job_type "query". {@link Job.status} fetches the whole registry record too; /// {@link Job.wait} records only the terminal state it establishes. #[napi(catch_unwind)] pub async fn refresh(&self) -> napi::Result<()> { @@ -76,8 +76,8 @@ impl Job { self.inner.state() } - /// The job's type, as the server names it. Null for an in-process job, - /// which has no server-side record. + /// "query" immediately for SQL-backed handles; otherwise the cached + /// server-reported job type. Null for in-process handles. #[napi(getter)] pub fn job_type(&self) -> Option { self.inner.job_type() diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 6ef23abc7..9e602cce8 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -320,6 +320,8 @@ class Job: @property def id(self) -> Optional[str]: ... @property + def job_type(self) -> Optional[str]: ... + @property def _state(self) -> Optional[str]: ... @property def _description(self) -> Optional[JobDescription]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 8e53eac04..1558dd52a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -674,10 +674,11 @@ class DBConnection(EnforceOverrides): ) -> Job[None]: """Submit materialized-view creation and return its job. - The job may already be complete for a local database. On LanceDB - Cloud and Enterprise, its ``id`` is the server job identifier from - the ``202 Accepted`` create response. Wait for the job before opening - or querying the view. + The job may already be complete for a local database. Remote creation + submits CREATE MATERIALIZED VIEW through the SQL server. Its job has + ``job_type == "query"`` and ``id`` identifies the connection-scoped + query. Registry event history is unavailable. + Wait for the job before opening or querying the view. """ raise NotImplementedError( "materialized views are not supported on this connection type" @@ -1265,7 +1266,7 @@ class DBConnection(EnforceOverrides): """ raise NotImplementedError("SQL is not supported for this connection type") - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" raise NotImplementedError("SQL is not supported for this connection type") @@ -2869,7 +2870,9 @@ class AsyncConnection(object): ) -> AsyncJob[None]: """Submit materialized-view creation and return its job. - Wait for the returned job before opening or querying the view. + Wait for the returned job before opening or querying the view. Remote + creation uses SQL; ``job_type`` is ``"query"`` and ``id`` + identifies that query on this connection. """ source, native_options = await self._materialized_view_source( source, select, where, limit @@ -3409,8 +3412,14 @@ class AsyncConnection(object): ) ) - async def describe_query(self, query_id: UUID) -> QueryDescription: - """Describe a submitted SQL query by its connection-scoped id.""" + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: + """Describe a submitted SQL query by its UUID or string ID. + + For a handle with ``job_type == "query"``, pass its ``id`` on the + connection that submitted it. Registry job IDs are not query IDs. + """ + if isinstance(query_id, str): + query_id = UUID(query_id) return await self._inner.describe_query(query_id) async def namespace_client(self) -> LanceNamespace: diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index e57fde0ea..e61a188f2 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Handles to operations a server may run asynchronously.""" +"""Handles to asynchronous registry jobs, SQL queries, and local operations.""" import asyncio import json @@ -44,12 +44,12 @@ class AsyncJob(Generic[T]): @property def id(self) -> Optional[str]: - """Identifies the operation on the server that is running it. + """The operation's opaque ID, interpreted according to :attr:`job_type`. - Returned for correlating with server logs or the jobs API. Operations - that run in this process have no server id and return `None`. The value - is opaque: parsing it or storing it to resume the job later is not - supported. + Registry jobs expose a registry job ID; SQL-backed jobs expose a + connection-scoped query ID as a string. Local operations return ``None``. + Query IDs cannot be used with registry job APIs or to resume after a + client restart. """ return self._inner.id if self._inner is not None else None @@ -94,9 +94,11 @@ class AsyncJob(Generic[T]): """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. + Registry details are `None` until this runs, because submitting an + operation does not fetch its full registry record. Query-backed jobs + expose only lifecycle state through this method. For registry jobs, + `status` fetches the whole record too; `wait` records only the terminal + state it establishes. """ if self._inner is None: return @@ -114,11 +116,13 @@ class AsyncJob(Generic[T]): @property def job_type(self) -> Optional[str]: - """The job's type, as the server names it. + """The job's kind of work, without contacting the backend. - `None` for an in-process job, which has no server-side record. + SQL-backed handles report ``"query"`` immediately. Registry handles + report the server's job type after :meth:`refresh` or :meth:`status`. + In-process handles return ``None``. """ - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: @@ -294,7 +298,7 @@ class Job(Generic[T]): @property def job_type(self) -> Optional[str]: """The job's type. See :attr:`AsyncJob.job_type`.""" - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 55baa39a3..992926769 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -1683,7 +1683,7 @@ class AsyncLanceNamespaceDBConnection: parameters=parameters, ) - async def describe_query(self, query_id: UUID) -> QueryDescription: + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query when supported.""" return await self._inner.describe_query(query_id) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 5e9288471..62d00f284 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1134,7 +1134,7 @@ class RemoteDBConnection(DBConnection): ) @override - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" return LOOP.run( self._conn.describe_query( diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 874034564..61e511345 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1729,7 +1729,11 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip(): ) registration = db.create_function_async(normalize_score) assert registration.id == "job-register" + assert registration.job_type is None + assert registration.status() == "finished" + assert registration.job_type == "create_function" created = registration.wait() + assert registration.job_type == "create_function" reopened = db.get_function("normalize_score", version=created.version) assert created == reopened diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 46972a15c..e2be9c0d4 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -8,6 +8,7 @@ import threading import lancedb import pytest +from lancedb.db import AsyncConnection from lancedb.materialized_view import MaterializedViewDefinition from lancedb.remote.db import RemoteDBConnection @@ -56,7 +57,7 @@ def mock_remote_materialized_views(): @contextlib.contextmanager -def mock_remote_materialized_view_create(): +def mock_remote_materialized_view_drop(): requests = [] class Handler(http.server.BaseHTTPRequestHandler): @@ -67,7 +68,8 @@ def mock_remote_materialized_view_create(): length = int(self.headers.get("Content-Length", 0)) body = json.loads(self.rfile.read(length) or b"{}") requests.append((self.path, body)) - job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123" + assert self.path.endswith("/drop") + job_id = "mv-drop-123" encoded = json.dumps({"job_id": job_id}).encode() self.send_response(202) self.send_header("Content-Type", "application/json") @@ -98,7 +100,7 @@ def test_remote_list_uses_namespace_route(): def test_remote_drop_async_returns_server_job(): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect( "db://dev", api_key="fake", @@ -192,11 +194,13 @@ def test_create_and_refresh_jobs(tmp_path): "adults", "people", where="age >= 18", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert create_job.wait() is None view = db.open_materialized_view("adults") refresh_job = view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 2 @@ -337,11 +341,13 @@ async def test_async_create_and_refresh_jobs(tmp_path): "adults", "people", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert await create_job.wait() is None view = await db.open_materialized_view("adults") refresh_job = await view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = await refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 1 @@ -567,6 +573,27 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): assert list(resumed) == [] +class QueryBackedCreateJob: + id = "0198f1b2-c3d4-7e5f-8123-456789abcdef" + + job_type = "query" + + async def wait(self): + return None + + async def status(self): + return "finished" + + +class CapturingCreateConnection: + def __init__(self): + self.calls = [] + + async def create_materialized_view_async(self, name, source, **options): + self.calls.append((name, source, options)) + return QueryBackedCreateJob() + + @pytest.mark.parametrize( "options", [ @@ -576,16 +603,28 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): ], ) def test_remote_create_forwards_refresh_options_without_interval_default(options): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect("db://dev", api_key="fake", host_override=host) - job = db.create_materialized_view_async("adults", "people", **options) - assert job.id == "mv-create-123" - assert requests == [ + native = CapturingCreateConnection() + db = RemoteDBConnection.__new__(RemoteDBConnection) + db._conn = AsyncConnection(native) + job = db.create_materialized_view_async("adults", "people", **options) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert job.wait() is None + assert job.status() == "finished" + assert native.calls == [ ( - "/v1/materialized_view/adults/create", + "adults", + "people", { - "query": 'SELECT * FROM "people"', - **options, + "projections": None, + "filter": None, + "limit": None, + "with_no_data": False, + "auto_refresh_enabled": options.get("auto_refresh_enabled"), + "auto_refresh_interval_seconds": options.get( + "auto_refresh_interval_seconds" + ), }, ) ] @@ -593,18 +632,18 @@ def test_remote_create_forwards_refresh_options_without_interval_default(options @pytest.mark.asyncio async def test_async_remote_create_keeps_interval_unspecified(): - with mock_remote_materialized_view_create() as (host, requests): - db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host) - job = await db.create_materialized_view_async( - "adults", - "people", - auto_refresh_enabled=True, - ) - assert job.id == "mv-create-123" - assert requests[0][1] == { - "query": 'SELECT * FROM "people"', - "auto_refresh_enabled": True, - } + native = CapturingCreateConnection() + db = AsyncConnection(native) + job = await db.create_materialized_view_async( + "adults", "people", auto_refresh_enabled=True + ) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert await job.wait() is None + assert await job.status() == "finished" + assert native.calls[0][2]["auto_refresh_enabled"] is True + assert native.calls[0][2]["auto_refresh_interval_seconds"] is None @pytest.mark.parametrize( @@ -625,7 +664,7 @@ def test_embedded_create_refuses_automatic_refresh_before_creating_view( @pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")]) def test_remote_create_rejects_invalid_interval_without_request(seconds): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect("db://dev", api_key="fake", host_override=host) with pytest.raises(ValueError, match="positive, finite"): db.create_materialized_view_async( diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index e6f5fd3c7..539b88647 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -137,9 +137,23 @@ async def test_async_namespace_connection_rejects_sql(tmp_path): await connection.describe_query(NIL_QUERY_ID) -def test_describe_query_requires_uuid(): - with pytest.raises(TypeError, match="UUID"): - remote_connection().describe_query(str(NIL_QUERY_ID)) +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", [NIL_QUERY_ID, str(NIL_QUERY_ID)]) +async def test_describe_query_accepts_uuid_and_job_id_string(identity): + calls = [] + + class NativeConnection: + async def describe_query(self, query_id): + calls.append(query_id) + + connection = AsyncConnection(NativeConnection()) + await connection.describe_query(identity) + assert calls == [NIL_QUERY_ID] + + +def test_describe_query_rejects_registry_job_id(): + with pytest.raises(ValueError, match="UUID"): + remote_connection().describe_query("registry-job-id") @pytest.mark.parametrize( diff --git a/python/src/job.rs b/python/src/job.rs index e22b2f897..1740b3e00 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -81,6 +81,11 @@ impl Job { self.inner.id().map(str::to_string) } + #[getter] + pub fn job_type(&self) -> Option { + self.inner.job_type() + } + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py( diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index d671365f9..5c27587dc 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -13,6 +13,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; +use uuid::Uuid; use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; @@ -55,10 +56,13 @@ fn job_detail_not_supported(what: &str) -> Result { /// Backend-specific tracking for an asynchronous operation. #[async_trait] pub(crate) trait JobHandle: Send + Sync { - /// Server-assigned id, when the backend has one. + /// Backend-specific id, when the operation has one. fn id(&self) -> Option<&str> { None } + fn job_type(&self) -> Option<&str> { + None + } async fn status(&self) -> Result; async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; @@ -167,14 +171,15 @@ struct JobCache { /// the endpoint's successful terminal result; unit-result operations use the /// default `Job<()>`. /// -/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the +/// Registry 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. +/// - Registry details are `None` until one of those runs. Query-backed handles +/// report [`Job::job_type`] as `"query"` immediately. pub struct Job where T: Clone + Send + Sync + 'static, @@ -273,13 +278,18 @@ where Self::new_typed(Box::new(SpawnedJob::new(task))) } - pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + pub(crate) fn from_query_task( + query_id: Uuid, + task: JoinHandle>, + cancel: F, + ) -> Self where F: Fn() -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + query_id, task, cancellation, ))) @@ -290,12 +300,10 @@ impl Job where T: Clone + Send + Sync + 'static, { - /// Identifies the operation on the server that is running it. - /// - /// Returned for correlating with server logs or the jobs API. Operations - /// that run in this process have no server id and return `None`. The - /// value is opaque: parsing it or storing it to resume the job later is - /// not supported. + /// The operation's opaque identity, interpreted according to [`Job::job_type`]. + /// Registry handles expose a registry job ID; query handles expose the + /// connection-scoped SQL query UUID as a string. Local handles return None. + /// A query ID cannot be used with registry APIs or to resume after restart. pub fn id(&self) -> Option<&str> { match &self.inner { JobInner::Handle { handle, .. } => handle.id(), @@ -369,8 +377,14 @@ where self.cache_read().description.clone() } - /// The job's type, as the server names it. `None` for an in-process job. + /// The job's type: `"query"` for SQL-backed handles, otherwise the cached + /// server-reported kind of work. `None` for an in-process job. pub fn job_type(&self) -> Option { + if let JobInner::Handle { handle, .. } = &self.inner + && let Some(job_type) = handle.job_type() + { + return Some(job_type.to_string()); + } self.with_description(|description| description.job_type.clone()) } @@ -536,6 +550,7 @@ struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, cancellation: Option, + id: Option, } type CancellationFuture = Pin> + Send>>; @@ -549,11 +564,17 @@ impl SpawnedJob { Self::new_inner(task, None) } - fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + fn new_with_cancellation( + query_id: Uuid, + task: JoinHandle>, + cancellation: Cancellation, + ) -> Self where T: Serialize + Send + 'static, { - Self::new_inner(task, Some(cancellation)) + let mut handle = Self::new_inner(task, Some(cancellation)); + handle.id = Some(query_id.to_string()); + handle } fn new_inner(task: JoinHandle>, cancellation: Option) -> Self @@ -582,12 +603,21 @@ impl SpawnedJob { outcome, abort, cancellation, + id: None, } } } #[async_trait] impl JobHandle for SpawnedJob { + fn id(&self) -> Option<&str> { + self.id.as_deref() + } + + fn job_type(&self) -> Option<&str> { + self.id.as_ref().map(|_| "query") + } + async fn status(&self) -> Result { let label = match &*self.outcome.borrow() { None => "running", @@ -648,11 +678,31 @@ mod tests { async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); + assert_eq!(job.job_type(), None); + assert_eq!(job.id(), None); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.status().await.unwrap(), "finished"); } + #[tokio::test] + async fn mapped_query_job_preserves_identity_and_type() { + let query_id = Uuid::new_v4(); + let job = Job::from_query_task(query_id, tokio::spawn(async { Ok(41_u64) }), || async { + Ok(()) + }) + .map(|value| value + 1); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + let completed = Job::new_done(); + assert_eq!(completed.job_type(), None); + assert_eq!(completed.id(), None); + } + #[tokio::test] async fn mapped_spawned_job_preserves_cancellation() { let job = Job::spawned(tokio::spawn(async { pending::>().await })) diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index ecdd1783b..7ee76a160 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -1716,7 +1716,8 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. + /// settles when the view is ready. Remote creation submits SQL; the handle's + /// [`Job::job_type`] is `"query"` and its ID identifies that query. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 768de79d2..d33565653 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -63,15 +63,6 @@ struct RemoteCloneTableRequest { is_shallow: Option, } -#[derive(serde::Serialize)] -struct RemoteCreateMaterializedViewRequest { - query: String, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_enabled: Option, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_interval_seconds: Option, -} - // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. @@ -594,6 +585,11 @@ mod test_utils { } } + pub(in crate::remote) fn with_sql_client(mut self, client: SqlClient) -> Self { + self.sql_client = Some(client); + self + } + pub fn new_mock_with_config(handler: F, config: ClientConfig) -> Self where F: Fn(reqwest::Request) -> http::Response + Send + Sync + 'static, @@ -885,8 +881,7 @@ impl Database for RemoteDatabase { ) -> Result { if request.with_no_data { return Err(Error::NotSupported { - message: "remote materialized views are always populated by their SQL job" - .to_string(), + message: "remote materialized views must be populated during creation".to_string(), }); } if let Some(seconds) = request.auto_refresh_interval_seconds @@ -896,36 +891,37 @@ impl Database for RemoteDatabase { message: "auto_refresh_interval_seconds must be positive, finite seconds".into(), }); } - let identifier = build_table_identifier(&request.name, &request.namespace_path)?; - let body = RemoteCreateMaterializedViewRequest { - query: request.query, - auto_refresh_enabled: request.auto_refresh_enabled, - auto_refresh_interval_seconds: request.auto_refresh_interval_seconds, - }; - let request = self - .client - .post(&format!("/v1/materialized_view/{identifier}/create")) - .json(&body); - let (request_id, response) = self.client.send(request).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view creation must return 202 Accepted".into(), - request_id, - status_code: Some(status), - }); + validate_table_name(&request.name)?; + validate_namespace(&request.namespace_path)?; + let client = self + .sql_client + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".into(), + })?; + let mut options = Vec::new(); + if let Some(enabled) = request.auto_refresh_enabled { + options.push(format!("auto_refresh_enabled = {enabled}")); } - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "materialized-view create response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), - })?; - Ok(Job::new(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + if let Some(seconds) = request.auto_refresh_interval_seconds { + options.push(format!("auto_refresh_interval_seconds = {seconds}")); + } + let name = request.name.replace('"', "\"\""); + let options = if options.is_empty() { + String::new() + } else { + format!(" WITH ({})", options.join(", ")) + }; + client + .submit_as_job( + format!( + "CREATE MATERIALIZED VIEW \"{name}\"{options} AS {}", + request.query + ), + request.namespace_path, + || async { Ok(()) }, + ) + .await } async fn drop_materialized_view_async( @@ -2016,72 +2012,22 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_uses_rest_and_preserves_optional_settings() { - for (enabled, interval) in [(None, None), (Some(true), None), (Some(false), Some(0.5))] { - let db = super::RemoteDatabase::new_mock(move |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/create" - ); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body["query"], "SELECT * FROM people"); - assert_eq!( - body.get("auto_refresh_enabled"), - enabled.map(serde_json::Value::from).as_ref() - ); - assert_eq!( - body.get("auto_refresh_interval_seconds"), - interval.map(serde_json::Value::from).as_ref() - ); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) - .unwrap() - }); - let job = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT * FROM people".into(), - with_no_data: false, - auto_refresh_enabled: enabled, - auto_refresh_interval_seconds: interval, - }) - .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-create")); - } - } - - #[tokio::test] - async fn test_create_materialized_view_rejects_invalid_acceptance() { - for (status, body) in [ - (200, r#"{"job_id":"j1"}"#), - (202, "{}"), - (202, r#"{"job_id":""}"#), - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(status).body(body).unwrap() - }); - let error = db + async fn test_create_materialized_view_validates_before_submission() { + let db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("materialized-view creation must not send REST requests") + }); + for seconds in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let result = db .create_materialized_view_async(CreateMaterializedViewRequest { name: "adults".into(), namespace_path: vec![], query: "SELECT * FROM people".into(), with_no_data: false, auto_refresh_enabled: None, - auto_refresh_interval_seconds: None, + auto_refresh_interval_seconds: Some(seconds), }) - .await - .err() - .unwrap(); - assert!(error.to_string().contains(if status == 200 { - "202 Accepted" - } else { - "job_id" - })); + .await; + assert!(matches!(result, Err(Error::InvalidInput { .. }))); } } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index d8f04119f..846e852f5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::fs; -use std::future::{Future, pending}; +use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -12,6 +12,7 @@ use arrow_array::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; use arrow_flight::flight_service_client::FlightServiceClient; +#[cfg(test)] use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; use arrow_flight::{ Action, CancelFlightInfoRequest, CancelFlightInfoResult, CancelStatus, FlightClient, @@ -22,7 +23,7 @@ use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; use serde::{Serialize, de::DeserializeOwned}; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; @@ -226,6 +227,7 @@ impl SqlClient { .await } + #[cfg(test)] pub(super) async fn submit( &self, query: &str, @@ -254,40 +256,36 @@ impl SqlClient { .await } - pub(super) fn submit_as_job( + pub(super) async fn submit_as_job( &self, statement: String, namespace: Vec, finish: F, - ) -> Job + ) -> Result> where T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, F: FnOnce() -> Fut + Send + 'static, Fut: Future> + Send + 'static, { - let client = self.clone(); - let (query_tx, query_rx) = watch::channel(None::>); + let query = Arc::new( + self.execute(ExecuteQueryRequest { + query: statement, + default_namespace_path: namespace, + parameters: None, + }) + .await?, + ); + let query_id = query.id(); + let execution = query.clone(); let task = tokio::spawn(async move { - let query = Arc::new(client.submit(&statement, &namespace).await?); - query_tx.send_replace(Some(query.clone())); - let mut reader = query.reader().await?; + let mut reader = execution.reader().await?; while reader.try_next().await?.is_some() {} finish().await }); - Job::spawned_with_cancellation(task, move || { - let mut query_rx = query_rx.clone(); - async move { - loop { - let query = query_rx.borrow().clone(); - if let Some(query) = query { - return query.cancel().await; - } - if query_rx.changed().await.is_err() { - pending::<()>().await; - } - } - } - }) + Ok(Job::from_query_task(query_id, task, move || { + let query = query.clone(); + async move { query.cancel().await } + })) } pub(super) async fn describe(&self, query_id: Uuid) -> Result { @@ -315,6 +313,7 @@ impl SqlClientInner { ) } + #[cfg(test)] async fn poll( &self, descriptor: FlightDescriptor, diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 4837baab7..1368dd753 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -22,6 +22,7 @@ use tonic::{Request, Response, Status, Streaming}; use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; +use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; use crate::remote::client::test_utils::client_with_handler; @@ -75,6 +76,7 @@ struct TestSqlService { defer: bool, invalid_decision: bool, query_count: Arc, + statements: Arc>>, do_get_count: Arc, cancel_count: Arc, cancel_denied_count: Arc, @@ -207,6 +209,7 @@ impl Default for TestSqlService { defer: false, invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), + statements: Arc::new(std::sync::Mutex::new(Vec::new())), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), cancel_denied_count: Arc::new(AtomicUsize::new(0)), @@ -269,6 +272,7 @@ impl FlightService for TestSqlService { .and_then(|any| any.unpack::().ok().flatten()); let (query, stage) = if let Some(command) = command { self.query_count.fetch_add(1, Ordering::SeqCst); + self.statements.lock().unwrap().push(command.query.clone()); (command.query, 0_u8) } else { let continuation = std::str::from_utf8(request.get_ref().cmd.as_ref()) @@ -289,6 +293,9 @@ impl FlightService for TestSqlService { .ok_or_else(|| Status::invalid_argument("invalid continuation"))?; (query.to_string(), stage) }; + if query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } if (query == "SELECT slow" || query == "SELECT cancelling") && stage > 0 { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -364,7 +371,11 @@ impl FlightService for TestSqlService { .ok() .and_then(|envelope| Any::decode(envelope.value).ok()) .and_then(|command| command.unpack::().ok().flatten()); - if initial.is_some() { + if let Some(command) = &initial { + self.statements.lock().unwrap().push(command.query.clone()); + if command.query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } self.capture_headers(request.metadata()); assert_eq!( request.metadata().get(PROTOCOL_HEADER).unwrap(), @@ -376,7 +387,8 @@ impl FlightService for TestSqlService { .map(|command| command.query.as_bytes()) .unwrap_or(request.get_ref().ticket.as_ref()); let ticket = query; - let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"" + || ticket == b"SELECT job cancel denied"; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -458,7 +470,9 @@ impl FlightService for TestSqlService { { return Err(Status::not_found("query cancellation completed")); } - let status = if query == "SELECT cancel unspecified" { + let status = if query == "SELECT job cancel denied" { + CancelStatus::NotCancellable + } else if query == "SELECT cancel unspecified" { CancelStatus::Unspecified } else if query == "SELECT cancelling" { if self @@ -880,6 +894,132 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_create_uses_sql_and_preserves_options_and_namespace() { + for defer in [false, true] { + let service = TestSqlService { + adaptive: true, + defer, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let count = service.query_count.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("CREATE must use SQL, not REST: {}", request.url().path()); + }) + .with_sql_client(sql_client.clone()); + for (enabled, interval, expected) in [ + (None, None, ""), + (Some(true), None, " WITH (auto_refresh_enabled = true)"), + ( + Some(false), + Some(0.5), + " WITH (auto_refresh_enabled = false, auto_refresh_interval_seconds = 0.5)", + ), + ] { + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: enabled, + auto_refresh_interval_seconds: interval, + }) + .await + .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + let query_id = + Uuid::parse_str(job.id().expect("SQL-backed job must expose its ID")).unwrap(); + job.wait().await.unwrap(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + assert!(job.events(Default::default()).await.is_err()); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!( + sql_client.describe(query_id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + statements.lock().unwrap().last().unwrap(), + &format!("CREATE MATERIALIZED VIEW \"adults\"{expected} AS SELECT * FROM people") + ); + } + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(statements.lock().unwrap().len(), 3); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "team$reports") + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn query_backed_jobs_propagate_submission_and_execution_failures() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let rejected = client + .submit_as_job("SELECT rejected".into(), vec![], || async { Ok(()) }) + .await; + assert!(rejected.is_err()); + let finished = Arc::new(AtomicBool::new(false)); + let callback_finished = finished.clone(); + let job = client + .submit_as_job("SELECT fail".into(), vec![], move || async move { + callback_finished.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert_eq!(job.status().await.unwrap(), "failed"); + assert!(!finished.load(Ordering::SeqCst)); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn rejected_job_cancellation_does_not_abort_sql_completion() { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let do_get_count = service.do_get_count.clone(); + let release = service.gate_refresh_release.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let job = client + .submit_as_job("SELECT job cancel denied".into(), vec![], || async { + Ok(()) + }) + .await + .unwrap(); + wait_for(&do_get_count, 1, "job did not begin reading its SQL result").await; + assert!(matches!( + job.cancel().await, + Err(Error::NotSupported { .. }) + )); + assert_eq!(job.status().await.unwrap(), "running"); + release.notify_one(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + job.cancel().await.unwrap(); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn remote_refresh_twice_keeps_latest_selector() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -969,7 +1109,11 @@ async fn cancelling_remote_refresh_cancels_sql_query() { let (socket, _) = listener.accept().await?; Ok::<_, std::io::Error>(Some((socket, listener))) }); - let service = TestSqlService::default(); + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; let cancel_count = service.cancel_count.clone(); let do_get_count = service.do_get_count.clone(); let gate_refresh_release = service.gate_refresh_release.clone(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b722dc4a1..4cbba04df 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2247,8 +2247,8 @@ impl BaseTable for RemoteTable { let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); let job_sql_client = sql_client.clone(); - Ok( - sql_client.submit_as_job(statement, namespace.clone(), move || async move { + sql_client + .submit_as_job(statement, namespace.clone(), move || async move { let table = Self::new_with_sql_client( client, name, @@ -2270,8 +2270,8 @@ impl BaseTable for RemoteTable { source_version: 0, version, }) - }), - ) + }) + .await } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?;