Skip to content

Commit

Permalink
feat(dump): Add support for dumps to an S3 bucket
Browse files Browse the repository at this point in the history
  • Loading branch information
gabe565 committed Mar 1, 2024
1 parent ecb9080 commit 31610b6
Show file tree
Hide file tree
Showing 7 changed files with 239 additions and 35 deletions.
46 changes: 41 additions & 5 deletions cmd/dump/dump.go
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -16,14 +22,27 @@ 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://clevyr-backups/dev/",
then the resulting filename might look like "s3://clevyr-backups/dev/` + dump.HelpFilename() + `".
The only exception is if a bucket name is provided without any sub-path (like "s3://backups"), then
the generated filename will be appended without requiring an ending "/".
`,

Args: cobra.MaximumNArgs(1),
ValidArgsFunction: validArgs,
Expand Down Expand Up @@ -102,7 +121,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)
}

Expand Down
21 changes: 17 additions & 4 deletions docs/kubedb_dump.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,25 @@ 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://clevyr-backups/dev/",
then the resulting filename might look like "s3://clevyr-backups/dev/clevyr_2022-01-09_094100.sql.gz".
The only exception is if a bucket name is provided without any sub-path (like "s3://backups"), then
the generated filename will be appended without requiring an ending "/".


```
kubedb dump [filename] [flags]
kubedb dump [filename | S3 URI] [flags]
```

### Options
Expand Down
18 changes: 18 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ 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/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
Expand All @@ -26,6 +28,22 @@ require (

require (
github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 // indirect
github.com/aws/aws-sdk-go-v2 v1.25.2 // 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
Expand Down
36 changes: 36 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
49 changes: 23 additions & 26 deletions internal/actions/dump/dump.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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),
Expand Down
54 changes: 54 additions & 0 deletions internal/s3/s3.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package s3

import (
"context"
"io"
"net/url"
"strings"

"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)
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
}
50 changes: 50 additions & 0 deletions internal/s3/s3_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
})
}
}

0 comments on commit 31610b6

Please sign in to comment.