diff --git a/tools/pubsys/src/aws/ami/mod.rs b/tools/pubsys/src/aws/ami/mod.rs index 81661499..93746b8d 100644 --- a/tools/pubsys/src/aws/ami/mod.rs +++ b/tools/pubsys/src/aws/ami/mod.rs @@ -17,7 +17,7 @@ use crate::Args; use aws_sdk_ebs::Client as EbsClient; use aws_sdk_ec2::error::{ProvideErrorMetadata, SdkError}; use aws_sdk_ec2::operation::copy_image::{CopyImageError, CopyImageOutput}; -use aws_sdk_ec2::types::OperationType; +use aws_sdk_ec2::types::{OperationType, ResourceType, Tag, TagSpecification}; use aws_sdk_ec2::{config::Region, Client as Ec2Client}; use aws_sdk_sts::operation::get_caller_identity::{ GetCallerIdentityError, GetCallerIdentityOutput, @@ -208,13 +208,19 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> "Registering '{}' in {}", tentative_amispec.name, base_region ); - let new_ids = register_image(ami_args, &base_region, base_ebs_client, &base_ec2_client) - .await - .context(error::RegisterImageSnafu { - name: &tentative_amispec.name, - arch: &arch, - region: base_region.as_ref(), - })?; + let new_ids = register_image( + ami_args, + &base_region, + base_ebs_client, + &base_ec2_client, + tentative_amispec.tags.clone(), + ) + .await + .context(error::RegisterImageSnafu { + name: &tentative_amispec.name, + arch: &arch, + region: base_region.as_ref(), + })?; info!( "Registered AMI '{}' in {}: {}", tentative_amispec.name, base_region, new_ids.image_id @@ -382,6 +388,7 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> } let ec2_client = &ec2_clients[®ion]; + let base_region = base_region.to_owned(); let copy_future = ec2_client .copy_image() @@ -389,6 +396,16 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result> .set_name(Some(tentative_amispec.name.clone())) .set_source_image_id(Some(ids_of_image.image_id.clone())) .set_source_region(Some(base_region.as_ref().to_string())) + .set_tag_specifications(tentative_amispec.tags.as_ref().map(|x| { + vec![TagSpecification::builder() + .resource_type(ResourceType::Image) + .set_tags(Some( + x.iter() + .map(|(key, value)| Tag::builder().key(key).value(value).build()) + .collect(), + )) + .build()] + })) .send(); // Store the region so we can output it to the user diff --git a/tools/pubsys/src/aws/ami/register/mk_amispec.rs b/tools/pubsys/src/aws/ami/register/mk_amispec.rs index afc870ea..3c5ed28d 100644 --- a/tools/pubsys/src/aws/ami/register/mk_amispec.rs +++ b/tools/pubsys/src/aws/ami/register/mk_amispec.rs @@ -128,6 +128,7 @@ pub(crate) struct MinimalAmiSpec { pub name: String, pub description: Option, pub architecture: Option, + pub tags: Option>, } impl From for MinimalAmiSpec { @@ -136,6 +137,7 @@ impl From for MinimalAmiSpec { name, description, architecture, + tags, .. } = amispec; @@ -143,6 +145,7 @@ impl From for MinimalAmiSpec { name, description, architecture, + tags, } } } diff --git a/tools/pubsys/src/aws/ami/register/mod.rs b/tools/pubsys/src/aws/ami/register/mod.rs index d8ec7fba..0ddd38e2 100644 --- a/tools/pubsys/src/aws/ami/register/mod.rs +++ b/tools/pubsys/src/aws/ami/register/mod.rs @@ -1,7 +1,7 @@ use super::{snapshot::snapshot_from_image, AmiArgs}; use crate::aws::ami::snapshot::build_progress_bar; use aws_sdk_ebs::Client as EbsClient; -use aws_sdk_ec2::types::Filter; +use aws_sdk_ec2::types::{Filter, ResourceType, Tag, TagSpecification}; use aws_sdk_ec2::{config::Region, Client as Ec2Client}; use coldsnap::{SnapshotUploader, SnapshotWaiter}; use futures::future::OptionFuture; @@ -9,6 +9,7 @@ use futures::TryFutureExt as _; use indicatif::{MultiProgress, ProgressBar}; use log::{debug, info, warn}; use snafu::{ensure, futures::TryFutureExt as _, OptionExt, ResultExt, Snafu}; +use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; use tokio::sync::Mutex; @@ -29,6 +30,7 @@ async fn _register_image( region: &Region, ebs_client: EbsClient, ec2_client: &Ec2Client, + tags: Option>, cleanup_snapshot_ids: Arc>>, ) -> Result { let bottlerocket_snapshots = BottlerocketSnapshots::create_snapshots( @@ -48,6 +50,16 @@ async fn _register_image( info!("Making register image call in {region}"); let register_response = amispec .as_register_image_call() + .set_tag_specifications(tags.map(|x| { + vec![TagSpecification::builder() + .resource_type(ResourceType::Image) + .set_tags(Some( + x.iter() + .map(|(key, value)| Tag::builder().key(key).value(value).build()) + .collect(), + )) + .build()] + })) .send_with(ec2_client) .await .context(error::RegisterImageSnafu { @@ -199,6 +211,7 @@ pub(crate) async fn register_image( region: &Region, ebs_client: EbsClient, ec2_client: &Ec2Client, + tags: Option>, ) -> Result { let cleanup_snapshot_ids = Arc::new(Mutex::new(Vec::new())); let register_result = _register_image( @@ -206,6 +219,7 @@ pub(crate) async fn register_image( region, ebs_client, ec2_client, + tags, Arc::clone(&cleanup_snapshot_ids), ) .await;