From 5858a79c3853633e50ab5f921cc3525c374cad74 Mon Sep 17 00:00:00 2001 From: Ammar Bandukwala Date: Sat, 4 Feb 2023 19:35:56 +0000 Subject: [PATCH] Stream template upload Resolves #5718 --- cli/root_test.go | 2 -- cli/templateinit.go | 3 +- cli/templatepush.go | 29 ++++++++++++++------ coderd/coderdtest/authorize.go | 3 +- coderd/coderdtest/coderdtest.go | 4 +-- coderd/files_test.go | 11 ++++---- coderd/templateversions_test.go | 5 ++-- codersdk/client.go | 10 ++++--- codersdk/files.go | 4 +-- enterprise/coderd/provisionerdaemons_test.go | 3 +- enterprise/coderd/templates_test.go | 3 +- provisionersdk/archive.go | 26 ++++++++---------- provisionersdk/archive_test.go | 14 ++++++---- 13 files changed, 69 insertions(+), 48 deletions(-) diff --git a/cli/root_test.go b/cli/root_test.go index 7f5442233b830..85863b656a2aa 100644 --- a/cli/root_test.go +++ b/cli/root_test.go @@ -30,8 +30,6 @@ var updateGoldenFiles = flag.Bool("update", false, "update .golden files") //nolint:tparallel,paralleltest // These test sets env vars. func TestCommandHelp(t *testing.T) { - t.Parallel() - commonEnv := map[string]string{ "CODER_CONFIG_DIR": "/tmp/coder-cli-test-config", } diff --git a/cli/templateinit.go b/cli/templateinit.go index 39a8f72662bde..b69fa1086bbb0 100644 --- a/cli/templateinit.go +++ b/cli/templateinit.go @@ -1,6 +1,7 @@ package cli import ( + "bytes" "fmt" "os" "path/filepath" @@ -70,7 +71,7 @@ func templateInit() *cobra.Command { if err != nil { return err } - err = provisionersdk.Untar(directory, archive) + err = provisionersdk.Untar(directory, bytes.NewReader(archive)) if err != nil { return err } diff --git a/cli/templatepush.go b/cli/templatepush.go index 3ba612d7e674f..11bf8209c38ed 100644 --- a/cli/templatepush.go +++ b/cli/templatepush.go @@ -1,6 +1,7 @@ package cli import ( + "bufio" "fmt" "io" "os" @@ -34,14 +35,17 @@ func (pf *templateUploadFlags) stdin() bool { func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Client) (*codersdk.UploadResponse, error) { var ( - content []byte - err error + content io.Reader + pipeErrCh = make(chan error, 1) ) if pf.stdin() { - content, err = io.ReadAll(cmd.InOrStdin()) + content = cmd.InOrStdin() + // No piping if reading from stdin. + pipeErrCh <- nil + close(pipeErrCh) } else { prettyDir := prettyDirectoryPath(pf.directory) - _, err = cliui.Prompt(cmd, cliui.PromptOptions{ + _, err := cliui.Prompt(cmd, cliui.PromptOptions{ Text: fmt.Sprintf("Upload %q?", prettyDir), IsConfirm: true, Default: cliui.ConfirmYes, @@ -50,10 +54,16 @@ func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Clien return nil, err } - content, err = provisionersdk.Tar(pf.directory, provisionersdk.TemplateArchiveLimit) - } - if err != nil { - return nil, xerrors.Errorf("read tar: %w", err) + pipeReader, pipeWriter := io.Pipe() + go func() { + defer pipeWriter.Close() + defer close(pipeErrCh) + bufWr := bufio.NewWriter(pipeWriter) + defer bufWr.Flush() + pipeErrCh <- provisionersdk.Tar(bufWr, pf.directory, provisionersdk.TemplateArchiveLimit) + }() + defer pipeReader.Close() + content = pipeReader } spin := spinner.New(spinner.CharSets[5], 100*time.Millisecond) @@ -66,6 +76,9 @@ func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Clien if err != nil { return nil, xerrors.Errorf("upload: %w", err) } + if err = <-pipeErrCh; err != nil { + return nil, xerrors.Errorf("pipe: %w", err) + } return &resp, nil } diff --git a/coderd/coderdtest/authorize.go b/coderd/coderdtest/authorize.go index 14019815d2a43..d763983391b86 100644 --- a/coderd/coderdtest/authorize.go +++ b/coderd/coderdtest/authorize.go @@ -1,6 +1,7 @@ package coderdtest import ( + "bytes" "context" "fmt" "io" @@ -378,7 +379,7 @@ func NewAuthTester(ctx context.Context, t *testing.T, client *codersdk.Client, a template := CreateTemplate(t, client, admin.OrganizationID, version.ID) workspace := CreateWorkspace(t, client, admin.OrganizationID, template.ID) AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID) - file, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024)) + file, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024))) require.NoError(t, err, "upload file") workspace, err = client.Workspace(ctx, workspace.ID) require.NoError(t, err, "workspace resources") diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 38ca73ec5cfad..0999355c1034d 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -525,7 +525,7 @@ func CreateTemplateVersion(t *testing.T, client *codersdk.Client, organizationID t.Helper() data, err := echo.Tar(res) require.NoError(t, err) - file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data) + file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) templateVersion, err := client.CreateTemplateVersion(context.Background(), organizationID, codersdk.CreateTemplateVersionRequest{ FileID: file.ID, @@ -572,7 +572,7 @@ func CreateTemplate(t *testing.T, client *codersdk.Client, organization uuid.UUI func UpdateTemplateVersion(t *testing.T, client *codersdk.Client, organizationID uuid.UUID, res *echo.Responses, templateID uuid.UUID) codersdk.TemplateVersion { data, err := echo.Tar(res) require.NoError(t, err) - file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data) + file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) templateVersion, err := client.CreateTemplateVersion(context.Background(), organizationID, codersdk.CreateTemplateVersionRequest{ TemplateID: templateID, diff --git a/coderd/files_test.go b/coderd/files_test.go index b3a3953a43972..0841785c8c660 100644 --- a/coderd/files_test.go +++ b/coderd/files_test.go @@ -1,6 +1,7 @@ package coderd_test import ( + "bytes" "context" "net/http" "testing" @@ -23,7 +24,7 @@ func TestPostFiles(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - _, err := client.Upload(ctx, "bad", []byte{'a'}) + _, err := client.Upload(ctx, "bad", bytes.NewReader([]byte{'a'})) require.Error(t, err) }) @@ -35,7 +36,7 @@ func TestPostFiles(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - _, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024)) + _, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024))) require.NoError(t, err) }) @@ -48,9 +49,9 @@ func TestPostFiles(t *testing.T) { defer cancel() data := make([]byte, 1024) - _, err := client.Upload(ctx, codersdk.ContentTypeTar, data) + _, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) - _, err = client.Upload(ctx, codersdk.ContentTypeTar, data) + _, err = client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) }) } @@ -79,7 +80,7 @@ func TestDownload(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - resp, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024)) + resp, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024))) require.NoError(t, err) data, contentType, err := client.Download(ctx, resp.ID) require.NoError(t, err) diff --git a/coderd/templateversions_test.go b/coderd/templateversions_test.go index adf35aeb1c21b..7e0dc7a4bea85 100644 --- a/coderd/templateversions_test.go +++ b/coderd/templateversions_test.go @@ -1,6 +1,7 @@ package coderd_test import ( + "bytes" "context" "net/http" "testing" @@ -108,7 +109,7 @@ func TestPostTemplateVersionsByOrganization(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - file, err := client.Upload(ctx, codersdk.ContentTypeTar, data) + file, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) version, err := client.CreateTemplateVersion(ctx, user.OrganizationID, codersdk.CreateTemplateVersionRequest{ Name: "bananas", @@ -895,7 +896,7 @@ func TestPaginatedTemplateVersions(t *testing.T) { templateVersionIDs := make([]uuid.UUID, total) data, err := echo.Tar(nil) require.NoError(t, err) - file, err := client.Upload(egCtx, codersdk.ContentTypeTar, data) + file, err := client.Upload(egCtx, codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) for i := 0; i < total; i++ { i := i diff --git a/codersdk/client.go b/codersdk/client.go index dac6fdc533f1c..c95fa4f318d1b 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -109,10 +109,13 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac var r io.Reader if body != nil { - if data, ok := body.([]byte); ok { + switch data := body.(type) { + case io.Reader: + r = data + case []byte: r = bytes.NewReader(data) - } else { - // Assume JSON if not bytes. + default: + // Assume JSON in all other cases. buf := bytes.NewBuffer(nil) enc := json.NewEncoder(buf) enc.SetEscapeHTML(false) @@ -120,7 +123,6 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac if err != nil { return nil, xerrors.Errorf("encode body: %w", err) } - r = buf } } diff --git a/codersdk/files.go b/codersdk/files.go index cc66132df95e7..3525e9d785d6e 100644 --- a/codersdk/files.go +++ b/codersdk/files.go @@ -21,8 +21,8 @@ type UploadResponse struct { // Upload uploads an arbitrary file with the content type provided. // This is used to upload a source-code archive. -func (c *Client) Upload(ctx context.Context, contentType string, content []byte) (UploadResponse, error) { - res, err := c.Request(ctx, http.MethodPost, "/api/v2/files", content, func(r *http.Request) { +func (c *Client) Upload(ctx context.Context, contentType string, rd io.Reader) (UploadResponse, error) { + res, err := c.Request(ctx, http.MethodPost, "/api/v2/files", rd, func(r *http.Request) { r.Header.Set("Content-Type", contentType) }) if err != nil { diff --git a/enterprise/coderd/provisionerdaemons_test.go b/enterprise/coderd/provisionerdaemons_test.go index 65c5eaebb26ca..73a72bc46f889 100644 --- a/enterprise/coderd/provisionerdaemons_test.go +++ b/enterprise/coderd/provisionerdaemons_test.go @@ -1,6 +1,7 @@ package coderd_test import ( + "bytes" "context" "net/http" "testing" @@ -119,7 +120,7 @@ func TestProvisionerDaemonServe(t *testing.T) { }}, }) require.NoError(t, err) - file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data) + file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) version, err := client.CreateTemplateVersion(context.Background(), user.OrganizationID, codersdk.CreateTemplateVersionRequest{ diff --git a/enterprise/coderd/templates_test.go b/enterprise/coderd/templates_test.go index 1e591fcd30b79..4a9935c86935c 100644 --- a/enterprise/coderd/templates_test.go +++ b/enterprise/coderd/templates_test.go @@ -1,6 +1,7 @@ package coderd_test import ( + "bytes" "context" "net/http" "testing" @@ -300,7 +301,7 @@ func TestTemplateACL(t *testing.T) { data, err := echo.Tar(nil) require.NoError(t, err) - file, err := client1.Upload(context.Background(), codersdk.ContentTypeTar, data) + file, err := client1.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data)) require.NoError(t, err) _, err = client1.CreateTemplateVersion(ctx, user.OrganizationID, codersdk.CreateTemplateVersionRequest{ diff --git a/provisionersdk/archive.go b/provisionersdk/archive.go index e5513d0f6e8b3..8aa0c1b46f92c 100644 --- a/provisionersdk/archive.go +++ b/provisionersdk/archive.go @@ -2,7 +2,6 @@ package provisionersdk import ( "archive/tar" - "bytes" "io" "os" "path/filepath" @@ -32,25 +31,24 @@ func dirHasExt(dir string, ext string) (bool, error) { } // Tar archives a Terraform directory. -func Tar(directory string, limit int64) ([]byte, error) { - var buffer bytes.Buffer - tarWriter := tar.NewWriter(&buffer) +func Tar(w io.Writer, directory string, limit int64) error { + tarWriter := tar.NewWriter(w) totalSize := int64(0) const tfExt = ".tf" hasTf, err := dirHasExt(directory, tfExt) if err != nil { - return nil, err + return err } if !hasTf { absPath, err := filepath.Abs(directory) if err != nil { - return nil, err + return err } // Show absolute path to aid in debugging. E.g. showing "." is // useless. - return nil, xerrors.Errorf( + return xerrors.Errorf( "%s is not a valid template since it has no %s files", absPath, tfExt, ) @@ -111,20 +109,20 @@ func Tar(directory string, limit int64) ([]byte, error) { return data.Close() }) if err != nil { - return nil, err + return err } err = tarWriter.Flush() if err != nil { - return nil, err + return err } - return buffer.Bytes(), nil + return nil } // Untar extracts the archive to a provided directory. -func Untar(directory string, archive []byte) error { - reader := tar.NewReader(bytes.NewReader(archive)) +func Untar(directory string, r io.Reader) error { + tarReader := tar.NewReader(r) for { - header, err := reader.Next() + header, err := tarReader.Next() if xerrors.Is(err, io.EOF) { return nil } @@ -149,7 +147,7 @@ func Untar(directory string, archive []byte) error { return err } // Max file size of 10MB. - _, err = io.CopyN(file, reader, (1<<20)*10) + _, err = io.CopyN(file, tarReader, (1<<20)*10) if xerrors.Is(err, io.EOF) { err = nil } diff --git a/provisionersdk/archive_test.go b/provisionersdk/archive_test.go index 4d37dd7ac5843..1d0b008b6a0c7 100644 --- a/provisionersdk/archive_test.go +++ b/provisionersdk/archive_test.go @@ -1,6 +1,8 @@ package provisionersdk_test import ( + "bytes" + "io" "os" "path/filepath" "testing" @@ -18,7 +20,7 @@ func TestTar(t *testing.T) { file, err := os.CreateTemp(dir, "") require.NoError(t, err) _ = file.Close() - _, err = provisionersdk.Tar(dir, 1024) + err = provisionersdk.Tar(io.Discard, dir, 1024) require.Error(t, err) }) t.Run("Valid", func(t *testing.T) { @@ -27,7 +29,7 @@ func TestTar(t *testing.T) { file, err := os.CreateTemp(dir, "*.tf") require.NoError(t, err) _ = file.Close() - _, err = provisionersdk.Tar(dir, 1024) + err = provisionersdk.Tar(io.Discard, dir, 1024) require.NoError(t, err) }) t.Run("HiddenFiles", func(t *testing.T) { @@ -71,10 +73,11 @@ func TestTar(t *testing.T) { file.Name, err = filepath.Rel(dir, tmpFile.Name()) require.NoError(t, err) } - content, err := provisionersdk.Tar(dir, 1024) + archive := new(bytes.Buffer) + err := provisionersdk.Tar(archive, dir, 1024) require.NoError(t, err) dir = t.TempDir() - err = provisionersdk.Untar(dir, content) + err = provisionersdk.Untar(dir, archive) require.NoError(t, err) for _, file := range files { _, err = os.Stat(filepath.Join(dir, file.Name)) @@ -94,7 +97,8 @@ func TestUntar(t *testing.T) { file, err := os.CreateTemp(dir, "*.tf") require.NoError(t, err) _ = file.Close() - archive, err := provisionersdk.Tar(dir, 1024) + archive := new(bytes.Buffer) + err = provisionersdk.Tar(archive, dir, 1024) require.NoError(t, err) dir = t.TempDir() err = provisionersdk.Untar(dir, archive)