Improved readability and maintainability

Author: marcleblanc2

cmd/src/snapshot_upload.go

@@ -6,8 +6,8 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"path"
 	"path/filepath"
+	"slices"
 	"strings"
 
 	"cloud.google.com/go/storage"
@@ -19,9 +19,20 @@ import (
 	"github.com/sourcegraph/src-cli/internal/pgdump"
 )
 
+// Package-level variables
 const srcSnapshotDir = "./src-snapshot"
 
-var srcSnapshotSummaryPath = path.Join(srcSnapshotDir, "summary.json")
+// summaryFile on its own, as it gets handled a little differently
+const summaryFile = "summary.json"
+var srcSnapshotSummaryPath = filepath.Join(srcSnapshotDir, summaryFile)
+
+// listOfValidFiles defines the valid snapshot filenames (with extensions) that can be uploaded
+var listOfValidFiles = []string{
+	"codeinsights.sql",
+	"codeintel.sql",
+	"pgsql.sql",
+	summaryFile,
+}
 
 // Define types
 type uploadArgs struct {
@@ -31,13 +42,14 @@ type uploadArgs struct {
 	filesToUpload   []string
 }
 
+// Google Cloud Storage upload client
 type gcsClient struct {
 	ctx           context.Context
 	out           *output.Output
 	storageClient *storage.Client
 }
 
-// uploadFile represents a file registered for upload
+// uploadFile represents a file opened for upload
 type uploadFile struct {
 	file      *os.File
 	stat      os.FileInfo
@@ -46,7 +58,7 @@ type uploadFile struct {
 
 
 func init() {
-	usage := `'src snapshot upload' uploads instance snapshot contents generated by 'src snapshot databases' and 'src snapshot summary' to the designated bucket.
+	usage := fmt.Sprintf(`'src snapshot upload' uploads instance snapshot contents generated by 'src snapshot databases' and 'src snapshot summary' to the designated bucket.
 
 USAGE
 	src snapshot upload -bucket=$MIGRATION_BUCKET_NAME -credentials=./migration_private_key.json [-file=$FILE]
@@ -55,15 +67,15 @@ BUCKET
 	In general, a Google Cloud Storage bucket and relevant credentials will be provided by Sourcegraph when using this functionality to share a snapshot with Sourcegraph.
 
 FILE
-	Optional: Specify files to upload as a comma-delimited list. Valid values: summary, primary, codeintel, codeinsights. Default: All valid values.
-	Examples: -file=summary,primary or -file=codeintel
-`
+	Optional: Specify files to upload as a comma-delimited list (with file extensions). Valid values: %s. Default: All files.
+	Examples: -file=summary.json,pgsql.sql or -file=codeintel.sql
+`, strings.Join(listOfValidFiles, ", "))
 
 	flagSet := flag.NewFlagSet("upload", flag.ExitOnError)
 	bucketName := flagSet.String("bucket", "", "destination Cloud Storage bucket name")
 	credentialsPath := flagSet.String("credentials", "", "JSON credentials file for Google Cloud service account")
+	fileFilter := flagSet.String("file", strings.Join(listOfValidFiles, ","), "comma-delimited list of files to upload")
 	filterSQL := flagSet.Bool("filter-sql", true, "filter incompatible SQL statements from database dumps for import to Google Cloud SQL")
-	fileFilter := flagSet.String("file", "summary,primary,codeintel,codeinsights", "comma-delimited list of files to upload")
 
 	// Register this command with the parent 'src snapshot' command.
 	// The parent snapshot.go command runs all registered subcommands via snapshotCommands.run().
@@ -76,14 +88,14 @@ FILE
 	})
 }
 
-// Helper function to keep init() succinct
+// Handler function to keep init() succinct
 func snapshotUploadHandler(flagSet *flag.FlagSet, bucketName, credentialsPath *string, filterSQL *bool, fileFilter *string) func([]string) error {
 	return func(args []string) error {
 		if err := flagSet.Parse(args); err != nil {
 			return err
 		}
 
-		// Validate and parse inputs
+		// Validate and parse inputs into an uploadArgs-type object
 		uploadArgs, err := validateUploadInputs(*bucketName, *credentialsPath, *fileFilter, *filterSQL)
 		if err != nil {
 			return err
@@ -131,43 +143,41 @@ func validateUploadInputs(bucketName, credentialsPath, fileFilter string, filter
 // Parse the --file arg values, and return a list of strings of the files to upload
 func parseFileFilter(fileFilter string) ([]string, error) {
 
-	validFiles := map[string]bool{
-		"summary":      true,
-		"primary":      true,
-		"codeintel":    true,
-		"codeinsights": true,
-	}
-
+	// Default: all files
 	if fileFilter == "" {
-		// Default: all files
-		return []string{"summary", "primary", "codeintel", "codeinsights"}, nil
+		return listOfValidFiles, nil
 	}
 
-	// Parse comma-delimited list
 	var filesToUpload []string
-	parts := strings.Split(fileFilter, ",")
-	for _, part := range parts {
 
-		// Normalize: trim spaces and strip file extensions
-		normalized := strings.TrimSpace(part)
-		normalized = strings.TrimSuffix(normalized, ".json")
-		normalized = strings.TrimSuffix(normalized, ".sql")
+	// Parse comma-delimited list
+	for _, part := range strings.Split(fileFilter, ",") {
 
-		// Validate
-		if !validFiles[normalized] {
-			return nil, errors.Newf("invalid -file value %q. Valid values: summary[.json], primary[.sql], codeintel[.sql], codeinsights[.sql]", part)
+		// Trim whitespace
+		filename := strings.TrimSpace(part)
+
+		// Validate against list of valid files
+		if !slices.Contains(listOfValidFiles, filename) {
+			return nil, errors.Newf("invalid -file value %q. Valid values: %s", part, strings.Join(listOfValidFiles, ", "))
 		}
 
-		filesToUpload = append(filesToUpload, normalized)
+		filesToUpload = append(filesToUpload, filename)
 	}
 
+	// Sort files alphabetically for consistent ordering
+	slices.Sort(filesToUpload)
+
+	// Remove duplicates (works on sorted slices by removing adjacent duplicates)
+	filesToUpload = slices.Compact(filesToUpload)
+
 	return filesToUpload, nil
 }
 
 func createGcsClient(flagSet *flag.FlagSet, credentialsPath string) (*gcsClient, error) {
 
-	out := output.NewOutput(flagSet.Output(), output.OutputOpts{Verbose: *verbose})
 	ctx := context.Background()
+	out := output.NewOutput(flagSet.Output(), output.OutputOpts{Verbose: *verbose})
+
 	// https://pkg.go.dev/cloud.google.com/go/storage#section-readme
 	client, err := storage.NewClient(ctx, option.WithCredentialsFile(credentialsPath))
 
@@ -186,22 +196,25 @@ func createGcsClient(flagSet *flag.FlagSet, credentialsPath string) (*gcsClient,
 // Returns arrays of uploadFile and progress bars (aligned by index).
 func openFilesAndCreateProgressBars(args *uploadArgs) ([]uploadFile, []output.ProgressBar, error) {
 	var (
-		openedFiles  []uploadFile          // Files opened from disk, ready for upload (aligned with progressBars)
+		openedFiles  []uploadFile         // Files opened from disk, ready for upload (aligned with progressBars)
 		progressBars []output.ProgressBar // Progress bars for UI (aligned with openedFiles)
 	)
 
 	// addFile opens a file from disk and registers it for upload.
 	// It adds the file to the openedFiles array and creates a corresponding progress bar.
 	// For database dumps (!isSummary), SQL filtering is enabled based on args.filterSQL.
-	addFile := func(filePath string, isSummary bool) error {
+	addFile := func(filePath string) error {
+
+		isSummary := strings.HasSuffix(filePath, summaryFile)
 
 		// Open the file
 		openFile, err := os.Open(filePath)
+
 		if err != nil {
 			if isSummary {
-				return errors.Wrap(err, "failed to open snapshot summary - generate one with 'src snapshot summary'")
+				return errors.Wrap(err, fmt.Sprintf("failed to open snapshot summary %s - Please generate it with 'src snapshot summary'", filePath))
 			}
-			return errors.Wrap(err, "failed to open database dump - generate one with 'src snapshot databases'")
+			return errors.Wrap(err, fmt.Sprintf("failed to open database dump %s - Please generate them with 'src snapshot databases'", filePath))
 		}
 
 		// Get file metadata (name, size)
@@ -228,17 +241,12 @@ func openFilesAndCreateProgressBars(args *uploadArgs) ([]uploadFile, []output.Pr
 	// Open files based on user's selection (via --file arg)
 	// Iterate through the user's selected files and open each one
 	for _, selectedFile := range args.filesToUpload {
-		if selectedFile == "summary" {
-			// Open summary.json
-			if err := addFile(srcSnapshotSummaryPath, true); err != nil {
-				return nil, nil, err
-			}
-		} else {
-			// Open database dump file (primary.sql, codeintel.sql, or codeinsights.sql)
-			dbFilePath := filepath.Join(srcSnapshotDir, selectedFile+".sql")
-			if err := addFile(dbFilePath, false); err != nil {
-				return nil, nil, err
-			}
+
+		// Construct full file path
+		filePath := filepath.Join(srcSnapshotDir, selectedFile)
+
+		if err := addFile(filePath); err != nil {
+			return nil, nil, err
 		}
 	}
 
@@ -248,6 +256,7 @@ func openFilesAndCreateProgressBars(args *uploadArgs) ([]uploadFile, []output.Pr
 // uploadFilesToBucket uploads the prepared files to Google Cloud Storage bucket.
 // Uploads are performed in parallel with progress tracking.
 func uploadFilesToBucket(client *gcsClient, args *uploadArgs, openedFiles []uploadFile, progressBars []output.ProgressBar) error {
+
 	// Start uploads with progress tracking
 	progress := client.out.Progress(progressBars, nil)
 	progress.WriteLine(output.Emoji(output.EmojiHourglass, "Starting uploads..."))
@@ -256,8 +265,10 @@ func uploadFilesToBucket(client *gcsClient, args *uploadArgs, openedFiles []uplo
 
 	// Upload each file in parallel
 	for fileIndex, openedFile := range openedFiles {
+
 		fileIndex := fileIndex
 		openedFile := openedFile
+
 		uploadPool.Go(func(ctx context.Context) error {
 			progressFn := func(bytesWritten int64) { progress.SetValue(fileIndex, float64(bytesWritten)) }
 
@@ -273,11 +284,11 @@ func uploadFilesToBucket(client *gcsClient, args *uploadArgs, openedFiles []uplo
 	errs := uploadPool.Wait()
 	progress.Complete()
 	if errs != nil {
-		client.out.WriteLine(output.Line(output.EmojiFailure, output.StyleFailure, "Some snapshot contents failed to upload."))
+		client.out.WriteLine(output.Line(output.EmojiFailure, output.StyleFailure, "Some file(s) failed to upload."))
 		return errs
 	}
 
-	client.out.WriteLine(output.Emoji(output.EmojiSuccess, "Summary contents uploaded!"))
+	client.out.WriteLine(output.Emoji(output.EmojiSuccess, "File(s) uploaded successfully!"))
 	return nil
 }
 
@@ -291,7 +302,7 @@ func streamFileToBucket(ctx context.Context, file *uploadFile, bucket *storage.B
 	// To assert against actual file size
 	var totalBytesWritten int64
 
-	// Do a partial copy, that filters out incompatible statements
+	// Start a partial copy, that filters out incompatible statements
 	if file.filterSQL {
 		bytesWritten, err := pgdump.FilterInvalidLines(writer, file.file, progressFn)
 		if err != nil {