Skip to content

Commit a422cc0

Browse files
ammarioammario
authored
Stream template upload (coder#6035)
Resolves coder#5718
1 parent 77fd34b commit a422cc0

12 files changed

+69
-46
lines changed

‎cli/templateinit.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli
22

33
import (
4+
"bytes"
45
"fmt"
56
"os"
67
"path/filepath"
@@ -70,7 +71,7 @@ func templateInit() *cobra.Command {
7071
if err != nil {
7172
return err
7273
}
73-
err = provisionersdk.Untar(directory, archive)
74+
err = provisionersdk.Untar(directory, bytes.NewReader(archive))
7475
if err != nil {
7576
return err
7677
}

‎cli/templatepush.go

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli
22

33
import (
4+
"bufio"
45
"fmt"
56
"io"
67
"os"
@@ -34,14 +35,17 @@ func (pf *templateUploadFlags) stdin() bool {
3435

3536
func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Client) (*codersdk.UploadResponse, error) {
3637
var (
37-
content []byte
38-
err error
38+
content io.Reader
39+
pipeErrCh = make(chan error, 1)
3940
)
4041
if pf.stdin() {
41-
content, err = io.ReadAll(cmd.InOrStdin())
42+
content = cmd.InOrStdin()
43+
// No piping if reading from stdin.
44+
pipeErrCh <- nil
45+
close(pipeErrCh)
4246
} else {
4347
prettyDir := prettyDirectoryPath(pf.directory)
44-
_, err = cliui.Prompt(cmd, cliui.PromptOptions{
48+
_, err := cliui.Prompt(cmd, cliui.PromptOptions{
4549
Text: fmt.Sprintf("Upload %q?", prettyDir),
4650
IsConfirm: true,
4751
Default: cliui.ConfirmYes,
@@ -50,10 +54,16 @@ func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Clien
5054
return nil, err
5155
}
5256

53-
content, err = provisionersdk.Tar(pf.directory, provisionersdk.TemplateArchiveLimit)
54-
}
55-
if err != nil {
56-
return nil, xerrors.Errorf("read tar: %w", err)
57+
pipeReader, pipeWriter := io.Pipe()
58+
go func() {
59+
defer pipeWriter.Close()
60+
defer close(pipeErrCh)
61+
bufWr := bufio.NewWriter(pipeWriter)
62+
defer bufWr.Flush()
63+
pipeErrCh <- provisionersdk.Tar(bufWr, pf.directory, provisionersdk.TemplateArchiveLimit)
64+
}()
65+
defer pipeReader.Close()
66+
content = pipeReader
5767
}
5868

5969
spin := spinner.New(spinner.CharSets[5], 100*time.Millisecond)
@@ -66,6 +76,9 @@ func (pf *templateUploadFlags) upload(cmd *cobra.Command, client *codersdk.Clien
6676
if err != nil {
6777
return nil, xerrors.Errorf("upload: %w", err)
6878
}
79+
if err = <-pipeErrCh; err != nil {
80+
return nil, xerrors.Errorf("pipe: %w", err)
81+
}
6982
return &resp, nil
7083
}
7184

‎coderd/coderdtest/authorize.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderdtest
22

33
import (
4+
"bytes"
45
"context"
56
"fmt"
67
"io"
@@ -384,7 +385,7 @@ func NewAuthTester(ctx context.Context, t *testing.T, client *codersdk.Client, a
384385
template := CreateTemplate(t, client, admin.OrganizationID, version.ID)
385386
workspace := CreateWorkspace(t, client, admin.OrganizationID, template.ID)
386387
AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
387-
file, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024))
388+
file, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024)))
388389
require.NoError(t, err, "upload file")
389390
workspace, err = client.Workspace(ctx, workspace.ID)
390391
require.NoError(t, err, "workspace resources")

‎coderd/coderdtest/coderdtest.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -525,7 +525,7 @@ func CreateTemplateVersion(t *testing.T, client *codersdk.Client, organizationID
525525
t.Helper()
526526
data, err := echo.Tar(res)
527527
require.NoError(t, err)
528-
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data)
528+
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data))
529529
require.NoError(t, err)
530530
templateVersion, err := client.CreateTemplateVersion(context.Background(), organizationID, codersdk.CreateTemplateVersionRequest{
531531
FileID: file.ID,
@@ -572,7 +572,7 @@ func CreateTemplate(t *testing.T, client *codersdk.Client, organization uuid.UUI
572572
func UpdateTemplateVersion(t *testing.T, client *codersdk.Client, organizationID uuid.UUID, res *echo.Responses, templateID uuid.UUID) codersdk.TemplateVersion {
573573
data, err := echo.Tar(res)
574574
require.NoError(t, err)
575-
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data)
575+
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data))
576576
require.NoError(t, err)
577577
templateVersion, err := client.CreateTemplateVersion(context.Background(), organizationID, codersdk.CreateTemplateVersionRequest{
578578
TemplateID: templateID,

‎coderd/files_test.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderd_test
22

33
import (
4+
"bytes"
45
"context"
56
"net/http"
67
"testing"
@@ -23,7 +24,7 @@ func TestPostFiles(t *testing.T) {
2324
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
2425
defer cancel()
2526

26-
_, err := client.Upload(ctx, "bad", []byte{'a'})
27+
_, err := client.Upload(ctx, "bad", bytes.NewReader([]byte{'a'}))
2728
require.Error(t, err)
2829
})
2930

@@ -35,7 +36,7 @@ func TestPostFiles(t *testing.T) {
3536
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
3637
defer cancel()
3738

38-
_, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024))
39+
_, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024)))
3940
require.NoError(t, err)
4041
})
4142

@@ -48,9 +49,9 @@ func TestPostFiles(t *testing.T) {
4849
defer cancel()
4950

5051
data := make([]byte, 1024)
51-
_, err := client.Upload(ctx, codersdk.ContentTypeTar, data)
52+
_, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data))
5253
require.NoError(t, err)
53-
_, err = client.Upload(ctx, codersdk.ContentTypeTar, data)
54+
_, err = client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data))
5455
require.NoError(t, err)
5556
})
5657
}
@@ -79,7 +80,7 @@ func TestDownload(t *testing.T) {
7980
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
8081
defer cancel()
8182

82-
resp, err := client.Upload(ctx, codersdk.ContentTypeTar, make([]byte, 1024))
83+
resp, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(make([]byte, 1024)))
8384
require.NoError(t, err)
8485
data, contentType, err := client.Download(ctx, resp.ID)
8586
require.NoError(t, err)

‎coderd/templateversions_test.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderd_test
22

33
import (
4+
"bytes"
45
"context"
56
"net/http"
67
"testing"
@@ -108,7 +109,7 @@ func TestPostTemplateVersionsByOrganization(t *testing.T) {
108109
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
109110
defer cancel()
110111

111-
file, err := client.Upload(ctx, codersdk.ContentTypeTar, data)
112+
file, err := client.Upload(ctx, codersdk.ContentTypeTar, bytes.NewReader(data))
112113
require.NoError(t, err)
113114
version, err := client.CreateTemplateVersion(ctx, user.OrganizationID, codersdk.CreateTemplateVersionRequest{
114115
Name: "bananas",
@@ -895,7 +896,7 @@ func TestPaginatedTemplateVersions(t *testing.T) {
895896
templateVersionIDs := make([]uuid.UUID, total)
896897
data, err := echo.Tar(nil)
897898
require.NoError(t, err)
898-
file, err := client.Upload(egCtx, codersdk.ContentTypeTar, data)
899+
file, err := client.Upload(egCtx, codersdk.ContentTypeTar, bytes.NewReader(data))
899900
require.NoError(t, err)
900901
for i := 0; i < total; i++ {
901902
i := i

‎codersdk/client.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,18 +109,20 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac
109109

110110
var r io.Reader
111111
if body != nil {
112-
if data, ok := body.([]byte); ok {
112+
switch data := body.(type) {
113+
case io.Reader:
114+
r = data
115+
case []byte:
113116
r = bytes.NewReader(data)
114-
} else {
115-
// Assume JSON if not bytes.
117+
default:
118+
// Assume JSON in all other cases.
116119
buf := bytes.NewBuffer(nil)
117120
enc := json.NewEncoder(buf)
118121
enc.SetEscapeHTML(false)
119122
err = enc.Encode(body)
120123
if err != nil {
121124
return nil, xerrors.Errorf("encode body: %w", err)
122125
}
123-
124126
r = buf
125127
}
126128
}

‎codersdk/files.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ type UploadResponse struct {
2121

2222
// Upload uploads an arbitrary file with the content type provided.
2323
// This is used to upload a source-code archive.
24-
func (c *Client) Upload(ctx context.Context, contentType string, content []byte) (UploadResponse, error) {
25-
res, err := c.Request(ctx, http.MethodPost, "/api/v2/files", content, func(r *http.Request) {
24+
func (c *Client) Upload(ctx context.Context, contentType string, rd io.Reader) (UploadResponse, error) {
25+
res, err := c.Request(ctx, http.MethodPost, "/api/v2/files", rd, func(r *http.Request) {
2626
r.Header.Set("Content-Type", contentType)
2727
})
2828
if err != nil {

‎enterprise/coderd/provisionerdaemons_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderd_test
22

33
import (
4+
"bytes"
45
"context"
56
"net/http"
67
"testing"
@@ -119,7 +120,7 @@ func TestProvisionerDaemonServe(t *testing.T) {
119120
}},
120121
})
121122
require.NoError(t, err)
122-
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data)
123+
file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data))
123124
require.NoError(t, err)
124125

125126
version, err := client.CreateTemplateVersion(context.Background(), user.OrganizationID, codersdk.CreateTemplateVersionRequest{

‎enterprise/coderd/templates_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package coderd_test
22

33
import (
4+
"bytes"
45
"context"
56
"net/http"
67
"testing"
@@ -300,7 +301,7 @@ func TestTemplateACL(t *testing.T) {
300301

301302
data, err := echo.Tar(nil)
302303
require.NoError(t, err)
303-
file, err := client1.Upload(context.Background(), codersdk.ContentTypeTar, data)
304+
file, err := client1.Upload(context.Background(), codersdk.ContentTypeTar, bytes.NewReader(data))
304305
require.NoError(t, err)
305306

306307
_, err = client1.CreateTemplateVersion(ctx, user.OrganizationID, codersdk.CreateTemplateVersionRequest{

‎provisionersdk/archive.go

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package provisionersdk
22

33
import (
44
"archive/tar"
5-
"bytes"
65
"io"
76
"os"
87
"path/filepath"
@@ -32,25 +31,24 @@ func dirHasExt(dir string, ext string) (bool, error) {
3231
}
3332

3433
// Tar archives a Terraform directory.
35-
func Tar(directory string, limit int64) ([]byte, error) {
36-
var buffer bytes.Buffer
37-
tarWriter := tar.NewWriter(&buffer)
34+
func Tar(w io.Writer, directory string, limit int64) error {
35+
tarWriter := tar.NewWriter(w)
3836
totalSize := int64(0)
3937

4038
const tfExt = ".tf"
4139
hasTf, err := dirHasExt(directory, tfExt)
4240
if err != nil {
43-
return nil, err
41+
return err
4442
}
4543
if !hasTf {
4644
absPath, err := filepath.Abs(directory)
4745
if err != nil {
48-
return nil, err
46+
return err
4947
}
5048

5149
// Show absolute path to aid in debugging. E.g. showing "." is
5250
// useless.
53-
return nil, xerrors.Errorf(
51+
return xerrors.Errorf(
5452
"%s is not a valid template since it has no %s files",
5553
absPath, tfExt,
5654
)
@@ -111,20 +109,20 @@ func Tar(directory string, limit int64) ([]byte, error) {
111109
return data.Close()
112110
})
113111
if err != nil {
114-
return nil, err
112+
return err
115113
}
116114
err = tarWriter.Flush()
117115
if err != nil {
118-
return nil, err
116+
return err
119117
}
120-
return buffer.Bytes(), nil
118+
return nil
121119
}
122120

123121
// Untar extracts the archive to a provided directory.
124-
func Untar(directory string, archive []byte) error {
125-
reader := tar.NewReader(bytes.NewReader(archive))
122+
func Untar(directory string, r io.Reader) error {
123+
tarReader := tar.NewReader(r)
126124
for {
127-
header, err := reader.Next()
125+
header, err := tarReader.Next()
128126
if xerrors.Is(err, io.EOF) {
129127
return nil
130128
}
@@ -149,7 +147,7 @@ func Untar(directory string, archive []byte) error {
149147
return err
150148
}
151149
// Max file size of 10MB.
152-
_, err = io.CopyN(file, reader, (1<<20)*10)
150+
_, err = io.CopyN(file, tarReader, (1<<20)*10)
153151
if xerrors.Is(err, io.EOF) {
154152
err = nil
155153
}

‎provisionersdk/archive_test.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package provisionersdk_test
22

33
import (
4+
"bytes"
5+
"io"
46
"os"
57
"path/filepath"
68
"testing"
@@ -18,7 +20,7 @@ func TestTar(t *testing.T) {
1820
file, err := os.CreateTemp(dir, "")
1921
require.NoError(t, err)
2022
_ = file.Close()
21-
_, err = provisionersdk.Tar(dir, 1024)
23+
err = provisionersdk.Tar(io.Discard, dir, 1024)
2224
require.Error(t, err)
2325
})
2426
t.Run("Valid", func(t *testing.T) {
@@ -27,7 +29,7 @@ func TestTar(t *testing.T) {
2729
file, err := os.CreateTemp(dir, "*.tf")
2830
require.NoError(t, err)
2931
_ = file.Close()
30-
_, err = provisionersdk.Tar(dir, 1024)
32+
err = provisionersdk.Tar(io.Discard, dir, 1024)
3133
require.NoError(t, err)
3234
})
3335
t.Run("HiddenFiles", func(t *testing.T) {
@@ -71,10 +73,11 @@ func TestTar(t *testing.T) {
7173
file.Name, err = filepath.Rel(dir, tmpFile.Name())
7274
require.NoError(t, err)
7375
}
74-
content, err := provisionersdk.Tar(dir, 1024)
76+
archive := new(bytes.Buffer)
77+
err := provisionersdk.Tar(archive, dir, 1024)
7578
require.NoError(t, err)
7679
dir = t.TempDir()
77-
err = provisionersdk.Untar(dir, content)
80+
err = provisionersdk.Untar(dir, archive)
7881
require.NoError(t, err)
7982
for _, file := range files {
8083
_, err = os.Stat(filepath.Join(dir, file.Name))
@@ -94,7 +97,8 @@ func TestUntar(t *testing.T) {
9497
file, err := os.CreateTemp(dir, "*.tf")
9598
require.NoError(t, err)
9699
_ = file.Close()
97-
archive, err := provisionersdk.Tar(dir, 1024)
100+
archive := new(bytes.Buffer)
101+
err = provisionersdk.Tar(archive, dir, 1024)
98102
require.NoError(t, err)
99103
dir = t.TempDir()
100104
err = provisionersdk.Untar(dir, archive)

0 commit comments

Comments
 (0)