remove internal tasks operations from cloud public API (#11447)

Author: andrea-reale

Cargo.lock

@@ -10548,7 +10548,6 @@ dependencies = [
  "mimalloc",
  "numpy",
  "parking_lot",
- "prost-types",
  "pyo3",
  "pyo3-build-config",
  "rand 0.8.5",

crates/store/re_protos/proto/rerun/v1alpha1/cloud.proto

@@ -154,9 +154,6 @@ service RerunCloudService {
   // Query the status of submitted tasks
   rpc QueryTasks(QueryTasksRequest) returns (QueryTasksResponse) {}
 
-  // Fetch the output of a completed task
-  rpc FetchTaskOutput(FetchTaskOutputRequest) returns (FetchTaskOutputResponse) {}
-
   // Query the status of submitted tasks as soon as they are no longer pending
   rpc QueryTasksOnCompletion(QueryTasksOnCompletionRequest) returns (stream QueryTasksOnCompletionResponse) {}
 
@@ -640,27 +637,6 @@ message DoGlobalMaintenanceResponse {}
 
 // --- Tasks ---
 
-// A task is a unit of work that can be submitted to the system
-message Task {
-  // Unique identifier for the task
-  rerun.common.v1alpha1.TaskId id = 1;
-  // Type of the task
-  string task_type = 2;
-  // Task-type dependant data necessary to de-serialize the task
-  bytes task_data = 3;
-}
-
-// `SubmitTasksRequest` is the request message for submitting tasks
-message SubmitTasksRequest {
-  repeated Task tasks = 1;
-}
-
-// `SubmitTaskResponse` contains, for each submitted task
-// its submission outcome, encoded as a `RecordBatch`
-message SubmitTasksResponse {
-  rerun.common.v1alpha1.DataframePart data = 1;
-}
-
 // `QueryTasksRequest` is the request message for querying tasks status
 message QueryTasksRequest {
   // Empty queries for all tasks if the server allows it.
@@ -689,18 +665,6 @@ message QueryTasksOnCompletionResponse {
   rerun.common.v1alpha1.DataframePart data = 1;
 }
 
-// `FetchTaskOutputRequest` is the request message for fetching task output
-message FetchTaskOutputRequest {
-  // Unique identifier for the task
-  rerun.common.v1alpha1.TaskId id = 1;
-}
-
-/// `FetchTaskOutputResponse` is the response message for fetching task output
-message FetchTaskOutputResponse {
-  // The output of the task, encoded as a record batch
-  rerun.common.v1alpha1.DataframePart data = 1;
-}
-
 // --- Catalog ---
 
 // FindEntries

crates/store/re_protos/src/v1alpha1/rerun.cloud.v1alpha1.ext.rs

@@ -23,7 +23,7 @@ use crate::common::v1alpha1::{
     ComponentDescriptor, DataframePart, TaskId,
     ext::{DatasetHandle, IfDuplicateBehavior, PartitionId},
 };
-use crate::{TypeConversionError, missing_field};
+use crate::{TypeConversionError, invalid_field, missing_field};
 
 // --- RegisterWithDatasetRequest ---
 
@@ -1605,6 +1605,57 @@ impl From<ComponentColumnDescriptor> for crate::cloud::v1alpha1::IndexColumn {
     }
 }
 
+// --- Tasks ---
+
+pub struct QueryTasksOnCompletionRequest {
+    pub task_ids: Vec<TaskId>,
+    pub timeout: std::time::Duration,
+}
+
+pub struct QueryTasksRequest {
+    pub task_ids: Vec<TaskId>,
+}
+
+impl TryFrom<QueryTasksOnCompletionRequest>
+    for crate::cloud::v1alpha1::QueryTasksOnCompletionRequest
+{
+    type Error = TypeConversionError;
+
+    fn try_from(
+        value: QueryTasksOnCompletionRequest,
+    ) -> Result<crate::cloud::v1alpha1::QueryTasksOnCompletionRequest, Self::Error> {
+        if value.task_ids.is_empty() {
+            return Err(missing_field!(
+                crate::cloud::v1alpha1::QueryTasksOnCompletionRequest,
+                "task_ids"
+            ));
+        }
+        let timeout: prost_types::Duration = value.timeout.try_into().map_err(|err| {
+            invalid_field!(
+                crate::cloud::v1alpha1::QueryTasksOnCompletionRequest,
+                "timeout",
+                err
+            )
+        })?;
+        Ok(Self {
+            ids: value.task_ids,
+            timeout: Some(timeout),
+        })
+    }
+}
+
+impl TryFrom<QueryTasksRequest> for crate::cloud::v1alpha1::QueryTasksRequest {
+    type Error = TypeConversionError;
+
+    fn try_from(
+        value: QueryTasksRequest,
+    ) -> Result<crate::cloud::v1alpha1::QueryTasksRequest, Self::Error> {
+        Ok(Self {
+            ids: value.task_ids,
+        })
+    }
+}
+
 // --
 
 #[cfg(test)]

crates/store/re_protos/src/v1alpha1/rerun.cloud.v1alpha1.rs

@@ -11,15 +11,16 @@ use re_protos::{
         CreateDatasetEntryRequest, DeleteEntryRequest, EntryFilter, EntryKind, FetchChunksRequest,
         FindEntriesRequest, GetDatasetManifestSchemaRequest, GetDatasetManifestSchemaResponse,
         GetPartitionTableSchemaRequest, GetPartitionTableSchemaResponse, QueryDatasetRequest,
-        QueryDatasetResponse, ReadDatasetEntryRequest, ReadTableEntryRequest,
-        RegisterWithDatasetResponse, ScanPartitionTableRequest, ScanPartitionTableResponse,
+        QueryDatasetResponse, QueryTasksOnCompletionResponse, QueryTasksResponse,
+        ReadDatasetEntryRequest, ReadTableEntryRequest, RegisterWithDatasetResponse,
+        ScanPartitionTableRequest, ScanPartitionTableResponse,
         ext::{
             CreateDatasetEntryResponse, DataSource, DataSourceKind, DatasetDetails, DatasetEntry,
             EntryDetails, EntryDetailsUpdate, LanceTable, ProviderDetails as _,
-            ReadDatasetEntryResponse, ReadTableEntryResponse, RegisterTableResponse,
-            RegisterWithDatasetRequest, RegisterWithDatasetTaskDescriptor, TableEntry,
-            UpdateDatasetEntryRequest, UpdateDatasetEntryResponse, UpdateEntryRequest,
-            UpdateEntryResponse,
+            QueryTasksOnCompletionRequest, QueryTasksRequest, ReadDatasetEntryResponse,
+            ReadTableEntryResponse, RegisterTableResponse, RegisterWithDatasetRequest,
+            RegisterWithDatasetTaskDescriptor, TableEntry, UpdateDatasetEntryRequest,
+            UpdateDatasetEntryResponse, UpdateEntryRequest, UpdateEntryResponse,
         },
         rerun_cloud_service_client::RerunCloudServiceClient,
     },
@@ -32,7 +33,7 @@ use re_protos::{
     missing_field,
 };
 
-use crate::{StreamEntryError, StreamError};
+use crate::{StreamEntryError, StreamError, StreamTasksError};
 
 pub type FetchChunksResponseStream = std::pin::Pin<
     Box<
@@ -548,4 +549,34 @@ where
             .map(|entry| entry.name.clone())
             .collect())
     }
+
+    // -- Tasks API --
+    pub async fn query_tasks_on_completion(
+        &mut self,
+        task_ids: Vec<TaskId>,
+        timeout: std::time::Duration,
+    ) -> Result<tonic::Streaming<QueryTasksOnCompletionResponse>, StreamError> {
+        let q = QueryTasksOnCompletionRequest { task_ids, timeout };
+        let response = self
+            .inner()
+            .query_tasks_on_completion(tonic::Request::new(q.try_into()?))
+            .await
+            .map_err(|err| StreamTasksError::StreamingTaskResults(err.into()))?
+            .into_inner();
+        Ok(response)
+    }
+
+    pub async fn query_tasks(
+        &mut self,
+        task_ids: Vec<TaskId>,
+    ) -> Result<QueryTasksResponse, StreamError> {
+        let q = QueryTasksRequest { task_ids };
+        let response = self
+            .inner()
+            .query_tasks(tonic::Request::new(q.try_into()?))
+            .await
+            .map_err(|err| StreamTasksError::StreamingTaskResults(err.into()))?
+            .into_inner();
+        Ok(response)
+    }
 }