Skip to content

Commit

Permalink
bib: show progress on ami upload
Browse files Browse the repository at this point in the history
Add a new `--progress` option that defaults to `text` and show
upload progress when uploading an AMI.
  • Loading branch information
mvo5 committed Jan 29, 2024
1 parent 9e7a536 commit eb51d1d
Show file tree
Hide file tree
Showing 9 changed files with 198 additions and 14 deletions.
16 changes: 15 additions & 1 deletion bib/cmd/bootc-image-builder/cloud.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package main

import (
"github.com/cheggaaa/pb"
"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
"github.com/osbuild/images/pkg/cloud/awscloud"
"github.com/spf13/pflag"
Expand All @@ -19,10 +20,23 @@ func uploadAMI(path string, flags *pflag.FlagSet) error {
if err != nil {
return err
}
progress, err := flags.GetString("progress")
if err != nil {
return err
}

client, err := awscloud.NewDefault(region)
if err != nil {
return err
}
return uploader.UploadAndRegister(client, path, bucketName, imageName)

// TODO: extract this as a helper once we add "uploadAzure" or
// similar.
var pbar *pb.ProgressBar
switch progress {
case "text":
pbar = pb.New(0)
}

return uploader.UploadAndRegister(client, path, bucketName, imageName, pbar)
}
1 change: 1 addition & 0 deletions bib/cmd/bootc-image-builder/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,7 @@ func run() error {
buildCmd.Flags().String("aws-region", "", "target region for AWS uploads (only for type=ami)")
buildCmd.Flags().String("aws-bucket", "", "target S3 bucket name for intermediate storage when creating AMI (only for type=ami)")
buildCmd.Flags().String("aws-ami-name", "", "name for the AMI in AWS (only for type=ami)")
buildCmd.Flags().String("progress", "text", "type of progress bar to use")

// flag rules
for _, dname := range []string{"output", "store", "rpmmd"} {
Expand Down
2 changes: 1 addition & 1 deletion bib/cmd/upload/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ func uploadAMI(cmd *cobra.Command, args []string) {
imageName, err := flags.GetString("ami-name")
check(err)

check(uploader.UploadAndRegister(client, filename, bucketName, imageName))
check(uploader.UploadAndRegister(client, filename, bucketName, imageName, nil))
}

func setupCLI() *cobra.Command {
Expand Down
1 change: 1 addition & 0 deletions bib/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.19

require (
github.com/aws/aws-sdk-go v1.50.5
github.com/cheggaaa/pb v1.0.29
github.com/google/uuid v1.6.0
github.com/osbuild/images v0.34.0
github.com/sirupsen/logrus v1.9.3
Expand Down
12 changes: 12 additions & 0 deletions bib/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ github.com/aws/aws-sdk-go v1.50.5/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3Tju
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
github.com/cheggaaa/pb v1.0.29 h1:FckUN5ngEk2LpvuG0fw1GEFx6LtyY2pWI/Z2QgCnEYo=
github.com/cheggaaa/pb v1.0.29/go.mod h1:W40334L7FMC5JKWldsTWbdGjLo0RxUKK73K+TuPxX30=
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
github.com/containerd/cgroups/v3 v3.0.2 h1:f5WFqIVSgo5IZmtTT3qVBo6TzI1ON6sycSBKkymb9L0=
Expand Down Expand Up @@ -72,6 +74,8 @@ github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7
github.com/facebookgo/clock v0.0.0-20150410010913-600d898af40a h1:yDWHCSQ40h88yih2JAcL6Ls/kVkSE8GFACTGVnMPruw=
github.com/facebookgo/limitgroup v0.0.0-20150612190941-6abd8d71ec01 h1:IeaD1VDVBPlx3viJT9Md8if8IxxJnO+x0JCGb054heg=
github.com/facebookgo/muster v0.0.0-20150708232844-fd3d7953fd52 h1:a4DFiKFJiDRGFD1qIcqGLX/WlUMD9dyLSLDt+9QZgt8=
github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
github.com/go-jose/go-jose/v3 v3.0.1 h1:pWmKFVtt+Jl0vBZTIpz/eAKwsm6LkIxDVVbFHKkchhA=
github.com/go-jose/go-jose/v3 v3.0.1/go.mod h1:RNkWWRld676jZEYoV3+XK8L2ZnNSvIsxFMht0mSX+u8=
github.com/go-openapi/analysis v0.21.2/go.mod h1:HZwRk4RRisyG8vx2Oe6aqeSQcoxRp47Xkp3+K6q+LdY=
Expand Down Expand Up @@ -230,6 +234,12 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
github.com/markbates/oncer v0.0.0-20181203154359-bf2de49a0be2/go.mod h1:Ld9puTsIW75CHf65OeIOkyKbteujpZVXDpWK6YGZbxE=
github.com/markbates/safe v1.0.1/go.mod h1:nAqgmRi7cY2nqMc92/bSEeQA+R4OheNU2T1kNSCBdG0=
github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
github.com/mattn/go-isatty v0.0.11/go.mod h1:PhnuNfih5lzO57/f3n+odYbM4JtupLOxQOAqxQCu2WE=
github.com/mattn/go-isatty v0.0.17 h1:BTarxUcIeDqL27Mc+vyvdWYSL28zpIhv3RoTdsLMPng=
github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U=
github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/mattn/go-shellwords v1.0.12 h1:M2zGm7EW6UQJvDeQxo4T51eKPurbeFbe8WtebGE2xrk=
Expand Down Expand Up @@ -416,11 +426,13 @@ golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
golang.org/x/sys v0.0.0-20190403152447-81d4e9dc473e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190419153524-e8e3143a4f4a/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20190531175056-4c3a928424d2/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
Expand Down
55 changes: 46 additions & 9 deletions bib/internal/uploader/aws.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,31 +2,68 @@ package uploader

import (
"fmt"
"io"
"os"
"path/filepath"

"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/aws/aws-sdk-go/service/s3/s3manager"
"github.com/cheggaaa/pb"
"github.com/google/uuid"

"github.com/osbuild/images/pkg/arch"
"github.com/osbuild/images/pkg/cloud/awscloud"
)

func UploadAndRegister(a *awscloud.AWS, filename, bucketName, imageName string) error {
keyName := fmt.Sprintf("%s-%s", uuid.New().String(), filepath.Base(filename))
var osStdout io.Writer = os.Stdout

type AwsUploader interface {
UploadFromReader(r io.Reader, bucketName, keyName string) (*s3manager.UploadOutput, error)
Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode *string) (*string, *string, error)
}

func doUpload(a AwsUploader, file *os.File, bucketName, keyName string, pbar *pb.ProgressBar) (*s3manager.UploadOutput, error) {
var r io.Reader = file

// TODO: extract this as a helper once we add "uploadAzure" or
// similar.
if pbar != nil {
st, err := file.Stat()
if err != nil {
return nil, fmt.Errorf("cannot stat upload: %v", err)
}
pbar.Total = st.Size()
pbar.Units = pb.U_BYTES
pbar.Output = osStdout
r = pbar.NewProxyReader(file)
pbar.Start()
defer pbar.Finish()
}

fmt.Printf("Uploading %s to %s:%s\n", filename, bucketName, keyName)
uploadOutput, err := a.Upload(filename, bucketName, keyName)
return a.UploadFromReader(r, bucketName, keyName)
}

func UploadAndRegister(a AwsUploader, filename, bucketName, imageName string, pbar *pb.ProgressBar) error {
file, err := os.Open(filename)
if err != nil {
return fmt.Errorf("cannot upload: %v", err)
}
defer file.Close()

keyName := fmt.Sprintf("%s-%s", uuid.New().String(), filepath.Base(filename))
fmt.Fprintf(osStdout, "Uploading %s to %s:%s\n", filename, bucketName, keyName)
uploadOutput, err := doUpload(a, file, bucketName, keyName, pbar)
if err != nil {
return err
}
fmt.Printf("File uploaded to %s\n", aws.StringValue(&uploadOutput.Location))
fmt.Fprintf(osStdout, "File uploaded to %s\n", aws.StringValue(&uploadOutput.Location))

hostArch := arch.Current()
bootMode := ec2.BootModeValuesUefiPreferred
fmt.Printf("Registering AMI %s\n", imageName)
fmt.Fprintf(osStdout, "Registering AMI %s\n", imageName)
ami, snapshot, err := a.Register(imageName, bucketName, keyName, nil, hostArch.String(), &bootMode)
fmt.Printf("Deleted S3 object %s:%s\n", bucketName, keyName)
fmt.Printf("AMI registered: %s\nSnapshot ID: %s\n", aws.StringValue(ami), aws.StringValue(snapshot))
fmt.Fprintf(osStdout, "Deleted S3 object %s:%s\n", bucketName, keyName)
fmt.Fprintf(osStdout, "AMI registered: %s\nSnapshot ID: %s\n", aws.StringValue(ami), aws.StringValue(snapshot))
if err != nil {
return err
}
Expand Down
87 changes: 87 additions & 0 deletions bib/internal/uploader/aws_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
package uploader_test

import (
"bytes"
"io"
"io/ioutil"

Check failure on line 6 in bib/internal/uploader/aws_test.go

View workflow job for this annotation

GitHub Actions / ⌨ Lint

SA1019: "io/ioutil" has been deprecated since Go 1.19: As of Go 1.16, the same functionality is now provided by package io or package os, and those implementations should be preferred in new code. See the specific function documentation for details. (staticcheck)
"os"
"path/filepath"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"github.com/aws/aws-sdk-go/service/s3/s3manager"
"github.com/cheggaaa/pb"

"github.com/osbuild/bootc-image-builder/bib/internal/uploader"
)

type FakeAwsUploader struct {
uploadCalled int
registerCalled int
}

func (f *FakeAwsUploader) UploadFromReader(r io.Reader, bucketName, keyName string) (*s3manager.UploadOutput, error) {
f.uploadCalled++

if _, err := ioutil.ReadAll(r); err != nil {
panic(err)
}

return &s3manager.UploadOutput{Location: "some-location"}, nil
}

func (f *FakeAwsUploader) Register(name, bucket, key string, shareWith []string, rpmArch string, bootMode *string) (*string, *string, error) {
f.registerCalled++

s1 := "ret1"
s2 := "ret2"
return &s1, &s2, nil
}

func TestUploadAndRegisterNoProgressBar(t *testing.T) {
fakeStdout := bytes.NewBuffer(nil)
restore := uploader.MockOsStdout(fakeStdout)
defer restore()

fakeDiskFile := filepath.Join(t.TempDir(), "fake-disk.img")
err := os.WriteFile(fakeDiskFile, nil, 0644)
require.Nil(t, err)
fakeUploader := &FakeAwsUploader{}

err = uploader.UploadAndRegister(fakeUploader, fakeDiskFile, "bucketName", "imageName", nil)
require.Nil(t, err)

assert.Equal(t, fakeUploader.uploadCalled, 1)
assert.Equal(t, fakeUploader.registerCalled, 1)

assert.Contains(t, fakeStdout.String(), "Uploading ")
assert.Contains(t, fakeStdout.String(), "Registering AMI ")
}

func TestUploadAndRegisterProgressBar(t *testing.T) {
fakeStdout := bytes.NewBuffer(nil)
restore := uploader.MockOsStdout(fakeStdout)
defer restore()

fakeDiskFile := filepath.Join(t.TempDir(), "fake-disk.img")
err := os.WriteFile(fakeDiskFile, nil, 0644)
require.Nil(t, err)
err = os.Truncate(fakeDiskFile, 10*1024*1024)
require.Nil(t, err)

fakeUploader := &FakeAwsUploader{}

pbar := pb.New(0)

err = uploader.UploadAndRegister(fakeUploader, fakeDiskFile, "bucketName", "imageName", pbar)
require.Nil(t, err)

assert.Equal(t, fakeUploader.uploadCalled, 1)
assert.Equal(t, fakeUploader.registerCalled, 1)

assert.Contains(t, fakeStdout.String(), "Uploading ")
assert.Contains(t, fakeStdout.String(), "10.00 MiB / 10.00 MiB [============================================] 100.00%")
assert.Contains(t, fakeStdout.String(), "Registering AMI ")
}
13 changes: 13 additions & 0 deletions bib/internal/uploader/export_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package uploader

import (
"io"
)

func MockOsStdout(new io.Writer) (restore func()) {
saved := osStdout
osStdout = new
return func() {
osStdout = saved
}
}
25 changes: 22 additions & 3 deletions test/test_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ class ImageBuildResult(NamedTuple):
img_path: str
username: str
password: str
bib_output: str
journal_output: str
metadata: dict = {}

Expand Down Expand Up @@ -113,7 +114,7 @@ def image_type_fixture(tmpdir_factory, build_container, request, force_aws_uploa
raise RuntimeError("AWS credentials not available (upload forced)")

# run container to deploy an image into a bootable disk and upload to a cloud service if applicable
subprocess.check_call([
p = subprocess.Popen([
"podman", "run", "--rm",
"--privileged",
"--security-opt", "label=type:unconfined_t",
Expand All @@ -125,7 +126,18 @@ def image_type_fixture(tmpdir_factory, build_container, request, force_aws_uploa
"--config", "/output/config.json",
"--type", image_type,
*upload_args,
])
], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
# not using subprocss.check_output() to ensure we get live output
# during the text
bib_output = ""
while True:
line = p.stdout.readline()
if not line:
break
print(line, end="")
bib_output += line
p.wait(timeout=10)

journal_output = testutil.journal_after_cursor(cursor)
metadata = {}
if image_type == "ami" and upload_args:
Expand All @@ -137,7 +149,7 @@ def del_ami():

journal_log_path.write_text(journal_output, encoding="utf8")

return ImageBuildResult(image_type, generated_img, username, password, journal_output, metadata)
return ImageBuildResult(image_type, generated_img, username, password, bib_output, journal_output, metadata)


def test_container_builds(build_container):
Expand Down Expand Up @@ -171,6 +183,13 @@ def test_ami_boots_in_aws(image_type, force_aws_upload):
raise RuntimeError("AWS credentials not available")
pytest.skip("AWS credentials not available (upload not forced)")

# check that upload progress is in the output log. Uploads looks like:
#
# Uploading /output/image/disk.raw to bootc-image-builder-ci:aac64b64-6e57-47df-9730-54763061d84b-disk.raw
# 0 B / 10.00 GiB 0.00%
# In the tests with no pty no progress bar is shown in the output just
# xx / yy zz%
assert " 100.00%\n" in image_type.bib_output
with AWS(image_type.metadata["ami_id"]) as test_vm:
exit_status, _ = test_vm.run("true", user=image_type.username, password=image_type.password)
assert exit_status == 0
Expand Down

0 comments on commit eb51d1d

Please sign in to comment.