github.com/kubeshop/testkube@v1.17.23/cmd/tcl/testworkflow-toolkit/artifacts/cloud_uploader.go (about)

     1  // Copyright 2024 Testkube.
     2  //
     3  // Licensed as a Testkube Pro file under the Testkube Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //	https://github.com/kubeshop/testkube/blob/main/licenses/TCL.txt
     8  
     9  package artifacts
    10  
    11  import (
    12  	"bytes"
    13  	"context"
    14  	"crypto/tls"
    15  	"encoding/json"
    16  	"fmt"
    17  	"io"
    18  	"net/http"
    19  	"sync"
    20  	"sync/atomic"
    21  	"time"
    22  
    23  	"github.com/pkg/errors"
    24  
    25  	"github.com/kubeshop/testkube/cmd/tcl/testworkflow-toolkit/env"
    26  	"github.com/kubeshop/testkube/pkg/cloud/data/artifact"
    27  	cloudexecutor "github.com/kubeshop/testkube/pkg/cloud/data/executor"
    28  	"github.com/kubeshop/testkube/pkg/ui"
    29  )
    30  
    31  type CloudUploaderRequestEnhancer = func(req *http.Request, path string, size int64)
    32  
    33  func NewCloudUploader(opts ...CloudUploaderOpt) Uploader {
    34  	uploader := &cloudUploader{
    35  		parallelism:  1,
    36  		reqEnhancers: make([]CloudUploaderRequestEnhancer, 0),
    37  	}
    38  	for _, opt := range opts {
    39  		opt(uploader)
    40  	}
    41  	return uploader
    42  }
    43  
    44  type cloudUploader struct {
    45  	client       cloudexecutor.Executor
    46  	wg           sync.WaitGroup
    47  	sema         chan struct{}
    48  	parallelism  int
    49  	error        atomic.Bool
    50  	reqEnhancers []CloudUploaderRequestEnhancer
    51  }
    52  
    53  func (d *cloudUploader) Start() (err error) {
    54  	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
    55  	defer cancel()
    56  	d.client = env.Cloud(ctx)
    57  	d.sema = make(chan struct{}, d.parallelism)
    58  	return err
    59  }
    60  
    61  func (d *cloudUploader) getSignedURL(name, contentType string) (string, error) {
    62  	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
    63  	defer cancel()
    64  	response, err := d.client.Execute(ctx, artifact.CmdScraperPutObjectSignedURL, &artifact.PutObjectSignedURLRequest{
    65  		Object:           name,
    66  		ExecutionID:      env.ExecutionId(),
    67  		TestWorkflowName: env.WorkflowName(),
    68  		ContentType:      contentType,
    69  	})
    70  	if err != nil {
    71  		return "", err
    72  	}
    73  	var commandResponse artifact.PutObjectSignedURLResponse
    74  	if err := json.Unmarshal(response, &commandResponse); err != nil {
    75  		return "", err
    76  	}
    77  	return commandResponse.URL, nil
    78  }
    79  
    80  func (d *cloudUploader) getContentType(path string, size int64) string {
    81  	req, err := http.NewRequestWithContext(context.Background(), http.MethodPut, "/", &bytes.Buffer{})
    82  	if err != nil {
    83  		return ""
    84  	}
    85  	for _, r := range d.reqEnhancers {
    86  		r(req, path, size)
    87  	}
    88  	contentType := req.Header.Get("Content-Type")
    89  	if contentType == "" {
    90  		return "application/octet-stream"
    91  	}
    92  	return contentType
    93  }
    94  
    95  func (d *cloudUploader) putObject(url string, path string, file io.Reader, size int64) error {
    96  	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
    97  	defer cancel()
    98  	req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, file)
    99  	if err != nil {
   100  		return err
   101  	}
   102  	for _, r := range d.reqEnhancers {
   103  		r(req, path, size)
   104  	}
   105  	req.ContentLength = size
   106  	if req.Header.Get("Content-Type") == "" {
   107  		req.Header.Set("Content-Type", "application/octet-stream")
   108  	}
   109  	tr := http.DefaultTransport.(*http.Transport).Clone()
   110  	tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
   111  	client := &http.Client{Transport: tr}
   112  	res, err := client.Do(req)
   113  	if err != nil {
   114  		return err
   115  	}
   116  	if res.StatusCode != http.StatusOK {
   117  		b, _ := io.ReadAll(res.Body)
   118  		return errors.Errorf("failed saving file: status code: %d / message: %s", res.StatusCode, string(b))
   119  	}
   120  	return nil
   121  }
   122  
   123  func (d *cloudUploader) upload(path string, file io.Reader, size int64) {
   124  	url, err := d.getSignedURL(path, d.getContentType(path, size))
   125  	if err != nil {
   126  		d.error.Store(true)
   127  		ui.Errf("%s: failed: get signed URL: %s", path, err.Error())
   128  		return
   129  	}
   130  	err = d.putObject(url, path, file, size)
   131  	if err != nil {
   132  		d.error.Store(true)
   133  		ui.Errf("%s: failed: store file: %s", path, err.Error())
   134  		return
   135  	}
   136  }
   137  
   138  func (d *cloudUploader) Add(path string, file io.ReadCloser, size int64) error {
   139  	d.wg.Add(1)
   140  	d.sema <- struct{}{}
   141  	go func() {
   142  		d.upload(path, file, size)
   143  		_ = file.Close()
   144  		d.wg.Done()
   145  		<-d.sema
   146  	}()
   147  	return nil
   148  }
   149  
   150  func (d *cloudUploader) End() error {
   151  	d.wg.Wait()
   152  	if d.error.Load() {
   153  		return fmt.Errorf("upload failed")
   154  	}
   155  	return nil
   156  }
   157  
   158  type CloudUploaderOpt = func(uploader *cloudUploader)
   159  
   160  func WithParallelismCloud(parallelism int) CloudUploaderOpt {
   161  	return func(uploader *cloudUploader) {
   162  		if parallelism < 1 {
   163  			parallelism = 1
   164  		}
   165  		uploader.parallelism = parallelism
   166  	}
   167  }
   168  
   169  func WithRequestEnhancerCloud(enhancer CloudUploaderRequestEnhancer) CloudUploaderOpt {
   170  	return func(uploader *cloudUploader) {
   171  		uploader.reqEnhancers = append(uploader.reqEnhancers, enhancer)
   172  	}
   173  }