feat: Parallel/concurrent collection

cmd/collect_github_repos.go

@@ -2,8 +2,14 @@ package cmd
 
 import (
 	"fmt"
+	"os"
 
+	"github.com/Snider/Borg/pkg/compress"
 	"github.com/Snider/Borg/pkg/github"
+	"github.com/Snider/Borg/pkg/tim"
+	"github.com/Snider/Borg/pkg/trix"
+	"github.com/Snider/Borg/pkg/ui"
+	"github.com/schollz/progressbar/v3"
 	"github.com/spf13/cobra"
 )
 
@@ -17,17 +23,80 @@ var collectGithubReposCmd = &cobra.Command{
 	Short: "Collects all public repositories for a user or organization",
 	Args:  cobra.ExactArgs(1),
 	RunE: func(cmd *cobra.Command, args []string) error {
+		parallel, _ := cmd.Flags().GetInt("parallel")
+		outputFile, _ := cmd.Flags().GetString("output")
+		format, _ := cmd.Flags().GetString("format")
+		compression, _ := cmd.Flags().GetString("compression")
+		password, _ := cmd.Flags().GetString("password")
+
 		repos, err := GithubClient.GetPublicRepos(cmd.Context(), args[0])
 		if err != nil {
 			return err
 		}
-		for _, repo := range repos {
-			fmt.Fprintln(cmd.OutOrStdout(), repo)
+
+		prompter := ui.NewNonInteractivePrompter(ui.GetVCSQuote)
+		prompter.Start()
+		defer prompter.Stop()
+		var bar *progressbar.ProgressBar
+		if prompter.IsInteractive() {
+			bar = ui.NewProgressBar(len(repos), "Cloning repositories")
 		}
+
+		downloader := github.NewDownloader(parallel, bar)
+		dn, err := downloader.DownloadRepositories(cmd.Context(), repos)
+		if err != nil {
+			return err
+		}
+
+		var data []byte
+		if format == "tim" {
+			tim, err := tim.FromDataNode(dn)
+			if err != nil {
+				return fmt.Errorf("error creating tim: %w", err)
+			}
+			data, err = tim.ToTar()
+			if err != nil {
+				return fmt.Errorf("error serializing tim: %w", err)
+			}
+		} else if format == "trix" {
+			data, err = trix.ToTrix(dn, password)
+			if err != nil {
+				return fmt.Errorf("error serializing trix: %w", err)
+			}
+		} else {
+			data, err = dn.ToTar()
+			if err != nil {
+				return fmt.Errorf("error serializing DataNode: %w", err)
+			}
+		}
+
+		compressedData, err := compress.Compress(data, compression)
+		if err != nil {
+			return fmt.Errorf("error compressing data: %w", err)
+		}
+
+		if outputFile == "" {
+			outputFile = args[0] + "." + format
+			if compression != "none" {
+				outputFile += "." + compression
+			}
+		}
+
+		err = os.WriteFile(outputFile, compressedData, 0644)
+		if err != nil {
+			return fmt.Errorf("error writing repos to file: %w", err)
+		}
+
+		fmt.Fprintln(cmd.OutOrStdout(), "Repositories saved to", outputFile)
 		return nil
 	},
 }
 
 func init() {
 	collectGithubCmd.AddCommand(collectGithubReposCmd)
+	collectGithubReposCmd.PersistentFlags().Int("parallel", 1, "Number of concurrent workers")
+	collectGithubReposCmd.PersistentFlags().String("output", "", "Output file for the DataNode")
+	collectGithubReposCmd.PersistentFlags().String("format", "datanode", "Output format (datanode, tim, or trix)")
+	collectGithubReposCmd.PersistentFlags().String("compression", "none", "Compression format (none, gz, or xz)")
+	collectGithubReposCmd.PersistentFlags().String("password", "", "Password for encryption")
 }

cmd/collect_website.go

@@ -35,6 +35,8 @@ func NewCollectWebsiteCmd() *cobra.Command {
 			websiteURL := args[0]
 			outputFile, _ := cmd.Flags().GetString("output")
 			depth, _ := cmd.Flags().GetInt("depth")
+			parallel, _ := cmd.Flags().GetInt("parallel")
+			rateLimit, _ := cmd.Flags().GetFloat64("rate-limit")
 			format, _ := cmd.Flags().GetString("format")
 			compression, _ := cmd.Flags().GetString("compression")
 			password, _ := cmd.Flags().GetString("password")
@@ -51,7 +53,7 @@ func NewCollectWebsiteCmd() *cobra.Command {
 				bar = ui.NewProgressBar(-1, "Crawling website")
 			}
 
-			dn, err := website.DownloadAndPackageWebsite(websiteURL, depth, bar)
+			dn, err := website.DownloadAndPackageWebsite(cmd.Context(), websiteURL, depth, parallel, rateLimit, bar)
 			if err != nil {
 				return fmt.Errorf("error downloading and packaging website: %w", err)
 			}
@@ -101,6 +103,8 @@ func NewCollectWebsiteCmd() *cobra.Command {
 	}
 	collectWebsiteCmd.PersistentFlags().String("output", "", "Output file for the DataNode")
 	collectWebsiteCmd.PersistentFlags().Int("depth", 2, "Recursion depth for downloading")
+	collectWebsiteCmd.PersistentFlags().Int("parallel", 1, "Number of concurrent workers")
+	collectWebsiteCmd.PersistentFlags().Float64("rate-limit", 0, "Max requests per second per domain")
 	collectWebsiteCmd.PersistentFlags().String("format", "datanode", "Output format (datanode, tim, or trix)")
 	collectWebsiteCmd.PersistentFlags().String("compression", "none", "Compression format (none, gz, or xz)")
 	collectWebsiteCmd.PersistentFlags().String("password", "", "Password for encryption")

cmd/collect_website_test.go

@@ -11,10 +11,14 @@ import (
 	"github.com/schollz/progressbar/v3"
 )
 
+import (
+	"context"
+)
+
 func TestCollectWebsiteCmd_Good(t *testing.T) {
 	// Mock the website downloader
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
-	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
+	website.DownloadAndPackageWebsite = func(ctx context.Context, startURL string, maxDepth, parallel int, rateLimit float64, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
 		return datanode.New(), nil
 	}
 	defer func() {
@@ -35,7 +39,7 @@ func TestCollectWebsiteCmd_Good(t *testing.T) {
 func TestCollectWebsiteCmd_Bad(t *testing.T) {
 	// Mock the website downloader to return an error
 	oldDownloadAndPackageWebsite := website.DownloadAndPackageWebsite
-	website.DownloadAndPackageWebsite = func(startURL string, maxDepth int, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
+	website.DownloadAndPackageWebsite = func(ctx context.Context, startURL string, maxDepth, parallel int, rateLimit float64, bar *progressbar.ProgressBar) (*datanode.DataNode, error) {
 		return nil, fmt.Errorf("website error")
 	}
 	defer func() {

examples/collect_website/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"os"
 
@@ -11,7 +12,7 @@ func main() {
 	log.Println("Collecting website...")
 
 	// Download and package the website.
-	dn, err := website.DownloadAndPackageWebsite("https://example.com", 2, nil)
+	dn, err := website.DownloadAndPackageWebsite(context.Background(), "https://example.com", 2, 1, 0, nil)
 	if err != nil {
 		log.Fatalf("Failed to collect website: %v", err)
 	}

go.mod

@@ -64,5 +64,6 @@ require (
 	golang.org/x/sys v0.38.0 // indirect
 	golang.org/x/term v0.37.0 // indirect
 	golang.org/x/text v0.31.0 // indirect
+	golang.org/x/time v0.8.0 // indirect
 	gopkg.in/warnings.v0 v0.1.2 // indirect
 )

go.sum

@@ -192,6 +192,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.31.0 h1:aC8ghyu4JhP8VojJ2lEHBnochRno1sgL6nEi9WGFGMM=
 golang.org/x/text v0.31.0/go.mod h1:tKRAlv61yKIjGGHX/4tP1LTbc13YSec1pxVEWXzfoeM=
+golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
+golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=

pkg/github/downloader.go

@@ -0,0 +1,128 @@
+package github
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"io/fs"
+	"net/url"
+	"strings"
+	"sync"
+
+	"github.com/Snider/Borg/pkg/datanode"
+	"github.com/Snider/Borg/pkg/vcs"
+	"github.com/schollz/progressbar/v3"
+)
+
+// Downloader manages a pool of workers for cloning repositories.
+type Downloader struct {
+	parallel int
+	bar      *progressbar.ProgressBar
+	cloner   vcs.GitCloner
+}
+
+// NewDownloader creates a new Downloader.
+func NewDownloader(parallel int, bar *progressbar.ProgressBar) *Downloader {
+	return &Downloader{
+		parallel: parallel,
+		bar:      bar,
+		cloner:   vcs.NewGitCloner(),
+	}
+}
+
+// DownloadRepositories downloads a list of repositories in parallel.
+func (d *Downloader) DownloadRepositories(ctx context.Context, repos []string) (*datanode.DataNode, error) {
+	var wg sync.WaitGroup
+	repoChan := make(chan string, len(repos))
+	errChan := make(chan error, len(repos))
+	mergedDN := datanode.New()
+	var mu sync.Mutex
+
+	for i := 0; i < d.parallel; i++ {
+		wg.Add(1)
+		go d.worker(ctx, &wg, repoChan, mergedDN, &mu, errChan)
+	}
+
+	for _, repo := range repos {
+		select {
+		case repoChan <- repo:
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
+	close(repoChan)
+
+	wg.Wait()
+	close(errChan)
+
+	var errs []error
+	for err := range errChan {
+		errs = append(errs, err)
+	}
+	if len(errs) > 0 {
+		return nil, fmt.Errorf("errors cloning repositories: %v", errs)
+	}
+
+	return mergedDN, nil
+}
+
+func (d *Downloader) worker(ctx context.Context, wg *sync.WaitGroup, repoChan <-chan string, mergedDN *datanode.DataNode, mu *sync.Mutex, errChan chan<- error) {
+	defer wg.Done()
+	for repoURL := range repoChan {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
+		repoName, err := GetRepoNameFromURL(repoURL)
+		if err != nil {
+			errChan <- err
+			continue
+		}
+
+		dn, err := d.cloner.CloneGitRepository(repoURL, nil)
+		if err != nil {
+			errChan <- fmt.Errorf("error cloning %s: %w", repoURL, err)
+			continue
+		}
+
+		err = dn.Walk(".", func(path string, de fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			if !de.IsDir() {
+				file, err := dn.Open(path)
+				if err != nil {
+					return err
+				}
+				defer file.Close()
+				content, err := io.ReadAll(file)
+				if err != nil {
+					return err
+				}
+				mu.Lock()
+				mergedDN.AddData(fmt.Sprintf("%s/%s", repoName, path), content)
+				mu.Unlock()
+			}
+			return nil
+		})
+		if err != nil {
+			errChan <- err
+		}
+
+		if d.bar != nil {
+			d.bar.Add(1)
+		}
+	}
+}
+
+// GetRepoNameFromURL extracts the repository name from a Git URL.
+func GetRepoNameFromURL(repoURL string) (string, error) {
+	u, err := url.Parse(repoURL)
+	if err != nil {
+		return "", err
+	}
+	path := strings.TrimSuffix(u.Path, ".git")
+	return strings.TrimPrefix(path, "/"), nil
+}