Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 25 additions & 8 deletions tools/pubsys/src/aws/ami/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -208,13 +208,19 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result<HashMap<String, Image>>
"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
Expand Down Expand Up @@ -382,13 +388,24 @@ async fn _run(args: &Args, ami_args: &AmiArgs) -> Result<HashMap<String, Image>>
}

let ec2_client = &ec2_clients[&region];

let base_region = base_region.to_owned();
let copy_future = ec2_client
.copy_image()
.set_description(tentative_amispec.description.clone())
.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
Expand Down
3 changes: 3 additions & 0 deletions tools/pubsys/src/aws/ami/register/mk_amispec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ pub(crate) struct MinimalAmiSpec {
pub name: String,
pub description: Option<String>,
pub architecture: Option<amispec::Architecture>,
pub tags: Option<HashMap<String, String>>,
}

impl From<AmiSpec> for MinimalAmiSpec {
Expand All @@ -136,13 +137,15 @@ impl From<AmiSpec> for MinimalAmiSpec {
name,
description,
architecture,
tags,
..
} = amispec;

Self {
name,
description,
architecture,
tags,
}
}
}
Expand Down
16 changes: 15 additions & 1 deletion tools/pubsys/src/aws/ami/register/mod.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
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;
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;
Expand All @@ -29,6 +30,7 @@ async fn _register_image(
region: &Region,
ebs_client: EbsClient,
ec2_client: &Ec2Client,
tags: Option<HashMap<String, String>>,
cleanup_snapshot_ids: Arc<Mutex<Vec<String>>>,
) -> Result<RegisteredIds> {
let bottlerocket_snapshots = BottlerocketSnapshots::create_snapshots(
Expand All @@ -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 {
Expand Down Expand Up @@ -199,13 +211,15 @@ pub(crate) async fn register_image(
region: &Region,
ebs_client: EbsClient,
ec2_client: &Ec2Client,
tags: Option<HashMap<String, String>>,
) -> Result<RegisteredIds> {
let cleanup_snapshot_ids = Arc::new(Mutex::new(Vec::new()));
let register_result = _register_image(
ami_args,
region,
ebs_client,
ec2_client,
tags,
Arc::clone(&cleanup_snapshot_ids),
)
.await;
Expand Down
Loading