diff --git a/internal/integration/postgres/postgres.go b/internal/integration/postgres/postgres.go index 97b2d64e..331bcad3 100644 --- a/internal/integration/postgres/postgres.go +++ b/internal/integration/postgres/postgres.go @@ -3,10 +3,12 @@ package postgres import ( "archive/zip" "bytes" + "context" "fmt" "io" "os" "os/exec" + "sync" "github.com/eduardolat/pgbackweb/internal/util/strutil" "github.com/orsinium-labs/enum" @@ -66,6 +68,27 @@ var ( type Client struct{} +type cancelableReadCloser struct { + reader io.ReadCloser + cancel context.CancelFunc + done <-chan struct{} + once sync.Once +} + +func (r *cancelableReadCloser) Read(p []byte) (int, error) { + return r.reader.Read(p) +} + +func (r *cancelableReadCloser) Close() error { + var err error + r.once.Do(func() { + r.cancel() + err = r.reader.Close() + <-r.done + }) + return err +} + func New() *Client { return &Client{} } @@ -138,10 +161,12 @@ type DumpParams struct { } // Dump runs the pg_dump command with the given parameters. It returns the SQL -// dump as an io.Reader. +// dump as an io.ReadCloser. Closing the returned reader cancels the running +// pg_dump process so the PostgreSQL session is not left hanging when a +// downstream consumer stops reading early. func (Client) Dump( version PGVersion, connString string, params ...DumpParams, -) io.Reader { +) io.ReadCloser { pickedParams := DumpParams{} if len(params) > 0 { pickedParams = params[0] @@ -168,51 +193,76 @@ func (Client) Dump( } errorBuffer := &bytes.Buffer{} + ctx, cancel := context.WithCancel(context.Background()) reader, writer := io.Pipe() - cmd := exec.Command(version.Value.PGDump, args...) + done := make(chan struct{}) + cmd := exec.CommandContext(ctx, version.Value.PGDump, args...) cmd.Stdout = writer cmd.Stderr = errorBuffer go func() { - defer writer.Close() + defer close(done) if err := cmd.Run(); err != nil { - writer.CloseWithError(fmt.Errorf( + _ = writer.CloseWithError(fmt.Errorf( "error running pg_dump v%s: %s", version.Value.Version, errorBuffer.String(), )) + return } + + _ = writer.Close() }() - return reader + return &cancelableReadCloser{ + reader: reader, + cancel: cancel, + done: done, + } } // DumpZip runs the pg_dump command with the given parameters and returns the -// ZIP-compressed SQL dump as an io.Reader. +// ZIP-compressed SQL dump as an io.ReadCloser. Closing the returned reader also +// cancels the underlying pg_dump process. func (c *Client) DumpZip( version PGVersion, connString string, params ...DumpParams, -) io.Reader { +) io.ReadCloser { dumpReader := c.Dump(version, connString, params...) reader, writer := io.Pipe() + done := make(chan struct{}) go func() { - defer writer.Close() + defer close(done) + defer dumpReader.Close() zipWriter := zip.NewWriter(writer) - defer zipWriter.Close() + defer func() { + if err := zipWriter.Close(); err != nil { + _ = writer.CloseWithError(fmt.Errorf("error closing zip file: %w", err)) + return + } + + _ = writer.Close() + }() fileWriter, err := zipWriter.Create("dump.sql") if err != nil { - writer.CloseWithError(fmt.Errorf("error creating zip file: %w", err)) + _ = writer.CloseWithError(fmt.Errorf("error creating zip file: %w", err)) return } if _, err := io.Copy(fileWriter, dumpReader); err != nil { - writer.CloseWithError(fmt.Errorf("error writing to zip file: %w", err)) + _ = writer.CloseWithError(fmt.Errorf("error writing to zip file: %w", err)) return } }() - return reader + return &cancelableReadCloser{ + reader: reader, + cancel: func() { + _ = dumpReader.Close() + }, + done: done, + } } // RestoreZip downloads or copies the ZIP from the given url or path, unzips it, diff --git a/internal/integration/storage/s3.go b/internal/integration/storage/s3.go index a2e9bed9..d218bdd8 100644 --- a/internal/integration/storage/s3.go +++ b/internal/integration/storage/s3.go @@ -71,10 +71,11 @@ func (Client) S3Test( return nil } -// S3Upload uploads a file to S3 from a reader. +// S3Upload uploads a file to S3 from a reader using the provided context. // // Returns the file size, in bytes. func (Client) S3Upload( + ctx context.Context, accessKey, secretKey, region, endpoint, bucketName, key string, fileReader io.Reader, ) (int64, error) { @@ -90,7 +91,7 @@ func (Client) S3Upload( uploader := manager.NewUploader(s3Client) _, err = uploader.Upload( - context.TODO(), + ctx, &s3.PutObjectInput{ Bucket: aws.String(bucketName), Key: aws.String(key), @@ -102,8 +103,11 @@ func (Client) S3Upload( return 0, fmt.Errorf("failed to upload file to S3: %w", err) } + headCtx, cancelHeadCtx := context.WithTimeout(context.Background(), 30*time.Second) + defer cancelHeadCtx() + fileHead, err := s3Client.HeadObject( - context.TODO(), + headCtx, &s3.HeadObjectInput{ Bucket: aws.String(bucketName), Key: aws.String(key), diff --git a/internal/service/executions/run_execution.go b/internal/service/executions/run_execution.go index db86651a..7b7f8929 100644 --- a/internal/service/executions/run_execution.go +++ b/internal/service/executions/run_execution.go @@ -107,6 +107,7 @@ func (s *Service) RunExecution(ctx context.Context, backupID uuid.UUID) error { NoComments: back.BackupOptNoComments, }, ) + defer dumpReader.Close() date := time.Now().Format(timeutil.LayoutSlashYYYYMMDD) file := fmt.Sprintf( @@ -133,6 +134,7 @@ func (s *Service) RunExecution(ctx context.Context, backupID uuid.UUID) error { if !back.BackupIsLocal { fileSize, err = s.ints.StorageClient.S3Upload( + ctx, back.DecryptedDestinationAccessKey, back.DecryptedDestinationSecretKey, back.DestinationRegion.String, back.DestinationEndpoint.String, back.DestinationBucketName.String, path, dumpReader,