fix: execute remote materialized view creation through SQL (#4665)

Remote materialized-view creation currently expects a REST response
containing a registry job ID, while refresh already submits SQL. Route
both through the SQL query backend and preserve Python's synchronous and
asynchronous return types.

### Query to Job mapping

`SqlClient::submit_as_job()` submits an `ExecuteQueryRequest` through
the same backend as `execute_query_async()`. Once submission succeeds,
it retains the returned `Query` and starts one client-side task that
consumes its result reader to completion. `Job::from_query_task()` wraps
that task, its cached outcome, a single ID and a cancellation callback.
Python exposes this as `Job[T]` or `AsyncJob[T]`.

The existing `job_type` property maps SQL-backed handles to the new
`"query"` kind. The shared `id` contains that query's UUID string, so
Python callers can inspect it with `db.describe_query(job.id)` on the
submitting connection. No separate `query_id`, backend enum or `type()`
API is added.

| Handle | `job_type` | `id` | Available lookup |
| --- | --- | --- | --- |
| SQL-backed | `"query"`, known at submission | Connection-scoped query
UUID string | `describe_query(job.id)` |
| Registry-backed | Existing server-reported kind, after `status()` or
`refresh()` | Registry job ID string | Registry job APIs, descriptions
and events |
| Local | `None` | `None` | The current in-process handle |

| Job concept | SQL-backed behavior |
| --- | --- |
| Submission | Await initial SQL submission before returning the handle.
Submission errors propagate immediately. |
| `status()` | Tracks the adapter task: running during result
consumption/finalization, finished after both succeed, failed if either
fails, or cancelled after cancellation. |
| `wait()` | Returns the operation's typed terminal result. One task
consumes the query reader once; repeated waits reuse its cached outcome
without resubmitting SQL or reopening the reader. |
| `cancel()` | Delegates to `Query.cancel()`. Rejected or unsupported
cancellation leaves the task running; an accepted request stops the
task. Cancelling a finished operation is a no-op. |
| Registry details | Query-backed handles expose no registry
description, event history, pause/resume or reattachment after a client
restart. `refresh()` updates cached lifecycle state. |

The Job tracks command completion rather than exposing SQL batches. A
wait timeout does not cancel the query. Query IDs remain scoped to the
submitting connection and cannot be used with registry APIs or as
durable resume tokens. Python `describe_query()` accepts both UUID
objects and the string `id` exposed by query-backed jobs. Rust, Python
and TypeScript expose the same classification through their existing
job-type accessors.

### Materialized-view operations

- **CREATE:** Render `CREATE MATERIALIZED VIEW` with the view namespace
and optional auto-refresh settings. After command completion, resolve
the job to `None`. Synchronous creation waits on the job, then opens and
returns the materialized view.
- **REFRESH:** Submit `REFRESH MATERIALIZED VIEW` through the same
adapter. After command completion, read the refreshed table's version
and row count to construct the existing typed refresh result.
Synchronous refresh waits on the same path.

Remote creation requires the SQL service. Execution placement remains a
server decision. REST endpoints are unchanged. The job guide documents
the ID/job-type contract and query-backed limitations.

Regression coverage exercises SQL submission, namespaces, refresh
options, job-type/ID preservation through typed result mapping, repeated
waits, failures and cancellation rejection. No local builds or tests
were run.
This commit is contained in:
Jack Ye
2026-10-25 01:20:13 -07:00
committed by GitHub
parent 24c2669a50
commit 9af7757053
20 changed files with 465 additions and 216 deletions
+8 -5
View File
@@ -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
+31 -6
View File
@@ -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.
@@ -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([]);
});
+8 -5
View File
@@ -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;
+6 -6
View File
@@ -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<String> {
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<String> {
self.inner.job_type()
+2
View File
@@ -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]: ...
+17 -8
View File
@@ -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:
+17 -13
View File
@@ -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]:
+1 -1
View File
@@ -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)
+1 -1
View File
@@ -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(
@@ -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
+63 -24
View File
@@ -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(
+17 -3
View File
@@ -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(
+5
View File
@@ -81,6 +81,11 @@ impl Job {
self.inner.id().map(str::to_string)
}
#[getter]
pub fn job_type(&self) -> Option<String> {
self.inner.job_type()
}
pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(
+63 -13
View File
@@ -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<T>(what: &str) -> Result<T> {
/// 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<String>;
async fn wait(&self) -> Result<TerminalResult>;
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<T = ()>
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<F, Fut>(task: JoinHandle<Result<T>>, cancel: F) -> Self
pub(crate) fn from_query_task<F, Fut>(
query_id: Uuid,
task: JoinHandle<Result<T>>,
cancel: F,
) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + 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<T> Job<T>
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<String> {
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<Option<Outcome>>,
abort: AbortHandle,
cancellation: Option<Cancellation>,
id: Option<String>,
}
type CancellationFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
@@ -549,11 +564,17 @@ impl SpawnedJob {
Self::new_inner(task, None)
}
fn new_with_cancellation<T>(task: JoinHandle<Result<T>>, cancellation: Cancellation) -> Self
fn new_with_cancellation<T>(
query_id: Uuid,
task: JoinHandle<Result<T>>,
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<T>(task: JoinHandle<Result<T>>, cancellation: Option<Cancellation>) -> 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<String> {
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::<Result<u64>>().await }))
+2 -1
View File
@@ -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<Job> {
if self.connection.uri().starts_with("db://") {
return self
+45 -99
View File
@@ -63,15 +63,6 @@ struct RemoteCloneTableRequest {
is_shallow: Option<bool>,
}
#[derive(serde::Serialize)]
struct RemoteCreateMaterializedViewRequest {
query: String,
#[serde(skip_serializing_if = "Option::is_none")]
auto_refresh_enabled: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
auto_refresh_interval_seconds: Option<f64>,
}
// 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<F, T>(handler: F, config: ClientConfig) -> Self
where
F: Fn(reqwest::Request) -> http::Response<T> + Send + Sync + 'static,
@@ -885,8 +881,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
) -> Result<Job> {
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<S: HttpSend> Database for RemoteDatabase<S> {
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<String> {
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 { .. })));
}
}
+22 -23
View File
@@ -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<T, F, Fut>(
pub(super) async fn submit_as_job<T, F, Fut>(
&self,
statement: String,
namespace: Vec<String>,
finish: F,
) -> Job<T>
) -> Result<Job<T>>
where
T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
{
let client = self.clone();
let (query_tx, query_rx) = watch::channel(None::<Arc<Query>>);
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<QueryDescription> {
@@ -315,6 +313,7 @@ impl SqlClientInner {
)
}
#[cfg(test)]
async fn poll(
&self,
descriptor: FlightDescriptor,
+148 -4
View File
@@ -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<AtomicUsize>,
statements: Arc<std::sync::Mutex<Vec<String>>>,
do_get_count: Arc<AtomicUsize>,
cancel_count: Arc<AtomicUsize>,
cancel_denied_count: Arc<AtomicUsize>,
@@ -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::<CommandStatementQuery>().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::<CommandStatementQuery>().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<String> {
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();
+4 -4
View File
@@ -2247,8 +2247,8 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
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<S: HttpSend> BaseTable for RemoteTable<S> {
source_version: 0,
version,
})
}),
)
})
.await
}
async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>> {
let description = self.describe().await?;