feat(commands): add ability to download attachment content for comments

reinfer · Jul 23, 2024 · ed345d4 · ed345d4
1 parent 2f92e5d
commit ed345d4
Show file tree

Hide file tree

Showing 5 changed files with 263 additions and 28 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,6 @@
 # Unreleased
 - Add `config parse-from-url` command for parsing configuration from a URL
+- Add ability to download attachments for comments
 
 # v0.28.0
 - Add general fields to `create datasets`

diff --git a/api/src/lib.rs b/api/src/lib.rs
@@ -14,7 +14,7 @@ use reqwest::{
 };
 use resources::{
     bucket_statistics::GetBucketStatisticsResponse,
-    comment::CommentTimestampFilter,
+    comment::{AttachmentReference, CommentTimestampFilter},
     dataset::{
         QueryRequestParams, QueryResponse,
         StatisticsRequestParams as DatasetStatisticsRequestParams, SummaryRequestParams,
@@ -36,7 +36,7 @@ use resources::{
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;
-use std::{cell::Cell, fmt::Display, path::Path};
+use std::{cell::Cell, fmt::Display, io::Read, path::Path};
 use url::Url;
 
 use crate::resources::{
@@ -622,6 +622,27 @@ impl Client {
         Ok(())
     }
 
+    pub fn get_attachment(&self, reference: &AttachmentReference) -> Result<Vec<u8>> {
+        let mut response = self.raw_request(
+            &Method::GET,
+            &self.endpoints.attachment(reference)?,
+            &None::<()>,
+            &None::<()>,
+            &Retry::Yes,
+        )?;
+
+        let mut buffer = Vec::new();
+
+        response
+            .read_to_end(&mut buffer)
+            .map_err(|source| Error::Unknown {
+                message: "Failed to read buffer".to_string(),
+                source: Box::new(source),
+            })?;
+
+        Ok(buffer)
+    }
+
     pub fn get_integrations(&self) -> Result<Vec<Integration>> {
         Ok(self
             .get::<_, GetIntegrationsResponse>(self.endpoints.integrations()?)?
@@ -1152,21 +1173,19 @@ impl Client {
         self.request(Method::PUT, url, Some(request), None::<()>, Retry::Yes)
     }
 
-    fn request<LocationT, RequestT, SuccessT, QueryT>(
+    fn raw_request<LocationT, RequestT, QueryT>(
         &self,
-        method: Method,
-        url: LocationT,
-        body: Option<RequestT>,
-        query: Option<QueryT>,
-        retry: Retry,
-    ) -> Result<SuccessT>
+        method: &Method,
+        url: &LocationT,
+        body: &Option<RequestT>,
+        query: &Option<QueryT>,
+        retry: &Retry,
+    ) -> Result<reqwest::blocking::Response>
     where
         LocationT: IntoUrl + Display + Clone,
         RequestT: Serialize,
         QueryT: Serialize,
-        for<'de> SuccessT: Deserialize<'de>,
     {
-        debug!("Attempting {} `{}`", method, url);
         let do_request = || {
             let request = self
                 .http_client
@@ -1192,6 +1211,26 @@ impl Client {
             message: format!("{method} operation failed."),
         })?;
 
+        Ok(http_response)
+    }
+
+    fn request<LocationT, RequestT, SuccessT, QueryT>(
+        &self,
+        method: Method,
+        url: LocationT,
+        body: Option<RequestT>,
+        query: Option<QueryT>,
+        retry: Retry,
+    ) -> Result<SuccessT>
+    where
+        LocationT: IntoUrl + Display + Clone,
+        RequestT: Serialize,
+        QueryT: Serialize,
+        for<'de> SuccessT: Deserialize<'de>,
+    {
+        debug!("Attempting {} `{}`", method, url);
+        let http_response = self.raw_request(&method, &url, &body, &query, &retry)?;
+
         let status = http_response.status();
 
         http_response
@@ -1489,6 +1528,10 @@ impl Endpoints {
         construct_endpoint(&self.base, &["api", "_private", "integrations", &name.0])
     }
 
+    fn attachment(&self, reference: &AttachmentReference) -> Result<Url> {
+        construct_endpoint(&self.base, &["api", "v1", "attachments", &reference.0])
+    }
+
     fn validation(
         &self,
         dataset_name: &DatasetFullName,

diff --git a/api/src/resources/comment.rs b/api/src/resources/comment.rs
@@ -292,11 +292,21 @@ pub enum Sentiment {
     Negative,
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+pub struct AttachmentReference(pub String);
+
+#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
+pub struct ContentHash(pub String);
+
 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
 pub struct AttachmentMetadata {
     pub name: String,
     pub size: u64,
     pub content_type: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub attachment_reference: Option<AttachmentReference>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub content_hash: Option<ContentHash>,
 }
 
 #[derive(Debug, Clone, PartialEq, Default, Eq)]