da36506b89
P0 fixes: - ModelError.Is(): use exact matching instead of substring contains() - shouldClearStickySession: add context param for cancellation/tracing P1 fixes: - TODO stubs: return 501 Not Implemented errors - validateInstanceSignature: deduplicate to shared validateCodeSignature() - Error messages: standardize to English only - http.go: remove pseudo if-else with duplicate branches
99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package repository
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"io"
|
||
"os/exec"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
)
|
||
|
||
// PgDumper implements service.DBDumper using pg_dump/psql
|
||
type PgDumper struct {
|
||
cfg *config.DatabaseConfig
|
||
}
|
||
|
||
// NewPgDumper creates a new PgDumper
|
||
func NewPgDumper(cfg *config.Config) service.DBDumper {
|
||
return &PgDumper{cfg: &cfg.Database}
|
||
}
|
||
|
||
// Dump executes pg_dump and returns a streaming reader of the output
|
||
func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||
args := []string{
|
||
"-h", d.cfg.Host,
|
||
"-p", fmt.Sprintf("%d", d.cfg.Port),
|
||
"-U", d.cfg.User,
|
||
"-d", d.cfg.DBName,
|
||
"--no-owner",
|
||
"--no-acl",
|
||
"--clean",
|
||
"--if-exists",
|
||
}
|
||
|
||
cmd := exec.CommandContext(ctx, "pg_dump", args...)
|
||
if d.cfg.Password != "" {
|
||
cmd.Env = append(cmd.Environ(), "PGPASSWORD="+d.cfg.Password)
|
||
}
|
||
if d.cfg.SSLMode != "" {
|
||
cmd.Env = append(cmd.Environ(), "PGSSLMODE="+d.cfg.SSLMode)
|
||
}
|
||
|
||
stdout, err := cmd.StdoutPipe()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("create stdout pipe: %w", err)
|
||
}
|
||
|
||
if err := cmd.Start(); err != nil {
|
||
return nil, fmt.Errorf("start pg_dump: %w", err)
|
||
}
|
||
|
||
// 返回一个 ReadCloser：读 stdout，关闭时等待进程退出
|
||
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
|
||
}
|
||
|
||
// Restore executes psql to restore from a streaming reader
|
||
func (d *PgDumper) Restore(ctx context.Context, data io.Reader) error {
|
||
args := []string{
|
||
"-h", d.cfg.Host,
|
||
"-p", fmt.Sprintf("%d", d.cfg.Port),
|
||
"-U", d.cfg.User,
|
||
"-d", d.cfg.DBName,
|
||
"--single-transaction",
|
||
}
|
||
|
||
cmd := exec.CommandContext(ctx, "psql", args...)
|
||
if d.cfg.Password != "" {
|
||
cmd.Env = append(cmd.Environ(), "PGPASSWORD="+d.cfg.Password)
|
||
}
|
||
if d.cfg.SSLMode != "" {
|
||
cmd.Env = append(cmd.Environ(), "PGSSLMODE="+d.cfg.SSLMode)
|
||
}
|
||
|
||
cmd.Stdin = data
|
||
|
||
output, err := cmd.CombinedOutput()
|
||
if err != nil {
|
||
return fmt.Errorf("%v: %s", err, string(output))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// cmdReadCloser wraps a command stdout pipe and waits for the process on Close
|
||
type cmdReadCloser struct {
|
||
io.ReadCloser
|
||
cmd *exec.Cmd
|
||
}
|
||
|
||
func (c *cmdReadCloser) Close() error {
|
||
// Close the pipe first
|
||
_ = c.ReadCloser.Close()
|
||
// Wait for the process to exit
|
||
if err := c.cmd.Wait(); err != nil {
|
||
return fmt.Errorf("pg_dump exited with error: %w", err)
|
||
}
|
||
return nil
|
||
}
|