S3 sidecars

crates/pile-dataset/src/source/s3.rs

@@ -56,11 +56,61 @@ impl S3DataSource {
 		})
 	}
 
-	fn make_item(self: &Arc<Self>, key: impl Into<SmartString<LazyCompact>>) -> Item {
+	async fn find_sidecar_key(&self, key: &str) -> Option<SmartString<LazyCompact>> {
+		// First try {key}.toml
+		let full_toml = format!("{key}.toml");
+		if self
+			.client
+			.head_object()
+			.bucket(self.bucket.as_str())
+			.key(&full_toml)
+			.send()
+			.await
+			.is_ok()
+		{
+			return Some(full_toml.into());
+		}
+
+		// Then try {key-with-extension-stripped}.toml
+		let stripped = std::path::Path::new(key).with_extension("toml");
+		if let Some(stripped_str) = stripped.to_str()
+			&& stripped_str != full_toml.as_str()
+			&& self
+				.client
+				.head_object()
+				.bucket(self.bucket.as_str())
+				.key(stripped_str)
+				.send()
+				.await
+				.is_ok()
+		{
+			return Some(stripped_str.into());
+		}
+
+		None
+	}
+
+	async fn make_item(self: &Arc<Self>, key: impl Into<SmartString<LazyCompact>>) -> Item {
+		let key: SmartString<LazyCompact> = key.into();
+
+		let sidecar = if self.sidecars {
+			self.find_sidecar_key(key.as_str())
+				.await
+				.map(|sidecar_key| {
+					Box::new(Item::S3 {
+						source: Arc::clone(self),
+						key: sidecar_key,
+						sidecar: None,
+					})
+				})
+		} else {
+			None
+		};
+
 		Item::S3 {
 			source: Arc::clone(self),
-			key: key.into(),
-			sidecar: None, // TODO: add sidecars
+			key,
+			sidecar,
 		}
 	}
 }
@@ -90,7 +140,7 @@ impl DataSource for Arc<S3DataSource> {
 				}
 				Err(std::io::Error::other(sdk_err))
 			}
-			Ok(_) => Ok(Some(self.make_item(key))),
+			Ok(_) => Ok(Some(self.make_item(key).await)),
 		}
 	}
 
@@ -136,11 +186,7 @@ impl DataSource for Arc<S3DataSource> {
 						continue;
 					}
 
-					let item = Item::S3 {
-						source: Arc::clone(&source),
-						key: key.into(),
-						sidecar: None, // TODO: add sidecars
-					};
+					let item = source.make_item(key).await;
 
 					if tx.send(Ok(item)).await.is_err() {
 						return;