From 08f25a1a89c30854d9310ce0182aee103713429d Mon Sep 17 00:00:00 2001 From: Gabe Cook Date: Fri, 1 Mar 2024 17:11:40 -0600 Subject: [PATCH] feat(dump): Add support for dumps to an S3 bucket --- cmd/dump/dump.go | 44 ++++++++++++++++++++++++---- docs/kubedb_dump.md | 19 +++++++++--- go.mod | 18 ++++++++++++ go.sum | 36 +++++++++++++++++++++++ internal/actions/dump/dump.go | 49 +++++++++++++++---------------- internal/s3/s3.go | 55 +++++++++++++++++++++++++++++++++++ internal/s3/s3_test.go | 50 +++++++++++++++++++++++++++++++ 7 files changed, 236 insertions(+), 35 deletions(-) create mode 100644 internal/s3/s3.go create mode 100644 internal/s3/s3_test.go diff --git a/cmd/dump/dump.go b/cmd/dump/dump.go index bab90500..48bfad85 100644 --- a/cmd/dump/dump.go +++ b/cmd/dump/dump.go @@ -1,11 +1,17 @@ package dump import ( + "net/url" "os" + "path" + "path/filepath" + "strings" + "time" "github.com/clevyr/kubedb/internal/actions/dump" "github.com/clevyr/kubedb/internal/config/flags" "github.com/clevyr/kubedb/internal/consts" + "github.com/clevyr/kubedb/internal/s3" "github.com/clevyr/kubedb/internal/util" "github.com/spf13/cobra" "github.com/spf13/viper" @@ -16,14 +22,25 @@ var action dump.Dump func New() *cobra.Command { cmd := &cobra.Command{ - Use: "dump [filename]", + Use: "dump [filename | S3 URI]", Aliases: []string{"d", "export"}, Short: "Dump a database to a sql file", Long: `Dump a database to a sql file. -If no filename is provided, the filename will be generated. -For example, if a dump is performed in the namespace "clevyr" with no extra flags, -the generated filename might look like "` + dump.HelpFilename() + `"`, +Filenames: + If a filename is provided, and it does not end with a "/", then it will be used verbatim. + Otherwise, the filename will be generated and appended to the given path. + For example, if a dump is performed in the namespace "clevyr" with no extra flags, + the generated filename might look like "` + dump.HelpFilename() + `". + Similarly, if the filename is passed as "backups/", then the generated path might look like + "backups/` + dump.HelpFilename() + `". + +S3: + If the filename begins with "s3://", then the dump will be directly uploaded to an S3 bucket. + S3 configuration will be loaded from the environment or from the current aws cli profile. + Note the above section on filenames. For example, if the filename is set to "s3://backups", + then the resulting filename might look like "s3://backups/` + dump.HelpFilename() + `". +`, Args: cobra.MaximumNArgs(1), ValidArgsFunction: validArgs, @@ -102,7 +119,24 @@ func preRun(cmd *cobra.Command, args []string) (err error) { return err } - if action.Filename != "" && !cmd.Flags().Lookup(consts.FormatFlag).Changed { + if action.Filename == "" || strings.HasSuffix(action.Filename, string(os.PathSeparator)) || s3.IsS3Dir(action.Filename) { + generated := dump.Filename{ + Database: action.Database, + Namespace: action.Client.Namespace, + Ext: action.Dialect.DumpExtension(action.Format), + Date: time.Now(), + }.Generate() + if s3.IsS3(action.Filename) { + u, err := url.Parse(action.Filename) + if err != nil { + return err + } + u.Path = path.Join(u.Path, generated) + action.Filename = u.String() + } else { + action.Filename = filepath.Join(action.Filename, generated) + } + } else if !cmd.Flags().Lookup(consts.FormatFlag).Changed { action.Format = action.Dialect.FormatFromFilename(action.Filename) } diff --git a/docs/kubedb_dump.md b/docs/kubedb_dump.md index 7534c81a..ba88a01e 100644 --- a/docs/kubedb_dump.md +++ b/docs/kubedb_dump.md @@ -6,12 +6,23 @@ Dump a database to a sql file Dump a database to a sql file. -If no filename is provided, the filename will be generated. -For example, if a dump is performed in the namespace "clevyr" with no extra flags, -the generated filename might look like "clevyr_2022-01-09_094100.sql.gz" +Filenames: + If a filename is provided, and it does not end with a "/", then it will be used verbatim. + Otherwise, the filename will be generated and appended to the given path. + For example, if a dump is performed in the namespace "clevyr" with no extra flags, + the generated filename might look like "clevyr_2022-01-09_094100.sql.gz". + Similarly, if the filename is passed as "backups/", then the generated path might look like + "backups/clevyr_2022-01-09_094100.sql.gz". + +S3: + If the filename begins with "s3://", then the dump will be directly uploaded to an S3 bucket. + S3 configuration will be loaded from the environment or from the current aws cli profile. + Note the above section on filenames. For example, if the filename is set to "s3://backups", + then the resulting filename might look like "s3://backups/clevyr_2022-01-09_094100.sql.gz". + ``` -kubedb dump [filename] [flags] +kubedb dump [filename | S3 URI] [flags] ``` ### Options diff --git a/go.mod b/go.mod index 22303a7c..cdd47875 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,9 @@ go 1.22.0 require ( github.com/AlecAivazis/survey/v2 v2.3.7 github.com/alessio/shellescape v1.4.2 + github.com/aws/aws-sdk-go-v2 v1.25.2 + github.com/aws/aws-sdk-go-v2/config v1.27.4 + github.com/aws/aws-sdk-go-v2/service/s3 v1.51.1 github.com/fatih/color v1.16.0 github.com/gabe565/go-spinners v1.0.1 github.com/jedib0t/go-pretty/v6 v6.5.4 @@ -26,6 +29,21 @@ require ( require ( github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.17.4 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.2 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.2 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.2 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.2 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.2 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.20.1 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.1 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.28.1 // indirect + github.com/aws/smithy-go v1.20.1 // indirect github.com/cpuguy83/go-md2man/v2 v2.0.3 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect github.com/emicklei/go-restful/v3 v3.11.0 // indirect diff --git a/go.sum b/go.sum index eb7af810..ee11568d 100644 --- a/go.sum +++ b/go.sum @@ -8,6 +8,42 @@ github.com/alessio/shellescape v1.4.2 h1:MHPfaU+ddJ0/bYWpgIeUnQUqKrlJ1S7BfEYPM4u github.com/alessio/shellescape v1.4.2/go.mod h1:PZAiSCk0LJaZkiCSkPv8qIobYglO3FPpyFjDCtHLS30= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= +github.com/aws/aws-sdk-go-v2 v1.25.2 h1:/uiG1avJRgLGiQM9X3qJM8+Qa6KRGK5rRPuXE0HUM+w= +github.com/aws/aws-sdk-go-v2 v1.25.2/go.mod h1:Evoc5AsmtveRt1komDwIsjHFyrP5tDuF1D1U+6z6pNo= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 h1:gTK2uhtAPtFcdRRJilZPx8uJLL2J85xK11nKtWL0wfU= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1/go.mod h1:sxpLb+nZk7tIfCWChfd+h4QwHNUR57d8hA1cleTkjJo= +github.com/aws/aws-sdk-go-v2/config v1.27.4 h1:AhfWb5ZwimdsYTgP7Od8E9L1u4sKmDW2ZVeLcf2O42M= +github.com/aws/aws-sdk-go-v2/config v1.27.4/go.mod h1:zq2FFXK3A416kiukwpsd+rD4ny6JC7QSkp4QdN1Mp2g= +github.com/aws/aws-sdk-go-v2/credentials v1.17.4 h1:h5Vztbd8qLppiPwX+y0Q6WiwMZgpd9keKe2EAENgAuI= +github.com/aws/aws-sdk-go-v2/credentials v1.17.4/go.mod h1:+30tpwrkOgvkJL1rUZuRLoxcJwtI/OkeBLYnHxJtVe0= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.2 h1:AK0J8iYBFeUk2Ax7O8YpLtFsfhdOByh2QIkHmigpRYk= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.2/go.mod h1:iRlGzMix0SExQEviAyptRWRGdYNo3+ufW/lCzvKVTUc= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.2 h1:bNo4LagzUKbjdxE0tIcR9pMzLR2U/Tgie1Hq1HQ3iH8= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.2/go.mod h1:wRQv0nN6v9wDXuWThpovGQjqF1HFdcgWjporw14lS8k= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.2 h1:EtOU5jsPdIQNP+6Q2C5e3d65NKT1PeCiQk+9OdzO12Q= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.2/go.mod h1:tyF5sKccmDz0Bv4NrstEr+/9YkSPJHrcO7UsUKf7pWM= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.2 h1:en92G0Z7xlksoOylkUhuBSfJgijC7rHVLRdnIlHEs0E= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.2/go.mod h1:HgtQ/wN5G+8QSlK62lbOtNwQ3wTSByJ4wH2rCkPt+AE= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 h1:EyBZibRTVAs6ECHZOw5/wlylS9OcTzwyjeQMudmREjE= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1/go.mod h1:JKpmtYhhPs7D97NL/ltqz7yCkERFW5dOlHyVl66ZYF8= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.2 h1:zSdTXYLwuXDNPUS+V41i1SFDXG7V0ITp0D9UT9Cvl18= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.2/go.mod h1:v8m8k+qVy95nYi7d56uP1QImleIIY25BPiNJYzPBdFE= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.2 h1:5ffmXjPtwRExp1zc7gENLgCPyHFbhEPwVTkTiH9niSk= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.2/go.mod h1:Ru7vg1iQ7cR4i7SZ/JTLYN9kaXtbL69UdgG0OQWQxW0= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.2 h1:1oY1AVEisRI4HNuFoLdRUB0hC63ylDAN6Me3MrfclEg= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.2/go.mod h1:KZ03VgvZwSjkT7fOetQ/wF3MZUvYFirlI1H5NklUNsY= +github.com/aws/aws-sdk-go-v2/service/s3 v1.51.1 h1:juZ+uGargZOrQGNxkVHr9HHR/0N+Yu8uekQnV7EAVRs= +github.com/aws/aws-sdk-go-v2/service/s3 v1.51.1/go.mod h1:SoR0c7Jnq8Tpmt0KSLXIavhjmaagRqQpe9r70W3POJg= +github.com/aws/aws-sdk-go-v2/service/sso v1.20.1 h1:utEGkfdQ4L6YW/ietH7111ZYglLJvS+sLriHJ1NBJEQ= +github.com/aws/aws-sdk-go-v2/service/sso v1.20.1/go.mod h1:RsYqzYr2F2oPDdpy+PdhephuZxTfjHQe7SOBcZGoAU8= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.1 h1:9/GylMS45hGGFCcMrUZDVayQE1jYSIN6da9jo7RAYIw= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.1/go.mod h1:YjAPFn4kGFqKC54VsHs5fn5B6d+PCY2tziEa3U/GB5Y= +github.com/aws/aws-sdk-go-v2/service/sts v1.28.1 h1:3I2cBEYgKhrWlwyZgfpSO2BpaMY1LHPqXYk/QGlu2ew= +github.com/aws/aws-sdk-go-v2/service/sts v1.28.1/go.mod h1:uQ7YYKZt3adCRrdCBREm1CD3efFLOUNH77MrUCvx5oA= +github.com/aws/smithy-go v1.20.1 h1:4SZlSlMr36UEqC7XOyRVb27XMeZubNcBNN+9IgEPIQw= +github.com/aws/smithy-go v1.20.1/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E= github.com/cpuguy83/go-md2man/v2 v2.0.3 h1:qMCsGGgs+MAzDFyp9LpAe1Lqy/fY/qCovCm0qnXZOBM= github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= diff --git a/internal/actions/dump/dump.go b/internal/actions/dump/dump.go index 8ba6de39..13d77b4d 100644 --- a/internal/actions/dump/dump.go +++ b/internal/actions/dump/dump.go @@ -14,6 +14,7 @@ import ( "github.com/clevyr/kubedb/internal/github" "github.com/clevyr/kubedb/internal/kubernetes" "github.com/clevyr/kubedb/internal/progressbar" + "github.com/clevyr/kubedb/internal/s3" gzip "github.com/klauspost/pgzip" log "github.com/sirupsen/logrus" "golang.org/x/sync/errgroup" @@ -24,22 +25,22 @@ type Dump struct { } func (action Dump) Run(ctx context.Context) (err error) { - if action.Filename == "" { - action.Filename = Filename{ - Database: action.Database, - Namespace: action.Client.Namespace, - Ext: action.Dialect.DumpExtension(action.Format), - Date: time.Now(), - }.Generate() - if err != nil { - return err - } - } + errGroup, ctx := errgroup.WithContext(ctx) var f io.WriteCloser - switch action.Filename { - case "-": + switch { + case action.Filename == "-": f = os.Stdout + case s3.IsS3(action.Filename): + pr, pw := io.Pipe() + f = pw + defer func(pw *io.PipeWriter) { + _ = pw.Close() + }(pw) + + errGroup.Go(func() error { + return s3.CreateUpload(ctx, pr, action.Filename) + }) default: if _, err := os.Stat(filepath.Dir(action.Filename)); os.IsNotExist(err) { err = os.MkdirAll(filepath.Dir(action.Filename), os.ModePerm) @@ -72,8 +73,6 @@ func (action Dump) Run(ctx context.Context) (err error) { bar, plogger := progressbar.New(os.Stderr, -1, "downloading", action.Spinner) defer bar.Close() - errGroup, ctx := errgroup.WithContext(ctx) - pr, pw := io.Pipe() // Begin database export errGroup.Go(func() error { @@ -130,8 +129,15 @@ func (action Dump) Run(ctx context.Context) (err error) { } } - _, err := io.Copy(io.MultiWriter(f, bar), r) - return err + if _, err := io.Copy(io.MultiWriter(f, bar), r); err != nil { + return err + } + + if err := f.Close(); err != nil && !errors.Is(err, os.ErrClosed) { + return err + } + + return nil }) if err := errGroup.Wait(); err != nil { @@ -140,15 +146,6 @@ func (action Dump) Run(ctx context.Context) (err error) { _ = bar.Finish() - // Close file - err = f.Close() - if err != nil { - // Ignore file already closed errors since w can be the same as f - if !errors.Is(err, os.ErrClosed) { - return err - } - } - log.WithFields(log.Fields{ "file": action.Filename, "in": time.Since(startTime).Truncate(10 * time.Millisecond), diff --git a/internal/s3/s3.go b/internal/s3/s3.go new file mode 100644 index 00000000..a525f408 --- /dev/null +++ b/internal/s3/s3.go @@ -0,0 +1,55 @@ +package s3 + +import ( + "context" + "io" + "net/url" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/s3" + "k8s.io/utils/ptr" +) + +const Schema = "s3://" + +func IsS3(path string) bool { + return strings.HasPrefix(path, Schema) +} + +func IsS3Dir(path string) bool { + if !IsS3(path) { + return false + } + if strings.HasSuffix(path, "/") { + return true + } + trimmed := strings.TrimPrefix(path, Schema) + return !strings.Contains(trimmed, "/") +} + +func CreateUpload(ctx context.Context, r io.ReadCloser, key string) error { + defer func(r io.ReadCloser) { + _ = r.Close() + }(r) + + sdkConfig, err := config.LoadDefaultConfig(ctx, config.WithEndpointDiscovery(aws.EndpointDiscoveryEnabled)) + if err != nil { + return err + } + + u, err := url.Parse(key) + if err != nil { + return err + } + u.Path = strings.TrimLeft(u.Path, "/") + + s3Client := s3.NewFromConfig(sdkConfig) + _, err = s3Client.PutObject(ctx, &s3.PutObjectInput{ + Bucket: ptr.To(u.Host), + Key: ptr.To(u.Path), + Body: r, + }) + return err +} diff --git a/internal/s3/s3_test.go b/internal/s3/s3_test.go new file mode 100644 index 00000000..9ec23ff6 --- /dev/null +++ b/internal/s3/s3_test.go @@ -0,0 +1,50 @@ +package s3 + +import "testing" + +func TestIsS3(t *testing.T) { + type args struct { + path string + } + tests := []struct { + name string + args args + want bool + }{ + {"relative local", args{"test.sql"}, false}, + {"absolute local", args{"/home/test/test.sql"}, false}, + {"s3 bucket", args{"s3://test"}, true}, + {"s3 bucket file", args{"s3://test/test.sql"}, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsS3(tt.args.path); got != tt.want { + t.Errorf("IsS3() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestIsS3Dir(t *testing.T) { + type args struct { + path string + } + tests := []struct { + name string + args args + want bool + }{ + {"relative local", args{"test.sql"}, false}, + {"absolute local", args{"/home/test/test.sql"}, false}, + {"s3 bucket", args{"s3://test"}, true}, + {"s3 bucket file", args{"s3://test/test.sql"}, false}, + {"s3 bucket dir", args{"s3://test/subdir/"}, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := IsS3Dir(tt.args.path); got != tt.want { + t.Errorf("IsS3Dir() = %v, want %v", got, tt.want) + } + }) + } +}