feat: Add dump/restore cloud storage completions

cmd/dump/cmd.go

@@ -3,10 +3,12 @@ package dump
 import (
 	"fmt"
 	"log/slog"
+	"maps"
 	"net/url"
 	"os"
 	"path"
 	"path/filepath"
+	"slices"
 	"time"
 
 	"github.com/clevyr/kubedb/internal/actions/dump"
@@ -64,7 +66,7 @@ func New() *cobra.Command {
 	return cmd
 }
 
-func validArgs(cmd *cobra.Command, args []string, _ string) ([]string, cobra.ShellCompDirective) {
+func validArgs(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 	if len(args) != 0 {
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}
@@ -82,6 +84,29 @@ func validArgs(cmd *cobra.Command, args []string, _ string) ([]string, cobra.She
 	}
 
 	formats := db.Formats()
+
+	if storage.IsCloud(toComplete) {
+		u, err := url.Parse(toComplete)
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		switch {
+		case storage.IsS3(toComplete):
+			if u.Host == "" || u.Path == "" {
+				return storage.CompleteBucketsS3(u)
+			} else {
+				return storage.CompleteObjectsS3(u, slices.Collect(maps.Values(formats)), true)
+			}
+		case storage.IsGCS(toComplete):
+			if u.Host == "" || u.Path == "" {
+				return storage.CompleteBucketsGCS(u, "")
+			} else {
+				return storage.CompleteObjectsGCS(u, slices.Collect(maps.Values(formats)), true)
+			}
+		}
+	}
+
 	exts := make([]string, 0, len(formats))
 	for _, ext := range formats {
 		exts = append(exts, ext[1:])

cmd/restore/cmd.go

@@ -3,7 +3,10 @@ package restore
 import (
 	"errors"
 	"fmt"
+	"maps"
+	"net/url"
 	"os"
+	"slices"
 
 	"github.com/charmbracelet/huh"
 	"github.com/clevyr/kubedb/internal/actions/restore"
@@ -16,7 +19,6 @@ import (
 	"github.com/clevyr/kubedb/internal/util"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"golang.org/x/exp/maps"
 	"k8s.io/kubectl/pkg/util/term"
 )
 
@@ -66,7 +68,7 @@ func New() *cobra.Command {
 	return cmd
 }
 
-func validArgs(cmd *cobra.Command, args []string, _ string) ([]string, cobra.ShellCompDirective) {
+func validArgs(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 	if len(args) != 0 {
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}
@@ -87,6 +89,29 @@ func validArgs(cmd *cobra.Command, args []string, _ string) ([]string, cobra.She
 	}
 
 	formats := db.Formats()
+
+	if storage.IsCloud(toComplete) {
+		u, err := url.Parse(toComplete)
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		switch {
+		case storage.IsS3(toComplete):
+			if u.Host == "" || u.Path == "" {
+				return storage.CompleteBucketsS3(u)
+			} else {
+				return storage.CompleteObjectsS3(u, slices.Collect(maps.Values(formats)), false)
+			}
+		case storage.IsGCS(toComplete):
+			if u.Host == "" || u.Path == "" {
+				return storage.CompleteBucketsGCS(u, "")
+			} else {
+				return storage.CompleteObjectsGCS(u, slices.Collect(maps.Values(formats)), false)
+			}
+		}
+	}
+
 	exts := make([]string, 0, len(formats))
 	for _, ext := range formats {
 		exts = append(exts, ext[1:])
@@ -142,7 +167,7 @@ func preRun(cmd *cobra.Command, args []string) error {
 				ShowSize(true).
 				ShowPermissions(false).
 				Height(15).
-				AllowedTypes(maps.Values(db.Formats())).
+				AllowedTypes(slices.Collect(maps.Values(db.Formats()))).
 				Value(&action.Filename),
 		))
 

go.mod

@@ -21,8 +21,8 @@ require (
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
-	golang.org/x/exp v0.0.0-20240909161429-701f63a606c0
 	golang.org/x/sync v0.8.0
+	google.golang.org/api v0.187.0
 	k8s.io/api v0.31.1
 	k8s.io/apimachinery v0.31.1
 	k8s.io/client-go v0.31.1
@@ -128,13 +128,13 @@ require (
 	go.opentelemetry.io/otel/trace v1.26.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 // indirect
 	golang.org/x/net v0.26.0 // indirect
 	golang.org/x/oauth2 v0.21.0 // indirect
 	golang.org/x/sys v0.25.0 // indirect
 	golang.org/x/term v0.22.0 // indirect
 	golang.org/x/text v0.18.0 // indirect
 	golang.org/x/time v0.5.0 // indirect
-	google.golang.org/api v0.187.0 // indirect
 	google.golang.org/genproto v0.0.0-20240624140628-dc46fd24d27d // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240617180043-68d350f18fd4 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240624140628-dc46fd24d27d // indirect

internal/storage/gcs.go

@@ -2,10 +2,13 @@ package storage
 
 import (
 	"context"
+	"errors"
+	"iter"
 	"net/url"
 	"strings"
 
 	"cloud.google.com/go/storage"
+	"google.golang.org/api/iterator"
 )
 
 const GCSSchema = "gs://"
@@ -25,6 +28,62 @@ func IsGCSDir(path string) bool {
 	return !strings.Contains(trimmed, "/")
 }
 
+func ListBucketsGCS(ctx context.Context, projectID string) iter.Seq2[*storage.BucketAttrs, error] {
+	return func(yield func(*storage.BucketAttrs, error) bool) {
+		client, err := storage.NewClient(ctx)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		objects := client.Buckets(ctx, projectID)
+		for {
+			attrs, err := objects.Next()
+			if err != nil && errors.Is(err, iterator.Done) {
+				return
+			}
+			if !yield(attrs, err) {
+				return
+			}
+		}
+	}
+}
+
+func ListObjectsGCS(ctx context.Context, key string) iter.Seq2[*storage.ObjectAttrs, error] {
+	return func(yield func(*storage.ObjectAttrs, error) bool) {
+		client, err := storage.NewClient(ctx)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		u, err := url.Parse(key)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		u.Path = strings.TrimLeft(u.Path, "/")
+
+		query := &storage.Query{
+			Delimiter:                "/",
+			Prefix:                   u.Path,
+			Projection:               storage.ProjectionNoACL,
+			IncludeFoldersAsPrefixes: true,
+		}
+
+		objects := client.Bucket(u.Host).Objects(ctx, query)
+		for {
+			attrs, err := objects.Next()
+			if err != nil && errors.Is(err, iterator.Done) {
+				return
+			}
+			if !yield(attrs, err) {
+				return
+			}
+		}
+	}
+}
+
 func UploadGCS(ctx context.Context, key string) (*storage.Writer, error) {
 	client, err := storage.NewClient(ctx)
 	if err != nil {

internal/storage/gcs_completions.go

@@ -0,0 +1,61 @@
+package storage
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"os"
+
+	"github.com/clevyr/kubedb/internal/util"
+	"github.com/dustin/go-humanize"
+	"github.com/spf13/cobra"
+)
+
+func CompleteBucketsGCS(u *url.URL, projectID string) ([]string, cobra.ShellCompDirective) {
+	if projectID == "" {
+		if val := os.Getenv("GOOGLE_CLOUD_PROJECT"); val != "" {
+			projectID = val
+		} else if val := os.Getenv("GCLOUD_PROJECT"); val != "" {
+			projectID = val
+		} else if val := os.Getenv("GCP_PROJECT"); val != "" {
+			projectID = val
+		}
+	}
+
+	u.Path = "/"
+
+	var names []string //nolint:prealloc
+	for bucket, err := range ListBucketsGCS(context.Background(), projectID) {
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		u.Host = bucket.Name
+		names = append(names, u.String())
+	}
+	return names, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveNoSpace
+}
+
+func CompleteObjectsGCS(u *url.URL, exts []string, dirOnly bool) ([]string, cobra.ShellCompDirective) {
+	var paths []string
+	for object, err := range ListObjectsGCS(context.Background(), u.String()) {
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		if object.Prefix != "" {
+			u.Path = object.Prefix
+			paths = append(paths, u.String())
+		} else if !dirOnly && util.FilterExts(exts, object.Name) {
+			u.Path = object.Name
+			paths = append(paths,
+				fmt.Sprintf("%s\t%s; %s",
+					u.String(),
+					object.Updated.Local().Format("Jan _2 15:04"), //nolint:gosmopolitan
+					humanize.IBytes(uint64(object.Size)),          //nolint:gosec
+				),
+			)
+		}
+	}
+	return paths, cobra.ShellCompDirectiveNoFileComp | cobra.ShellCompDirectiveNoSpace
+}