Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/repository/backup_pg_dumper.go

@@ -0,0 +1,98 @@
+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
+}