fix: retry if the url is invalid

Author: cimandef

client/src/s3.rs

@@ -188,7 +188,7 @@ impl Client {
     }
 
     pub async fn s3_head(&self, url: Url) -> Result<S3HeadResponse> {
-        self.validate_url_host(&url)?;
+        self.validate_host(&url)?;
 
         let body = Body::from(r#"{"head": true}"#.to_string());
 
@@ -197,7 +197,9 @@ impl Client {
             .uri(url.to_string());
 
         if let Some(content_length) = body.content_length() {
-            if let Some(headers) = request.headers_mut() { headers.insert(reqwest::header::CONTENT_LENGTH, content_length.into()); }
+            if let Some(headers) = request.headers_mut() {
+                headers.insert(reqwest::header::CONTENT_LENGTH, content_length.into());
+            }
         }
 
         let request = request.body(body)?;

src/commands/dataset/download.rs

@@ -43,6 +43,14 @@ pub struct Download {
 )]
 pub struct GetDatasetVersionFiles;
 
+#[derive(GraphQLQuery)]
+#[graphql(
+    query_path = "src/graphql/get_dataset_version_file_by_partition.graphql",
+    schema_path = "schema.graphql",
+    response_derives = "Debug"
+)]
+pub struct GetDatasetVersionFileByPartition;
+
 pub async fn download(args: Download, global: GlobalArgs) -> Result<()> {
     let m = MultiProgress::new();
 
@@ -103,9 +111,9 @@ pub async fn download(args: Download, global: GlobalArgs) -> Result<()> {
 
     stream::iter(nodes)
         .map(|node| {
-            let client = &client;
-            let m = &m;
-            let multipart_options = &multipart_options;
+            let client = client.to_owned();
+            let m = m.to_owned();
+            let multipart_options = multipart_options.to_owned();
             let dataset_dir = dataset_dir.to_owned();
             let dataset_name = dataset_name.to_owned();
 
@@ -138,7 +146,44 @@ async fn download_partition_file(
     dataset_name: &str,
     file_node: GetDatasetVersionFilesNodeOnDatasetVersionFilesNodes,
 ) -> Result<()> {
-    let metadata = client.s3_head(file_node.url.clone()).await?;
+    let (metadata, url) = match client.s3_head(file_node.url.clone()).await {
+        Ok(metadata) => (metadata, file_node.url.clone()),
+        // retry if presigned url expired due to long dataset download time
+        Err(e) => {
+            tracing::warn!(error = %e, "Retrying: failed to fetch object header");
+            let response = client
+                .send::<GetDatasetVersionFileByPartition>(
+                    get_dataset_version_file_by_partition::Variables {
+                        dataset_version_id: file_node.dataset_version.id,
+                        partition_num: file_node.partition_num,
+                    },
+                )
+                .await?;
+
+            let dataset_version_file = match response.node {
+            get_dataset_version_file_by_partition::GetDatasetVersionFileByPartitionNode::DatasetVersion(v) => v,
+            _ => {
+                return Err(error::system(
+                    "Invalid node type",
+                    "Unexpected GraphQL response",
+                ));
+            }
+        };
+            let file_by_partition_num = match dataset_version_file.file_by_partition_num {
+                Some(file_by_partition_num) => file_by_partition_num,
+                None => {
+                    return Err(error::system(
+                        "Invalid partition number",
+                        "The partition does not exist",
+                    ))
+                }
+            };
+
+            let file_url = file_by_partition_num.url;
+            (client.s3_head(file_url.clone()).await?, file_url)
+        }
+    };
+
     let filename = format!("{}-{}.parquet", dataset_name, file_node.partition_num);
     let output_path = output_dir.join(&filename);
 
@@ -159,7 +204,7 @@ async fn download_partition_file(
     multipart_download(
         client,
         metadata.size,
-        file_node.url.clone(),
+        url,
         multipart_options,
         &temp_path,
         &pb,

src/graphql/get_dataset_version_file_by_partition.graphql

@@ -0,0 +1,14 @@
+query GetDatasetVersionFileByPartition(
+  $datasetVersionId: ID!
+  $partitionNum: Int!
+) {
+  node(id: $datasetVersionId) {
+    __typename
+    ... on DatasetVersion {
+      fileByPartitionNum(partitionNum: $partitionNum) {
+        id
+        url
+      }
+    }
+  }
+}

src/graphql/get_dataset_version_files.graphql

@@ -12,6 +12,9 @@ query GetDatasetVersionFiles($datasetVersionId: ID!) {
           rows
           url
           partitionNum
+          datasetVersion {
+            id
+          }
         }
       }
     }