mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 05:26:14 +00:00
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:
@@ -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
|
||||
|
||||
|
||||
@@ -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([]);
|
||||
});
|
||||
|
||||
@@ -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
@@ -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()
|
||||
|
||||
@@ -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]: ...
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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 }))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 { .. })));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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?;
|
||||
|
||||
Reference in New Issue
Block a user