github.com/emcfarlane/larking@v0.0.0-20220605172417-1704b45ee6c3/starlib/starlarkrule/archive.go (about)

     1  package starlarkrule
     2  
     3  import (
     4  	"archive/tar"
     5  	"compress/gzip"
     6  	"fmt"
     7  	"io"
     8  	"path"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/emcfarlane/larking/starlib/starext"
    13  	"github.com/emcfarlane/larking/starlib/starlarkthread"
    14  	"go.starlark.net/starlark"
    15  	"go.starlark.net/starlarkstruct"
    16  	"gocloud.dev/blob"
    17  )
    18  
    19  var archiveModule = &starlarkstruct.Module{
    20  	Name: "packaging",
    21  	Members: starlark.StringDict{
    22  		"tar": starext.MakeBuiltin("packaging.tar", makeTar),
    23  	},
    24  }
    25  
    26  // makeTar creates a tarball.
    27  func makeTar(thread *starlark.Thread, fnname string, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    28  	ctx := starlarkthread.GetContext(thread)
    29  
    30  	var (
    31  		name        string
    32  		srcs        *starlark.List
    33  		packageDir  string
    34  		stripPrefix string
    35  	)
    36  	if err := starlark.UnpackArgs(
    37  		fnname, args, kwargs,
    38  		"name", &name,
    39  		"srcs", &srcs,
    40  		"package_dir?", &packageDir,
    41  		"strip_prefix?", &stripPrefix,
    42  	); err != nil {
    43  		return nil, err
    44  	}
    45  
    46  	var (
    47  		creationTime time.Time // zero
    48  		blobs        starext.Blobs
    49  	)
    50  	defer blobs.Close()
    51  
    52  	createTar := func(l *Label) error {
    53  		fmt.Println("createTar", l)
    54  		bktURL := l.BucketURL()
    55  		key := l.Key()
    56  
    57  		w, err := blobs.NewWriter(ctx, bktURL, key, nil)
    58  		if err != nil {
    59  			return err
    60  		}
    61  		defer w.Close()
    62  
    63  		// compress writer
    64  		var cw io.WriteCloser
    65  		switch {
    66  		case strings.HasSuffix(key, ".tar.gz"):
    67  			cw = gzip.NewWriter(w)
    68  			defer cw.Close()
    69  		default:
    70  			cw = w
    71  		}
    72  
    73  		tw := tar.NewWriter(cw)
    74  		defer tw.Close()
    75  
    76  		addFile := func(l *Label, key string) error {
    77  			bktURL := l.BucketURL()
    78  			bkt, err := blob.OpenBucket(ctx, bktURL)
    79  			if err != nil {
    80  				return err
    81  			}
    82  
    83  			r, err := blobs.NewReader(ctx, bktURL, l.Key(), nil)
    84  			if err != nil {
    85  				return err
    86  			}
    87  			defer r.Close()
    88  
    89  			attrs, err := bkt.Attributes(ctx, l.Key())
    90  			if err != nil {
    91  				return err
    92  			}
    93  
    94  			header := &tar.Header{
    95  				Name:     key,
    96  				Size:     attrs.Size,
    97  				Typeflag: tar.TypeReg,
    98  				Mode:     0600,
    99  				ModTime:  creationTime,
   100  			}
   101  			// write the header to the tarball archive
   102  			if err := tw.WriteHeader(header); err != nil {
   103  				return err
   104  			}
   105  			// copy the file data to the tarball
   106  			if _, err := io.Copy(tw, r); err != nil {
   107  				return err
   108  			}
   109  			return nil
   110  		}
   111  
   112  		for i, n := 0, srcs.Len(); i < n; i++ {
   113  			src := srcs.Index(i)
   114  			l, ok := src.(*Label)
   115  			if !ok {
   116  				return fmt.Errorf("invalid src type: %T", src.Type())
   117  			}
   118  
   119  			// Form the key path of the file in the tar fs.
   120  			key := path.Join(
   121  				packageDir,
   122  				strings.TrimPrefix(l.Key(), stripPrefix),
   123  			)
   124  
   125  			// TODO: strip_prefix
   126  			if err := addFile(l, key); err != nil {
   127  				return err
   128  			}
   129  		}
   130  		return nil
   131  	}
   132  	l, err := ParseRelativeLabel(thread.Name, name)
   133  	if err != nil {
   134  		return nil, err
   135  	}
   136  
   137  	if err := createTar(l); err != nil {
   138  		return nil, err
   139  	}
   140  	fmt.Println("makeTar", l)
   141  	return l, nil
   142  }