PR comments

Author: rambleraptor

pyiceberg/catalog/rest/planning_models.py

@@ -22,6 +22,7 @@
 from pydantic import Field
 
 from pyiceberg.catalog.rest.expression import Expression
+from pyiceberg.catalog.rest.response import ErrorResponse as IcebergErrorResponse
 from pyiceberg.typedef import IcebergBaseModel, IcebergRootModel
 
 
@@ -345,3 +346,37 @@ class ScanTasks(IcebergBaseModel):
     )
     file_scan_tasks: Optional[List[FileScanTask]] = Field(None, alias="file-scan-tasks")
     plan_tasks: Optional[List[PlanTask]] = Field(None, alias="plan-tasks")
+
+
+class FailedPlanningResult(IcebergErrorResponse):
+    """Failed server-side planning result."""
+
+    status: Literal["failed"]
+
+
+class AsyncPlanningResult(IcebergBaseModel):
+    status: Literal["submitted"]
+    plan_id: str = Field(..., alias="plan-id", description="ID used to track a planning request")
+
+
+class CancelledPlanningResult(IcebergBaseModel):
+    """A cancelled planning result."""
+
+    status: Literal["cancelled"]
+
+
+class CompletedPlanningWithIDResult(ScanTasks):
+    """Completed server-side planning result."""
+
+    status: Literal["completed"]
+    plan_id: Optional[str] = Field(None, alias="plan-id", description="ID used to track a planning request")
+
+
+class PlanTableScanResult(
+    IcebergRootModel[Union[CompletedPlanningWithIDResult, FailedPlanningResult, AsyncPlanningResult, CancelledPlanningResult]]
+):
+    """Result of server-side scan planning for planTableScan."""
+
+    root: Union[CompletedPlanningWithIDResult, FailedPlanningResult, AsyncPlanningResult, CancelledPlanningResult] = Field(
+        ..., discriminator="status"
+    )

tests/catalog/test_rest_serializers.py

@@ -14,21 +14,29 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import json
+
 from pyiceberg.catalog.rest.expression import (
     AndOrExpression,
     Expression,
     LiteralExpression,
     Term,
 )
 from pyiceberg.catalog.rest.planning_models import (
+    AsyncPlanningResult,
+    CancelledPlanningResult,
+    CompletedPlanningWithIDResult,
     DataFile,
     DeleteFile,
     EqualityDeleteFile,
+    FailedPlanningResult,
     FileScanTask,
     PlanTableScanRequest,
+    PlanTableScanResult,
     PositionDeleteFile,
     ScanTasks,
 )
+from pyiceberg.catalog.rest.response import ErrorResponseMessage
 
 
 def test_serialize_plan_table_scan_request() -> None:
@@ -165,3 +173,118 @@ def snapshot_json_for_plan_table_scan_request() -> str:
 
 def snapshot_json_for_scan_tasks() -> str:
     return """{"delete-files":[{"content":"position-deletes","file-path":"/path/to/delete-a.parquet","file-format":"parquet","spec-id":0,"partition":[],"file-size-in-bytes":256,"record-count":10},{"content":"equality-deletes","file-path":"/path/to/delete-b.parquet","file-format":"parquet","spec-id":0,"partition":[],"file-size-in-bytes":256,"record-count":10,"equality-ids":[1,2]}],"file-scan-tasks":[{"data-file":{"content":"data","file-path":"/path/to/data-a.parquet","file-format":"parquet","spec-id":0,"partition":[],"file-size-in-bytes":1024,"record-count":56},"delete-file-references":[0,1]}]}"""
+
+
+def test_deserialize_async_planning_result() -> None:
+    """Test deserializing a dict to an AsyncPlanningResult"""
+    result = PlanTableScanResult.model_validate_json(snapshot_json_for_async_planning_result())
+    expected = AsyncPlanningResult(status="submitted", plan_id="plan-123")
+    # Assert that deserialized dict == Python object
+    assert result.root == expected
+
+
+def test_serialize_async_planning_result() -> None:
+    """Test serializing an AsyncPlanningResult to a dict"""
+    result = PlanTableScanResult(root=AsyncPlanningResult(status="submitted", plan_id="plan-123"))
+    # Assert that JSON matches
+    assert json.loads(result.model_dump_json(by_alias=True)) == json.loads(snapshot_json_for_async_planning_result())
+
+
+def snapshot_json_for_async_planning_result() -> str:
+    return """{"status":"submitted","plan-id":"plan-123"}"""
+
+
+def test_deserialize_failed_planning_result() -> None:
+    """Test deserializing a dict to a FailedPlanningResult"""
+    result = PlanTableScanResult.model_validate_json(snapshot_json_for_failed_planning_result())
+    expected = FailedPlanningResult(
+        status="failed",
+        error=ErrorResponseMessage(
+            message="The plan is invalid",
+            type="NoSuchPlanException",
+            code=404,
+        ),
+    )
+    # Assert that deserialized dict == Python object
+    assert result.root == expected
+
+
+def test_serialize_failed_planning_result() -> None:
+    """Test serializing a FailedPlanningResult to a dict"""
+    result = PlanTableScanResult(
+        root=FailedPlanningResult(
+            status="failed",
+            error=ErrorResponseMessage(
+                message="The plan is invalid",
+                type="NoSuchPlanException",
+                code=404,
+            ),
+        )
+    )
+    # Assert that JSON matches
+    assert json.loads(result.model_dump_json(by_alias=True, exclude_none=True)) == json.loads(
+        snapshot_json_for_failed_planning_result()
+    )
+
+
+def snapshot_json_for_failed_planning_result() -> str:
+    return """{"status":"failed","error":{"message":"The plan is invalid","type":"NoSuchPlanException","code":404}}"""
+
+
+def test_deserialize_cancelled_planning_result() -> None:
+    """Test deserializing a dict to an CancelledPlanningResult"""
+    result = PlanTableScanResult.model_validate_json(snapshot_json_for_cancelled_planning_result())
+    expected = CancelledPlanningResult(status="cancelled")
+    # Assert that deserialized dict == Python object
+    assert result.root == expected
+
+
+def test_serialize_cancelled_planning_result() -> None:
+    """Test serializing an CancelledPlanningResult to a dict"""
+    result = PlanTableScanResult(root=CancelledPlanningResult(status="cancelled"))
+    # Assert that JSON matches
+    assert json.loads(result.model_dump_json(by_alias=True)) == json.loads(snapshot_json_for_cancelled_planning_result())
+
+
+def snapshot_json_for_cancelled_planning_result() -> str:
+    return """{"status":"cancelled"}"""
+
+
+def test_deserialize_completed_planning_with_id_result() -> None:
+    """Test deserializing a dict to a CompletedPlanningWithIDResult"""
+    scan_tasks_dict = json.loads(snapshot_json_for_scan_tasks())
+    scan_tasks_dict["status"] = "completed"
+    scan_tasks_dict["plan-id"] = "plan-456"
+    json_str = json.dumps(scan_tasks_dict)
+
+    result = PlanTableScanResult.model_validate_json(json_str)
+    expected_scan_tasks = ScanTasks.model_validate_json(snapshot_json_for_scan_tasks())
+
+    expected = CompletedPlanningWithIDResult(
+        status="completed",
+        plan_id="plan-456",
+        file_scan_tasks=expected_scan_tasks.file_scan_tasks,
+        delete_files=expected_scan_tasks.delete_files,
+    )
+    # Assert that deserialized dict == Python object
+    assert result.root == expected
+
+
+def test_serialize_completed_planning_with_id_result() -> None:
+    """Test serializing a CompletedPlanningWithIDResult to a dict"""
+    expected_scan_tasks = ScanTasks.model_validate_json(snapshot_json_for_scan_tasks())
+    result = PlanTableScanResult(
+        root=CompletedPlanningWithIDResult(
+            status="completed",
+            plan_id="plan-456",
+            file_scan_tasks=expected_scan_tasks.file_scan_tasks,
+            delete_files=expected_scan_tasks.delete_files,
+        )
+    )
+
+    scan_tasks_dict = json.loads(snapshot_json_for_scan_tasks())
+    scan_tasks_dict["status"] = "completed"
+    scan_tasks_dict["plan-id"] = "plan-456"
+
+    # Assert that JSON matches
+    assert json.loads(result.model_dump_json(exclude_none=True, by_alias=True)) == scan_tasks_dict