zakariya-s commented on code in PR #2976:
URL: https://github.com/apache/iceberg-rust/pull/2976#discussion_r4205003272
##########
crates/storage/opendal/src/lib.rs:
##########
@@ -1098,11 +1354,353 @@ mod tests {
}
}
+ /// Vends a credential for the whole `s3` scheme until it has served
+ /// `table-c`, then credentials scoped to each table. The scheme-wide
+ /// credential expires once the `table-c` batch is signed.
+ #[cfg(feature = "opendal-s3")]
+ #[derive(Debug, Default)]
+ struct ChangingScopeProvider {
+ scoped: AtomicBool,
+ scheme_wide_expired: AtomicBool,
+ }
+
+ #[cfg(feature = "opendal-s3")]
+ #[async_trait]
+ impl StorageCredentialProvider for ChangingScopeProvider {
+ fn supports_path(&self, _path: &str) -> bool {
+ true
+ }
+
+ async fn load_credential(&self, path: &str) ->
Result<StorageCredential> {
+ if path.starts_with("s3://bucket/table-c") {
+ self.scoped.store(true, Ordering::SeqCst);
+ }
+ // Signing the scoped batch means the scheme-wide credential is
gone.
+ if path == "s3://bucket/table-c" {
+ self.scheme_wide_expired.store(true, Ordering::SeqCst);
+ }
+ match path.strip_prefix("s3://bucket/") {
+ Some(rest) if self.scoped.load(Ordering::SeqCst) &&
!rest.is_empty() => {
+ let table = rest.split('/').next().unwrap_or_default();
+ Ok(s3_credential(format!("s3://bucket/{table}"),
"SCOPED_AK"))
+ }
+ _ if self.scheme_wide_expired.load(Ordering::SeqCst) =>
Err(Error::new(
+ ErrorKind::Unexpected,
+ "the scheme-wide credential expired",
+ )),
+ _ => Ok(s3_credential("s3", "SCHEME_AK")),
+ }
+ }
+ }
+
+ #[cfg(feature = "opendal-s3")]
+ #[tokio::test]
+ async fn test_delete_stream_signs_each_batch_with_its_vended_credential() {
+ let mut server = mockito::Server::new_async().await;
+ let mut delete = |path: &str, access_key: &str| {
+ server
+ .mock("DELETE", path)
+ .match_header(
+ "authorization",
+
mockito::Matcher::Regex(format!("Credential={access_key}/")),
+ )
+ .expect(1)
+ .with_status(204)
+ };
+ // The scheme-wide batch is flushed, with its own credential, before
+ // the scoped credential that replaces it is used.
+ let scheme_wide = delete("/bucket/table-b/f.parquet", "SCHEME_AK")
+ .create_async()
+ .await;
+ let scoped = delete("/bucket/table-c/f.parquet", "SCOPED_AK")
+ .create_async()
+ .await;
+
+ let mut config = S3Config::default();
+ config.endpoint = Some(server.url());
+ config.region = Some("us-east-1".to_string());
+ config.disable_config_load = true;
+ config.disable_ec2_metadata = true;
+ let storage = OpenDalStorage::S3 {
+ config: Arc::new(config),
+ customized_credential_load: None,
+ credential_provider:
Some(Arc::new(ChangingScopeProvider::default())),
+ client_config: OpenDalClientConfig::default(),
+ };
+ storage
+ .delete_stream(
+ futures::stream::iter([
+ "s3://bucket/table-b/f.parquet".to_string(),
+ "s3://bucket/table-c/f.parquet".to_string(),
+ ])
+ .boxed(),
+ )
+ .await
+ .unwrap();
+
+ scheme_wide.assert_async().await;
+ scoped.assert_async().await;
+ }
+
+ /// Vends `token` as a GCS credential, or fails when it is `None`.
+ #[cfg(feature = "opendal-gcs")]
+ #[derive(Debug)]
+ struct GcsTokenProvider(Option<&'static str>);
+
+ #[cfg(feature = "opendal-gcs")]
+ #[async_trait]
+ impl StorageCredentialProvider for GcsTokenProvider {
+ fn supports_path(&self, _path: &str) -> bool {
+ true
+ }
+
+ async fn load_credential(&self, _path: &str) ->
Result<StorageCredential> {
+ let token = self
+ .0
+ .ok_or_else(|| Error::new(ErrorKind::Unexpected, "refresh
failed"))?;
+ Ok(StorageCredential::new(
+ "gs",
+ HashMap::from([(GCS_TOKEN.to_string(), token.to_string())]),
+ ))
+ }
+ }
+
+ #[cfg(feature = "opendal-gcs")]
+ fn gcs_storage(server: &mockito::Server, provider: GcsTokenProvider) ->
OpenDalStorage {
+ let config = gcs_config_parse(HashMap::from([
+ (GCS_SERVICE_HOST.to_string(), server.url()),
+ (GCS_TOKEN.to_string(), "static-token".to_string()),
+ ]))
+ .unwrap();
+ OpenDalStorage::Gcs {
+ config: Arc::new(config),
+ credential_provider: Some(Arc::new(provider)),
+ client_config: OpenDalClientConfig::default(),
+ }
+ }
+
+ #[cfg(feature = "opendal-gcs")]
+ #[tokio::test]
+ async fn test_gcs_signs_with_the_vended_token() {
+ let mut server = mockito::Server::new_async().await;
+ let mock = server
+ .mock("GET", mockito::Matcher::Any)
+ .match_header("authorization", "Bearer vended-token")
+ .expect_at_least(1)
+ .with_status(404)
+ .create_async()
+ .await;
+
+ let storage = gcs_storage(&server,
GcsTokenProvider(Some("vended-token")));
+ assert!(!storage.exists("gs://bucket/file.parquet").await.unwrap());
+ mock.assert_async().await;
+ }
Review Comment:
yes good catch, made both an s3
[test_s3_operator_reloads_an_expiring_credential](https://github.com/apache/iceberg-rust/pull/2976/changes/7ff3d887594502ca2dd04d4b7b95ce3b8c46c6f6#diff-3c899976a8c7dc30d93ea1fea29403afdf76b1efcf175770cdfef47de4a22990R1511-R1548)
and a gcs one
[test_gcs_operator_reloads_an_expiring_token](https://github.com/apache/iceberg-rust/pull/2976/changes/7ff3d887594502ca2dd04d4b7b95ce3b8c46c6f6#diff-3c899976a8c7dc30d93ea1fea29403afdf76b1efcf175770cdfef47de4a22990R1639-R1662)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
