Add context support and cancelable readers to PostgreSQL+S3 integration

[ambyte]

Enhance PostgreSQL and S3 integration with context support and cancelable readers. Update Dump and DumpZip methods to return io.ReadCloser, allowing for process cancellation. Modify S3Upload to accept context for better control over uploads.

Summary by CodeRabbit

• Bug Fixes
  • Resolved hanging PostgreSQL sessions that occurred when dump operations were interrupted or cancelled.
  • Improved timeout handling for remote storage uploads.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds context-aware cancellation and proper resource cleanup across PostgreSQL dumps and S3 uploads. A new cancelableReadCloser wrapper enables cancellation of running pg_dump processes when consumers stop reading. Both Dump and DumpZip now return io.ReadCloser for proper cleanup. S3Upload accepts a context parameter with a dedicated 30-second timeout for metadata checks. The execution service ensures dump readers are closed and propagates context through the storage API.

Changes

Context-aware cancellation flow

Layer / File(s)	Summary
Cancelable dump reader foundation internal/integration/postgres/postgres.go	Introduces cancelableReadCloser wrapper that coordinates context cancellation via sync.Once and a done channel. Client.Dump now returns io.ReadCloser, uses exec.CommandContext, and streams via io.Pipe to enable cancellation of running pg_dump processes when the reader is closed.
Cascading cancellation through zip internal/integration/postgres/postgres.go	Client.DumpZip now returns io.ReadCloser and wraps the ZIP stream with cancelableReadCloser, so closing the ZIP reader propagates cancellation to the underlying dump and ultimately cancels the pg_dump process.
Context propagation in S3 upload internal/integration/storage/s3.go	S3Upload method signature now accepts context.Context and passes it to the S3 uploader. The HeadObject metadata check is wrapped with a dedicated 30-second timeout context instead of using a fixed context.
Integration in execution service internal/service/executions/run_execution.go	Dump reader is deferred closed immediately after creation to ensure cleanup. StorageClient.S3Upload call now passes request context, enabling timeout and cancellation propagation through the S3 upload.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐰 A cancellation charm takes flight,
Dump streams close when readers close tight,
Context flows through S3 and more,
Resources freed, no hangs to endure—
Graceful shutdowns, bunny's delight!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically summarizes the main changes: adding context support and cancelable readers to PostgreSQL and S3 integration, which directly aligns with the changeset modifications across all three files.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

internal/integration/storage/s3.go (1)

106-110: 💤 Low value

HeadObject ignores parent context cancellation.

The timeout context derives from context.Background() rather than ctx. If the caller cancels ctx after the upload completes, HeadObject continues for up to 30 seconds before returning.

If immediate cancellation is desired, derive from the parent context:

Suggested fix

-	headCtx, cancelHeadCtx := context.WithTimeout(context.Background(), 30*time.Second)
+	headCtx, cancelHeadCtx := context.WithTimeout(ctx, 30*time.Second)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@internal/integration/storage/s3.go` around lines 106 - 110, The HeadObject
call uses a timeout context derived from context.Background() so it won't
observe caller cancellation; change the context creation to derive from the
parent ctx (use context.WithTimeout(ctx, 30*time.Second)) so headCtx and
cancelHeadCtx for s3Client.HeadObject will be canceled if the incoming ctx is
canceled, keeping defer cancelHeadCtx() as-is.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@internal/integration/storage/s3.go`:
- Around line 106-110: The HeadObject call uses a timeout context derived from
context.Background() so it won't observe caller cancellation; change the context
creation to derive from the parent ctx (use context.WithTimeout(ctx,
30*time.Second)) so headCtx and cancelHeadCtx for s3Client.HeadObject will be
canceled if the incoming ctx is canceled, keeping defer cancelHeadCtx() as-is.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: d9a3024a-1593-46f9-b1c4-6b9c5a9cca67

📥 Commits

Reviewing files that changed from the base of the PR and between 70e0d9e and b0392c9.

📒 Files selected for processing (3)

• internal/integration/postgres/postgres.go
• internal/integration/storage/s3.go
• internal/service/executions/run_execution.go