From d0e6c2a7c71d6e71b60b95122ae13a6edc83e785 Mon Sep 17 00:00:00 2001 From: Luca Falsina Date: Fri, 31 Jul 2026 17:51:37 +0200 Subject: [PATCH] feat: Add async REST scan planning poll and plan storage credentials Catalogs that return status=submitted from planTableScan can now be polled via GET .../plan/{plan-id}, with best-effort cancel and scan-scoped FileIO rebuilt from plan storage-credentials. Public RestCatalog.plan_scan still returns list[FileScanTask]. Co-authored-by: Cursor --- mkdocs/docs/api.md | 2 + mkdocs/docs/configuration.md | 4 + pyiceberg/catalog/rest/__init__.py | 207 ++++++++++++++++++--- pyiceberg/catalog/rest/scan_planning.py | 15 +- pyiceberg/exceptions.py | 8 + pyiceberg/table/__init__.py | 7 +- tests/catalog/test_scan_planning_models.py | 144 +++++++++++++- 7 files changed, 355 insertions(+), 32 deletions(-) diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index cc4290f2d9..945887e1c1 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -1695,6 +1695,8 @@ scan = table.scan( [task.file.file_path for task in scan.plan_files()] ``` +When the REST catalog is configured with `scan-planning-mode=server` and advertises the plan endpoint, `plan_files()` / `to_arrow()` use server-side scan planning. Catalogs that return async plans (`status=submitted`) are polled automatically; see [REST Catalog configuration](configuration.md#rest-catalog). + The low level API `plan_files` methods returns a set of tasks that provide the files that might contain matching rows: ```json diff --git a/mkdocs/docs/configuration.md b/mkdocs/docs/configuration.md index 44b5e395a9..6ba98924aa 100644 --- a/mkdocs/docs/configuration.md +++ b/mkdocs/docs/configuration.md @@ -386,6 +386,10 @@ catalog: | snapshot-loading-mode | refs | The snapshots to return in the body of the metadata. Setting the value to `all` would return the full set of snapshots currently valid for the table. Setting the value to `refs` would load all snapshots referenced by branches or tags. | | `header.X-Iceberg-Access-Delegation` | `vended-credentials` | Signal to the server that the client supports delegated access via a comma-separated list of access mechanisms. The server may choose to supply access via any or none of the requested mechanisms. When using `vended-credentials`, the server provides temporary credentials to the client. When using `remote-signing`, the server signs requests on behalf of the client. (default: `vended-credentials`) | | view-endpoints-supported | false | For backwards compatibility with older REST servers. Set to `true` if the server supports view endpoints but doesn't send the `endpoints` field in the ConfigResponse. | +| scan-planning-mode | client | When set to `server`, and the catalog advertises the plan-table-scan endpoint, `table.scan()` uses REST server-side scan planning. Async plans (`status=submitted`) are polled via `GET .../plan/{plan-id}` until completion. | +| rest-scan-planning.poll-timeout-ms | 300000 | Maximum time to wait for an async scan plan to complete before failing (default: 5 minutes). | + +When server-side planning returns `storage-credentials` on a completed plan, PyIceberg applies them to the scan-scoped FileIO (layered on top of the existing table/load-time IO properties) so planned data and delete files can be read. #### Headers in REST Catalog diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py index 97a71b429b..ef1ff56c5f 100644 --- a/pyiceberg/catalog/rest/__init__.py +++ b/pyiceberg/catalog/rest/__init__.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import time from collections import deque from enum import Enum from typing import ( @@ -38,6 +39,7 @@ PlanCancelled, PlanCompleted, PlanFailed, + PlannedScanResult, PlanningResponse, PlanSubmitted, PlanTableScanRequest, @@ -52,9 +54,11 @@ NamespaceNotEmptyError, NoSuchIdentifierError, NoSuchNamespaceError, + NoSuchPlanIdError, NoSuchPlanTaskError, NoSuchTableError, NoSuchViewError, + RemotePlanTimeoutError, TableAlreadyExistsError, UnauthorizedError, ViewAlreadyExistsError, @@ -160,6 +164,9 @@ class Endpoints: drop_view: str = "namespaces/{namespace}/views/{view}" view_exists: str = "namespaces/{namespace}/views/{view}" plan_table_scan: str = "namespaces/{namespace}/tables/{table}/plan" + # Use plan_id (underscore) for str.format; Capability paths use {plan-id} to match the REST spec. + fetch_planning_result: str = "namespaces/{namespace}/tables/{table}/plan/{plan_id}" + cancel_planning: str = "namespaces/{namespace}/tables/{table}/plan/{plan_id}" fetch_scan_tasks: str = "namespaces/{namespace}/tables/{table}/tasks" @@ -190,6 +197,13 @@ class Capability: V1_REGISTER_VIEW = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.register_view}") V1_DELETE_VIEW = Endpoint(http_method=HttpMethod.DELETE, path=f"{API_PREFIX}/{Endpoints.drop_view}") V1_SUBMIT_TABLE_SCAN_PLAN = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.plan_table_scan}") + # Spec advertises {plan-id}; must match ConfigResponse endpoint strings from servers. + V1_FETCH_TABLE_SCAN_PLAN = Endpoint( + http_method=HttpMethod.GET, path=f"{API_PREFIX}/namespaces/{{namespace}}/tables/{{table}}/plan/{{plan-id}}" + ) + V1_CANCEL_TABLE_SCAN_PLAN = Endpoint( + http_method=HttpMethod.DELETE, path=f"{API_PREFIX}/namespaces/{{namespace}}/tables/{{table}}/plan/{{plan-id}}" + ) V1_TABLE_SCAN_PLAN_TASKS = Endpoint(http_method=HttpMethod.POST, path=f"{API_PREFIX}/{Endpoints.fetch_scan_tasks}") @@ -264,6 +278,12 @@ class ScanPlanningMode(Enum): CUSTOM = "custom" SCAN_PLANNING_MODE = "scan-planning-mode" SCAN_PLANNING_MODE_DEFAULT = ScanPlanningMode.CLIENT.value +REST_SCAN_PLANNING_POLL_TIMEOUT_MS = "rest-scan-planning.poll-timeout-ms" +REST_SCAN_PLANNING_POLL_TIMEOUT_MS_DEFAULT = 5 * 60 * 1000 # 5 minutes, matches Java +REST_SCAN_PLANNING_POLL_MIN_SLEEP_MS = 1000 +REST_SCAN_PLANNING_POLL_MAX_SLEEP_MS = 60 * 1000 +REST_SCAN_PLANNING_POLL_SCALE_FACTOR = 2.0 +REST_SCAN_PLANNING_POLL_MAX_RETRIES = 10 # for backwards compatibility with older REST servers where it can be assumed that a particular # server supports view endpoints but doesn't send the "endpoints" field in the ConfigResponse VIEW_ENDPOINTS_SUPPORTED = "view-endpoints-supported" @@ -551,45 +571,115 @@ def _fetch_scan_tasks(self, identifier: str | Identifier, plan_task: str) -> Sca return ScanTasks.model_validate_json(response.text) - def plan_scan(self, identifier: str | Identifier, request: PlanTableScanRequest) -> list[FileScanTask]: - """Plan a table scan and return FileScanTasks. - - Handles the full scan planning lifecycle including pagination. + @retry(**_RETRY_ARGS) + def _fetch_planning_result(self, identifier: str | Identifier, plan_id: str) -> PlanningResponse: + """Fetch the result of an async scan plan by plan-id. Args: identifier: Table identifier. - request: The scan plan request parameters. + plan_id: Plan id returned from a submitted planTableScan response. Returns: - List of FileScanTask objects ready for execution. + PlanningResponse with the current plan status. Raises: - RuntimeError: If planning fails, is cancelled, or returns unexpected response. - NotImplementedError: If async planning is required but not yet supported. + NoSuchPlanIdError: If the plan-id does not exist. + NoSuchTableError: If the table does not exist. """ - response = self._plan_table_scan(identifier, request) + self._check_endpoint(Capability.V1_FETCH_TABLE_SCAN_PLAN) + response = self._session.get( + self.url( + Endpoints.fetch_planning_result, + prefixed=True, + plan_id=quote(plan_id, safe=""), + **self._split_identifier_for_path(identifier), + ), + ) + try: + response.raise_for_status() + except HTTPError as exc: + _handle_non_200_response(exc, {404: NoSuchPlanIdError}) - if isinstance(response, PlanFailed): - error_msg = response.error.message if response.error else "unknown error" - raise RuntimeError(f"Received status: failed: {error_msg}") + return _PLANNING_RESPONSE_ADAPTER.validate_json(response.text) - if isinstance(response, PlanCancelled): - raise RuntimeError("Received status: cancelled") + def _cancel_planning(self, identifier: str | Identifier, plan_id: str) -> bool: + """Best-effort cancel of an async scan plan. - if isinstance(response, PlanSubmitted): - # TODO: implement polling for async planning - raise NotImplementedError(f"Async scan planning not yet supported for planId: {response.plan_id}") + Returns: + True if the cancel request was accepted, False otherwise. + """ + if Capability.V1_CANCEL_TABLE_SCAN_PLAN not in self._supported_endpoints: + return False - if not isinstance(response, PlanCompleted): - raise RuntimeError(f"Invalid planStatus for response: {type(response).__name__}") + try: + response = self._session.delete( + self.url( + Endpoints.cancel_planning, + prefixed=True, + plan_id=quote(plan_id, safe=""), + **self._split_identifier_for_path(identifier), + ), + ) + response.raise_for_status() + return True + except Exception: + # Plan may have already completed, failed, or been cancelled. + return False + + def _poll_until_completed(self, identifier: str | Identifier, plan_id: str) -> PlanCompleted: + """Poll fetchPlanningResult until the plan completes or times out. + + Uses exponential backoff matching Java RESTTableScan defaults. + """ + max_wait_ms = property_as_int( + self.properties, + REST_SCAN_PLANNING_POLL_TIMEOUT_MS, + REST_SCAN_PLANNING_POLL_TIMEOUT_MS_DEFAULT, + ) + if max_wait_ms is None or max_wait_ms <= 0: + raise ValueError(f"Invalid value for {REST_SCAN_PLANNING_POLL_TIMEOUT_MS}: {max_wait_ms} (must be positive)") + + sleep_ms = float(REST_SCAN_PLANNING_POLL_MIN_SLEEP_MS) + start = time.monotonic() + retries = 0 + + while True: + response = self._fetch_planning_result(identifier, plan_id) + + if isinstance(response, PlanCompleted): + return response + + if isinstance(response, PlanFailed): + error_msg = response.error.message if response.error else "unknown error" + self._cancel_planning(identifier, plan_id) + raise RuntimeError(f"Remote scan planning failed for planId: {plan_id}: {error_msg}") + + if isinstance(response, PlanCancelled): + raise RuntimeError(f"Remote scan planning cancelled for planId: {plan_id}") + + if not isinstance(response, PlanSubmitted): + self._cancel_planning(identifier, plan_id) + raise RuntimeError(f"Invalid planStatus for planId: {plan_id}: {type(response).__name__}") + + elapsed_ms = (time.monotonic() - start) * 1000 + if retries >= REST_SCAN_PLANNING_POLL_MAX_RETRIES or elapsed_ms >= max_wait_ms: + self._cancel_planning(identifier, plan_id) + raise RemotePlanTimeoutError( + f"Remote scan planning for planId: {plan_id} did not complete within configured limits " + f"(timeout={max_wait_ms} ms, maxRetries={REST_SCAN_PLANNING_POLL_MAX_RETRIES})" + ) + time.sleep(sleep_ms / 1000.0) + sleep_ms = min(sleep_ms * REST_SCAN_PLANNING_POLL_SCALE_FACTOR, REST_SCAN_PLANNING_POLL_MAX_SLEEP_MS) + retries += 1 + + def _expand_plan_tasks(self, identifier: str | Identifier, response: PlanCompleted) -> list[FileScanTask]: + """Expand a completed plan response into FileScanTask objects, including pagination.""" tasks: list[FileScanTask] = [] - # Collect tasks from initial response for task in response.file_scan_tasks: tasks.append(FileScanTask.from_rest_response(task, response.delete_files)) - # Fetch and collect from additional batches pending_tasks = deque(response.plan_tasks) while pending_tasks: plan_task = pending_tasks.popleft() @@ -600,6 +690,81 @@ def plan_scan(self, identifier: str | Identifier, request: PlanTableScanRequest) return tasks + def _plan_scan_result(self, identifier: str | Identifier, request: PlanTableScanRequest) -> PlannedScanResult: + """Plan a table scan and return tasks with optional plan storage credentials. + + Handles the full scan planning lifecycle including async polling and pagination. + """ + response = self._plan_table_scan(identifier, request) + + if isinstance(response, PlanFailed): + error_msg = response.error.message if response.error else "unknown error" + raise RuntimeError(f"Received status: failed: {error_msg}") + + if isinstance(response, PlanCancelled): + raise RuntimeError("Received status: cancelled") + + if isinstance(response, PlanSubmitted): + if not response.plan_id: + raise ValueError("Async scan planning submitted without plan-id") + response = self._poll_until_completed(identifier, response.plan_id) + + if not isinstance(response, PlanCompleted): + raise RuntimeError(f"Invalid planStatus for response: {type(response).__name__}") + + tasks = self._expand_plan_tasks(identifier, response) + return PlannedScanResult( + tasks=tasks, + storage_credentials=list(response.storage_credentials or []), + plan_id=response.plan_id, + ) + + def plan_scan(self, identifier: str | Identifier, request: PlanTableScanRequest) -> list[FileScanTask]: + """Plan a table scan and return FileScanTasks. + + Handles the full scan planning lifecycle including async polling and pagination. + + Args: + identifier: Table identifier. + request: The scan plan request parameters. + + Returns: + List of FileScanTask objects ready for execution. + + Raises: + RuntimeError: If planning fails, is cancelled, or returns unexpected response. + RemotePlanTimeoutError: If async planning does not complete in time. + ValueError: If a submitted plan is missing plan-id. + """ + return self._plan_scan_result(identifier, request).tasks + + def _file_io_from_plan( + self, + existing_properties: Properties, + storage_credentials: list[StorageCredential], + location: str | None = None, + ) -> FileIO | None: + """Build a scan-scoped FileIO from plan storage credentials. + + Layers resolved plan credentials on top of the existing scan FileIO properties so + load-time settings (for example custom S3 endpoints) are retained. + """ + if not storage_credentials: + return None + + resolve_location = location + if resolve_location is None and storage_credentials[0].prefix: + resolve_location = storage_credentials[0].prefix + + credential_config = self._resolve_storage_credentials(storage_credentials, resolve_location) + if not credential_config and resolve_location is None: + credential_config = dict(storage_credentials[0].config) + + if not credential_config: + return None + + return self._load_file_io({**existing_properties, **credential_config}, resolve_location) + def _create_legacy_oauth2_auth_manager(self, session: Session) -> AuthManager: """Create the LegacyOAuth2AuthManager by fetching required properties. diff --git a/pyiceberg/catalog/rest/scan_planning.py b/pyiceberg/catalog/rest/scan_planning.py index 1a2204341b..5aa70bc5a8 100644 --- a/pyiceberg/catalog/rest/scan_planning.py +++ b/pyiceberg/catalog/rest/scan_planning.py @@ -16,9 +16,10 @@ # under the License. from __future__ import annotations +from dataclasses import dataclass, field from datetime import date, datetime, time from decimal import Decimal -from typing import Annotated, Generic, Literal, TypeAlias, TypeVar +from typing import TYPE_CHECKING, Annotated, Generic, Literal, TypeAlias, TypeVar from uuid import UUID from pydantic import Field, model_validator @@ -28,6 +29,9 @@ from pyiceberg.manifest import FileFormat from pyiceberg.typedef import IcebergBaseModel +if TYPE_CHECKING: + from pyiceberg.table import FileScanTask + # Primitive types that can appear in partition values and bounds PrimitiveTypeValue: TypeAlias = bool | int | float | str | Decimal | UUID | date | time | datetime | bytes @@ -207,3 +211,12 @@ class FetchScanTasksRequest(IcebergBaseModel): """Request body for fetching scan tasks endpoint.""" plan_task: str = Field(alias="plan-task") + + +@dataclass(frozen=True) +class PlannedScanResult: + """Result of REST server-side scan planning, including optional storage credentials.""" + + tasks: list[FileScanTask] + storage_credentials: list[StorageCredential] = field(default_factory=list) + plan_id: str | None = None diff --git a/pyiceberg/exceptions.py b/pyiceberg/exceptions.py index 4a22f1bb21..019ccff894 100644 --- a/pyiceberg/exceptions.py +++ b/pyiceberg/exceptions.py @@ -60,6 +60,14 @@ class NoSuchPlanTaskError(Exception): """Raised when a scan plan task is not found.""" +class NoSuchPlanIdError(Exception): + """Raised when a scan plan-id is not found.""" + + +class RemotePlanTimeoutError(Exception): + """Raised when async remote scan planning does not complete within configured limits.""" + + class RESTError(Exception): """Raises when there is an unknown response from the REST Catalog.""" diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index c0adce84dc..1cbbb03b66 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -2274,7 +2274,12 @@ def _plan_files_server_side(self) -> Iterable[FileScanTask]: case_sensitive=self.case_sensitive, ) - return self.catalog.plan_scan(self.table_identifier, request) + result = self.catalog._plan_scan_result(self.table_identifier, request) + location = result.tasks[0].file.file_path if result.tasks else None + plan_io = self.catalog._file_io_from_plan(self.io.properties, result.storage_credentials, location) + if plan_io is not None: + self.io = plan_io + return result.tasks def _plan_files_local(self) -> Iterable[FileScanTask]: """Plan files locally by reading manifests.""" diff --git a/tests/catalog/test_scan_planning_models.py b/tests/catalog/test_scan_planning_models.py index f2c80cfb9b..a23a2aa5ef 100644 --- a/tests/catalog/test_scan_planning_models.py +++ b/tests/catalog/test_scan_planning_models.py @@ -35,6 +35,7 @@ RESTFileScanTask, RESTPositionDeleteFile, ScanTasks, + StorageCredential, ValueMap, ) from pyiceberg.expressions import AlwaysTrue, EqualTo, Reference @@ -52,6 +53,8 @@ def rest_scan_catalog(requests_mock: Mocker) -> RestCatalog: "overrides": {}, "endpoints": [ "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan", + "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", + "DELETE /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/tasks", ], }, @@ -431,7 +434,7 @@ def test_plan_scan_completed_single_batch(rest_scan_catalog: RestCatalog, reques ) request = PlanTableScanRequest() - tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = rest_scan_catalog.plan_scan(("db", "tbl"), request) assert len(tasks) == 1 assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" @@ -465,7 +468,7 @@ def test_plan_scan_with_pagination(rest_scan_catalog: RestCatalog, requests_mock request = PlanTableScanRequest() - tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = rest_scan_catalog.plan_scan(("db", "tbl"), request) assert len(tasks) == 2 assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" @@ -493,14 +496,19 @@ def test_plan_scan_with_delete_files(rest_scan_catalog: RestCatalog, requests_mo ) request = PlanTableScanRequest() - tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = rest_scan_catalog.plan_scan(("db", "tbl"), request) assert len(tasks) == 1 assert tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" assert len(tasks[0].delete_files) == 1 -def test_plan_scan_async_not_supported(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: +def test_plan_scan_async_poll_completes( + rest_scan_catalog: RestCatalog, requests_mock: Mocker, monkeypatch: pytest.MonkeyPatch +) -> None: + file_one = _rest_data_file(file_path="s3://bucket/tbl/data/file1.parquet") + monkeypatch.setattr("time.sleep", lambda _s: None) + requests_mock.post( f"{TEST_URI}v1/namespaces/db/tables/tbl/plan", json={ @@ -509,10 +517,128 @@ def test_plan_scan_async_not_supported(rest_scan_catalog: RestCatalog, requests_ }, status_code=200, ) + requests_mock.get( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-456", + [ + {"json": {"status": "submitted", "plan-id": "plan-456"}, "status_code": 200}, + { + "json": { + "status": "completed", + "plan-id": "plan-456", + "delete-files": [], + "file-scan-tasks": [{"data-file": file_one}], + "plan-tasks": [], + "storage-credentials": [ + { + "prefix": "s3://bucket/tbl/data", + "config": { + "s3.access-key-id": "plan-key", + "s3.secret-access-key": "plan-secret", + }, + } + ], + }, + "status_code": 200, + }, + ], + ) + + request = PlanTableScanRequest() + result = rest_scan_catalog._plan_scan_result(("db", "tbl"), request) + + assert len(result.tasks) == 1 + assert result.tasks[0].file.file_path == "s3://bucket/tbl/data/file1.parquet" + assert result.plan_id == "plan-456" + assert len(result.storage_credentials) == 1 + assert result.storage_credentials[0].config["s3.access-key-id"] == "plan-key" + + +def test_plan_storage_credentials_keep_scan_io_properties(rest_scan_catalog: RestCatalog) -> None: + existing_properties = {"s3.endpoint": "https://minio:9000", "s3.access-key-id": "table-key"} + storage_credentials = [ + StorageCredential(prefix="s3://bucket/tbl/data", config={"s3.access-key-id": "plan-key"}), + ] + + plan_io = rest_scan_catalog._file_io_from_plan( + existing_properties, + storage_credentials, + "s3://bucket/tbl/data/file1.parquet", + ) + + assert plan_io is not None + # Plan credentials take precedence, but the load-time IO configuration is retained. + assert plan_io.properties["s3.access-key-id"] == "plan-key" + assert plan_io.properties["s3.endpoint"] == "https://minio:9000" + + +def test_plan_scan_async_timeout(rest_scan_catalog: RestCatalog, requests_mock: Mocker, monkeypatch: pytest.MonkeyPatch) -> None: + from pyiceberg.exceptions import RemotePlanTimeoutError + + requests_mock.post( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan", + json={"status": "submitted", "plan-id": "plan-timeout"}, + status_code=200, + ) + requests_mock.get( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-timeout", + json={"status": "submitted", "plan-id": "plan-timeout"}, + status_code=200, + ) + requests_mock.delete( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-timeout", + status_code=204, + ) + + import pyiceberg.catalog.rest as rest_module + + monkeypatch.setitem(rest_scan_catalog.properties, "rest-scan-planning.poll-timeout-ms", "1") + monkeypatch.setattr(rest_module, "REST_SCAN_PLANNING_POLL_MAX_RETRIES", 0) + monkeypatch.setattr("time.sleep", lambda _s: None) + + request = PlanTableScanRequest() + with pytest.raises(RemotePlanTimeoutError, match="plan-timeout"): + rest_scan_catalog.plan_scan(("db", "tbl"), request) + + +def test_plan_scan_async_failed(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: + requests_mock.post( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan", + json={"status": "submitted", "plan-id": "plan-fail"}, + status_code=200, + ) + requests_mock.get( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-fail", + json={ + "status": "failed", + "error": {"message": "planning blew up", "type": "PlanningError", "code": 500}, + }, + status_code=200, + ) + requests_mock.delete( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-fail", + status_code=204, + ) + + request = PlanTableScanRequest() + with pytest.raises(RuntimeError, match="planning blew up"): + rest_scan_catalog.plan_scan(("db", "tbl"), request) + + +def test_plan_scan_async_cancelled(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: + requests_mock.post( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan", + json={"status": "submitted", "plan-id": "plan-cancel"}, + status_code=200, + ) + requests_mock.get( + f"{TEST_URI}v1/namespaces/db/tables/tbl/plan/plan-cancel", + json={"status": "cancelled"}, + status_code=200, + ) request = PlanTableScanRequest() - with pytest.raises(NotImplementedError, match="Async scan planning not yet supported"): - list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + with pytest.raises(RuntimeError, match="cancelled"): + rest_scan_catalog.plan_scan(("db", "tbl"), request) def test_plan_scan_empty_result(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: @@ -529,7 +655,7 @@ def test_plan_scan_empty_result(rest_scan_catalog: RestCatalog, requests_mock: M ) request = PlanTableScanRequest() - tasks = list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + tasks = rest_scan_catalog.plan_scan(("db", "tbl"), request) assert len(tasks) == 0 @@ -542,7 +668,7 @@ def test_plan_scan_cancelled(rest_scan_catalog: RestCatalog, requests_mock: Mock request = PlanTableScanRequest() with pytest.raises(RuntimeError, match="Received status: cancelled"): - list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + rest_scan_catalog.plan_scan(("db", "tbl"), request) def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog, requests_mock: Mocker) -> None: @@ -567,4 +693,4 @@ def test_plan_scan_equality_deletes_not_supported(rest_scan_catalog: RestCatalog request = PlanTableScanRequest() with pytest.raises(NotImplementedError, match="PyIceberg does not yet support equality deletes"): - list(rest_scan_catalog.plan_scan(("db", "tbl"), request)) + rest_scan_catalog.plan_scan(("db", "tbl"), request)