End-to-end out-of-band entry headers: datasets (part 2) (#11411)

With this second part, everything related to the manifest registry is
now fully migrated.
As usual, the interesting stuff happens in the sibling.

* Sibling: rerun-io/dataplatform#1796

Author: teh-cmc

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

@@ -46,7 +46,9 @@ service RerunCloudService {
 
   /* Write data */
 
-  // Register new partitions with the Dataset
+  // Register new partitions with the Dataset.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc RegisterWithDataset(RegisterWithDatasetRequest) returns (RegisterWithDatasetResponse) {}
 
   // Write chunks to one or more partitions.
@@ -71,22 +73,30 @@ service RerunCloudService {
   // Inspect the contents of the partition table.
   //
   // The data will follow the schema returned by `GetPartitionTableSchema`.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc ScanPartitionTable(ScanPartitionTableRequest) returns (stream ScanPartitionTableResponse) {}
 
   // Returns the schema of the dataset.
   //
   // This is the union of all the schemas from all the underlying partitions. It will contain all the indexes,
   // entities and components present in the dataset.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc GetDatasetSchema(GetDatasetSchemaRequest) returns (GetDatasetSchemaResponse) {}
 
   /* Indexing */
 
   // Creates a custom index for a specific column (vector search, full-text search, etc).
+  //
+  // This endpoint requires the standard dataset headers.
   rpc CreateIndex(CreateIndexRequest) returns (CreateIndexResponse) {}
 
   /* Queries */
 
   // Search a previously created index.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc SearchDataset(SearchDatasetRequest) returns (stream SearchDatasetResponse) {}
 
   // Perform Rerun-native queries on a dataset, returning the matching chunk IDs, as well
@@ -103,6 +113,8 @@ service RerunCloudService {
   // To fetch the actual chunks themselves, see `GetChunks`.
   //
   // Passing chunk IDs to this method effectively acts as a IF_EXIST filter.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc QueryDataset(QueryDatasetRequest) returns (stream QueryDatasetResponse) {}
 
   // Perform Rerun-native queries on a dataset, returning the underlying chunks.
@@ -113,6 +125,8 @@ service RerunCloudService {
   // * Arbitrary Lance filters.
   //
   // To fetch only the actual chunk IDs rather than the chunks themselves, see `QueryDataset`.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc GetChunks(GetChunksRequest) returns (stream GetChunksResponse) {}
 
   // Fetch specific chunks from Rerun Cloud. In a 2-step query process, result of 1st phase,
@@ -146,6 +160,8 @@ service RerunCloudService {
   // --- Utilities ---
 
   // Rerun Manifests maintenance operations: scalar index creation, compaction, etc.
+  //
+  // This endpoint requires the standard dataset headers.
   rpc DoMaintenance(DoMaintenanceRequest) returns (DoMaintenanceResponse) {}
 
   // Run global maintenance operations on the platform: this includes optimization
@@ -515,8 +531,6 @@ message QueryRange {
 }
 
 message GetChunksRequest {
-  rerun.common.v1alpha1.EntryId dataset_id = 1;
-
   // Client can specify from which partitions to get chunks. If left unspecified (empty list),
   // data from all partition (that match other query parameters) will be included.
   repeated rerun.common.v1alpha1.PartitionId partition_ids = 2;
@@ -562,6 +576,9 @@ message GetChunksRequest {
 
   // Query details
   Query query = 5;
+
+  reserved 1;
+  reserved "dataset_id";
 }
 
 message GetChunksResponse {
@@ -609,8 +626,6 @@ message ScanTableResponse {
 // --- Maintenance ---
 
 message DoMaintenanceRequest {
-  rerun.common.v1alpha1.EntryId dataset_id = 1;
-
   // Optimize all builtin and user-defined indexes on this dataset.
   //
   // This merges all individual index deltas back in the main index, improving runtime performance
@@ -645,6 +660,9 @@ message DoMaintenanceRequest {
   //
   // ⚠️ Do not ever use this unless you know exactly what you're doing. Improper use will lead to data loss.
   bool unsafe_allow_recent_cleanup = 5;
+
+  reserved 1;
+  reserved "dataset_id";
 }
 
 message DoMaintenanceResponse {

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

@@ -63,7 +63,6 @@ impl From<RegisterWithDatasetRequest> for crate::cloud::v1alpha1::RegisterWithDa
 
 #[derive(Debug, Clone)]
 pub struct GetChunksRequest {
-    pub dataset_id: EntryId,
     pub partition_ids: Vec<crate::common::v1alpha1::ext::PartitionId>,
     pub chunk_ids: Vec<re_chunk::ChunkId>,
     pub entity_paths: Vec<EntityPath>,
@@ -75,11 +74,6 @@ impl TryFrom<crate::cloud::v1alpha1::GetChunksRequest> for GetChunksRequest {
 
     fn try_from(value: crate::cloud::v1alpha1::GetChunksRequest) -> Result<Self, Self::Error> {
         Ok(Self {
-            dataset_id: value
-                .dataset_id
-                .ok_or_else(|| tonic::Status::invalid_argument("dataset_id is required"))?
-                .try_into()?,
-
             partition_ids: value
                 .partition_ids
                 .into_iter()
@@ -112,7 +106,6 @@ impl TryFrom<crate::cloud::v1alpha1::GetChunksRequest> for GetChunksRequest {
 
 #[derive(Debug, Clone)]
 pub struct DoMaintenanceRequest {
-    pub dataset_id: Option<crate::common::v1alpha1::EntryId>,
     pub optimize_indexes: bool,
     pub retrain_indexes: bool,
     pub compact_fragments: bool,
@@ -123,7 +116,6 @@ pub struct DoMaintenanceRequest {
 impl From<DoMaintenanceRequest> for crate::cloud::v1alpha1::DoMaintenanceRequest {
     fn from(value: DoMaintenanceRequest) -> Self {
         Self {
-            dataset_id: value.dataset_id,
             optimize_indexes: value.optimize_indexes,
             retrain_indexes: value.retrain_indexes,
             compact_fragments: value.compact_fragments,

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

@@ -370,6 +370,23 @@ where
         Ok(response.table_entry)
     }
 
+    pub async fn get_chunks(
+        &mut self,
+        dataset_id: EntryId,
+        req: re_protos::cloud::v1alpha1::GetChunksRequest,
+    ) -> Result<tonic::Streaming<re_protos::cloud::v1alpha1::GetChunksResponse>, StreamError> {
+        Ok(self
+            .inner()
+            .get_chunks(
+                tonic::Request::new(req)
+                    .with_entry_id(dataset_id)
+                    .map_err(|err| StreamEntryError::InvalidId(err.into()))?,
+            )
+            .await
+            .map_err(|err| crate::StreamPartitionError::StreamingChunks(err.into()))?
+            .into_inner())
+    }
+
     #[allow(clippy::fn_params_excessive_bools)]
     pub async fn do_maintenance(
         &mut self,
@@ -381,17 +398,20 @@ where
         unsafe_allow_recent_cleanup: bool,
     ) -> Result<(), StreamError> {
         self.inner()
-            .do_maintenance(tonic::Request::new(
-                re_protos::cloud::v1alpha1::ext::DoMaintenanceRequest {
-                    dataset_id: Some(dataset_id.into()),
-                    optimize_indexes,
-                    retrain_indexes,
-                    compact_fragments,
-                    cleanup_before,
-                    unsafe_allow_recent_cleanup,
-                }
-                .into(),
-            ))
+            .do_maintenance(
+                tonic::Request::new(
+                    re_protos::cloud::v1alpha1::ext::DoMaintenanceRequest {
+                        optimize_indexes,
+                        retrain_indexes,
+                        compact_fragments,
+                        cleanup_before,
+                        unsafe_allow_recent_cleanup,
+                    }
+                    .into(),
+                )
+                .with_entry_id(dataset_id)
+                .map_err(|err| StreamEntryError::InvalidId(err.into()))?,
+            )
             .await
             .map_err(|err| StreamEntryError::Maintenance(err.into()))?;