From fced33bfbf785f4f4304511e7145bcbc73d90f65 Mon Sep 17 00:00:00 2001 From: brendan-coughlan Date: Fri, 13 Dec 2024 06:07:08 -0800 Subject: [PATCH] some updates from feedback --- cmd/createSnapshot.go | 151 +++++++++++++++++++++++++++++++++++++ cmd/restoreSnapshot.go | 98 ++++++++++++++++++++++++ cmd/snapshot-create.go | 161 ---------------------------------------- cmd/snapshot-restore.go | 103 ------------------------- 4 files changed, 249 insertions(+), 264 deletions(-) create mode 100644 cmd/createSnapshot.go create mode 100644 cmd/restoreSnapshot.go delete mode 100644 cmd/snapshot-create.go delete mode 100644 cmd/snapshot-restore.go diff --git a/cmd/createSnapshot.go b/cmd/createSnapshot.go new file mode 100644 index 00000000..12209bd2 --- /dev/null +++ b/cmd/createSnapshot.go @@ -0,0 +1,151 @@ +package cmd + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + + "github.com/Layr-Labs/sidecar/internal/logger" + "github.com/spf13/cobra" +) + +var ( + output string + snapshotType string + compressionType string +) + +var createSnapshotCmd = &cobra.Command{ + Use: "create-snapshot", + Short: "Create a snapshot of the database", + Long: `Create a snapshot of the database. + +Currently available Type levels: +- archive (default): includes chain data, EigenModel state, rewards, and staker-operator table data. +`, + RunE: func(cmd *cobra.Command, args []string) error { + // Initialize logger + l, err := logger.NewLogger(&logger.LoggerConfig{Debug: true}) + if err != nil { + // If logger can't be initialized, just print and exit + fmt.Fprintf(os.Stderr, "Failed to initialize logger: %v\n", err) + os.Exit(1) + } + + // Validate compression type + if compressionType != "" && compressionType != "custom" { + l.Sugar().Fatalw("Invalid compression type", "compressionType", compressionType) + os.Exit(1) + } + + // Retrieve database connection details from global flags + dbHost, err := cmd.Flags().GetString("database.host") + if err != nil { + l.Sugar().Fatalw("Failed to get database host", "error", err) + os.Exit(1) + } + dbName, err := cmd.Flags().GetString("database.db_name") + if err != nil { + l.Sugar().Fatalw("Failed to get database name", "error", err) + os.Exit(1) + } + dbUser, err := cmd.Flags().GetString("database.user") + if err != nil { + l.Sugar().Fatalw("Failed to get database user", "error", err) + os.Exit(1) + } + dbPassword, err := cmd.Flags().GetString("database.password") + if err != nil { + l.Sugar().Fatalw("Failed to get database password", "error", err) + os.Exit(1) + } + dbPort, err := cmd.Flags().GetInt("database.port") + if err != nil { + l.Sugar().Fatalw("Failed to get database port", "error", err) + os.Exit(1) + } + schemaName, err := cmd.Flags().GetString("database.schema_name") + if err != nil { + l.Sugar().Fatalw("Failed to get schema name", "error", err) + os.Exit(1) + } + + // Validate Type + if snapshotType != "archive" { + l.Sugar().Warnw("Unsupported Type specified; falling back to 'archive'.", + "requested", snapshotType, + "used", "archive", + ) + snapshotType = "archive" + } + + // Log database connection details without password + l.Sugar().Infow("Database connection details", + "host", dbHost, + "name", dbName, + "user", dbUser, + "port", dbPort, + "schema", schemaName, + ) + + // Prepare pg_dump command using connection string + connectionString := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s", dbUser, dbPassword, dbHost, dbPort, dbName) + dumpCmd := []string{ + "pg_dump", + connectionString, + } + + // If a schema is specified, tell pg_dump to limit dump to that schema + if schemaName != "" { + dumpCmd = append(dumpCmd, "-n", schemaName) + } + + // If compression is requested, use the -Fc option + if compressionType == "custom" { + dumpCmd = append(dumpCmd, "-Fc") + } + + if output != "" { + // If output is specified, we consider output as a full file path + dumpDir := filepath.Dir(output) + + // Ensure the dump directory exists + if err := os.MkdirAll(dumpDir, 0755); err != nil { + l.Sugar().Fatalw("Failed to create output directory", "directory", dumpDir, "error", err) + os.Exit(1) + } + + dumpCmd = append(dumpCmd, "-f", output) + } else { + // If no output file is specified, dump to stdout + dumpCmd = append(dumpCmd, "-f", "/dev/stdout") + } + + // Log starting snapshot without including the password + safeConnectionString := fmt.Sprintf("postgresql://%s:****@%s:%d/%s?sslmode=disable", dbUser, dbHost, dbPort, dbName) + l.Sugar().Infow("Starting database snapshot", + "connection", safeConnectionString, + "command", dumpCmd, + ) + + // Execute pg_dump command + cmdExec := exec.Command(dumpCmd[0], dumpCmd[1:]...) + outputBytes, err := cmdExec.CombinedOutput() + if err != nil { + l.Sugar().Fatalw("Failed to create database snapshot", "error", err, "output", string(outputBytes)) + os.Exit(1) + } + + l.Sugar().Infow("Successfully created snapshot", "file", output) + + return nil + }, +} + +func init() { + rootCmd.AddCommand(createSnapshotCmd) + createSnapshotCmd.Flags().StringVarP(&output, "output", "f", "", "Path to save the snapshot file to (default is stdout if not specified)") + createSnapshotCmd.Flags().StringVar(&snapshotType, "snapshot-type", "archive", "The type of the snapshot: 'archive' only currently supported.") + createSnapshotCmd.Flags().StringVar(&compressionType, "compression-type", "", "Compression type for the snapshot: 'none' (default if empty), 'custom' for -Fc") +} diff --git a/cmd/restoreSnapshot.go b/cmd/restoreSnapshot.go new file mode 100644 index 00000000..e2d70c80 --- /dev/null +++ b/cmd/restoreSnapshot.go @@ -0,0 +1,98 @@ +package cmd + +import ( + "fmt" + "os" + "os/exec" + + "github.com/Layr-Labs/sidecar/internal/logger" + "github.com/spf13/cobra" +) + +var restoreSnapshotFromPath string +var cleanRestoreFlag bool + +var restoreSnapshotCmd = &cobra.Command{ + Use: "restore-snapshot", + Short: "Restore from a snapshot", + RunE: func(cmd *cobra.Command, args []string) error { + // Initialize logger + l, err := logger.NewLogger(&logger.LoggerConfig{Debug: true}) + if err != nil { + fmt.Fprintf(os.Stderr, "Failed to initialize logger: %v\n", err) + os.Exit(1) + } + + // Retrieve database connection details from global flags + dbHost, err := cmd.Flags().GetString("database.host") + if err != nil { + l.Sugar().Fatalw("Failed to get database host", "error", err) + os.Exit(1) + } + dbName, err := cmd.Flags().GetString("database.db_name") + if err != nil { + l.Sugar().Fatalw("Failed to get database name", "error", err) + os.Exit(1) + } + dbUser, err := cmd.Flags().GetString("database.user") + if err != nil { + l.Sugar().Fatalw("Failed to get database user", "error", err) + os.Exit(1) + } + dbPassword, err := cmd.Flags().GetString("database.password") + if err != nil { + l.Sugar().Fatalw("Failed to get database password", "error", err) + os.Exit(1) + } + dbPort, err := cmd.Flags().GetInt("database.port") + if err != nil { + l.Sugar().Fatalw("Failed to get database port", "error", err) + os.Exit(1) + } + + // Log the database connection details without password + l.Sugar().Infow("Database connection details", + "host", dbHost, + "name", dbName, + "user", dbUser, + "port", dbPort, + ) + + // Log the snapshot path + l.Sugar().Infow("Snapshot path", "path", restoreSnapshotFromPath) + + // Prepare pg_restore command using connection string + connectionString := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s", dbUser, dbPassword, dbHost, dbPort, dbName) + restoreCmd := []string{ + "pg_restore", + "--dbname", connectionString, + } + + if cleanRestoreFlag { + restoreCmd = append(restoreCmd, "--clean") + } + + if restoreSnapshotFromPath != "" { + restoreCmd = append(restoreCmd, restoreSnapshotFromPath) + } else { + restoreCmd = append(restoreCmd, "/dev/stdin") + } + + // Execute pg_restore command + cmdExec := exec.Command(restoreCmd[0], restoreCmd[1:]...) + outputBytes, err := cmdExec.CombinedOutput() + if err != nil { + l.Sugar().Errorw("Failed to restore from snapshot", "error", err, "output", string(outputBytes)) + return fmt.Errorf("restore snapshot failed: %v", err) + } + + l.Sugar().Infow("Successfully restored from snapshot") + return nil + }, +} + +func init() { + rootCmd.AddCommand(restoreSnapshotCmd) + restoreSnapshotCmd.Flags().StringVarP(&restoreSnapshotFromPath, "restore-snapshot-from-path", "p", "", "Path to snapshot file (use standard input if not specified)") + restoreSnapshotCmd.Flags().BoolVar(&cleanRestoreFlag, "clean", false, "Clean restore (default false)") +} diff --git a/cmd/snapshot-create.go b/cmd/snapshot-create.go deleted file mode 100644 index 752235dd..00000000 --- a/cmd/snapshot-create.go +++ /dev/null @@ -1,161 +0,0 @@ -package cmd - -import ( - "context" - "fmt" - "os" - "os/exec" - "path/filepath" - "time" - - "github.com/Layr-Labs/sidecar/internal/logger" - "github.com/spf13/cobra" -) - -var ( - snapshotCreateToPath string - snapshotFidelity string - s3Path string - compressSnapshot bool -) - -var snapshotCreateCmd = &cobra.Command{ - Use: "snapshot-create", - Short: "Create a snapshot of the database", - Long: `Create a snapshot of the database. - -Currently available fidelity levels: -- archive (default): includes chain data, EigenModel state, rewards, and staker-operator table data. -`, - RunE: func(cmd *cobra.Command, args []string) error { - // Initialize logger (in a real scenario, you'd likely pass in a config) - l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: true}) - - // Retrieve database connection details from global flags - dbHost, _ := cmd.Flags().GetString("database.host") - dbName, _ := cmd.Flags().GetString("database.db_name") - dbUser, _ := cmd.Flags().GetString("database.user") - dbPassword, _ := cmd.Flags().GetString("database.password") - dbPort, _ := cmd.Flags().GetInt("database.port") - schemaName, _ := cmd.Flags().GetString("database.schema_name") - - // Retrieve chain and fidelity - chain, _ := cmd.Flags().GetString("chain") - - // Validate fidelity - if snapshotFidelity != "archive" { - l.Sugar().Warnw("Unsupported fidelity specified; falling back to 'archive'.", - "requested", snapshotFidelity, - "used", "archive", - ) - snapshotFidelity = "archive" - } - - // Log database connection details at info level - l.Sugar().Infow("Database connection details", - "host", dbHost, - "name", dbName, - "user", dbUser, - "port", dbPort, - "schema", schemaName, - ) - - // If no specific file name is provided, generate one - if snapshotCreateToPath == "" { - fileNameSuffix := time.Now().UTC().Format("2006-01-02-15-04-05") - snapshotCreateToPath = fmt.Sprintf("sidecar-snapshot-%s-%s-%s.dump", chain, snapshotFidelity, fileNameSuffix) - } - - // Log the snapshot file path - l.Sugar().Infow("Snapshot will be created at this file", "file", snapshotCreateToPath) - - // Set the password in an environment variable for pg_dump - if err := os.Setenv("PGPASSWORD", dbPassword); err != nil { - l.Sugar().Fatalw("Failed to set environment variable", "error", err) - return err - } - - // Construct the pg_dump command with a context for timeout - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) - defer cancel() - - pgDumpCmd := exec.CommandContext(ctx, - "pg_dump", - "--host", dbHost, - "--port", fmt.Sprintf("%d", dbPort), - "--username", dbUser, - "--format", "custom", - "--file", snapshotCreateToPath, - "--schema", schemaName, - dbName, - ) - - // Direct command output to stdout/stderr for visibility - pgDumpCmd.Stdout = os.Stdout - pgDumpCmd.Stderr = os.Stderr - - l.Sugar().Infow("Running pg_dump command", - "command", pgDumpCmd.String(), - ) - - // Run the pg_dump command - if err := pgDumpCmd.Run(); err != nil { - if ctx.Err() == context.DeadlineExceeded { - l.Sugar().Errorw("pg_dump command timed out", - "error", err, - "command", pgDumpCmd.String(), - ) - return fmt.Errorf("pg_dump command timed out: %w", err) - } - l.Sugar().Errorw("Failed to run pg_dump", - "error", err, - "command", pgDumpCmd.String(), - ) - return fmt.Errorf("failed to run pg_dump: %w", err) - } - - l.Sugar().Infow("Successfully created snapshot", "file", snapshotCreateToPath) - - // Compress the snapshot if the flag is set - if compressSnapshot { - compressedSnapshot := snapshotCreateToPath + ".tar.gz" - l.Sugar().Infow("Compressing snapshot", "file", compressedSnapshot) - tarCmd := exec.Command("tar", "-czf", compressedSnapshot, "-C", filepath.Dir(snapshotCreateToPath), filepath.Base(snapshotCreateToPath)) - tarCmd.Stdout = os.Stdout - tarCmd.Stderr = os.Stderr - - if err := tarCmd.Run(); err != nil { - l.Sugar().Errorw("Snapshot compression failed", "error", err) - return fmt.Errorf("snapshot compression failed: %w", err) - } - - l.Sugar().Infow("Successfully compressed snapshot", "file", compressedSnapshot) - snapshotCreateToPath = compressedSnapshot - } - - // Upload to S3 if s3Path is specified - if s3Path != "" { - l.Sugar().Infow("Uploading to S3", "path", s3Path) - s3Cmd := exec.Command("aws", "s3", "cp", snapshotCreateToPath, s3Path) - s3Cmd.Stdout = os.Stdout - s3Cmd.Stderr = os.Stderr - - if err := s3Cmd.Run(); err != nil { - l.Sugar().Errorw("Upload to S3 failed", "error", err) - return fmt.Errorf("upload to S3 failed: %w", err) - } - - l.Sugar().Infow("Successfully uploaded snapshot to S3", "file", snapshotCreateToPath) - } - - return nil - }, -} - -func init() { - rootCmd.AddCommand(snapshotCreateCmd) - snapshotCreateCmd.Flags().StringVarP(&snapshotCreateToPath, "snapshot-create-to-path", "f", "", "Path to save the snapshot file to (default is generated based on timestamp and chain)") - snapshotCreateCmd.Flags().StringVar(&snapshotFidelity, "fidelity", "archive", "The fidelity level of the snapshot: 'archive' only currently supported.") - snapshotCreateCmd.Flags().StringVar(&s3Path, "s3-path", "", "S3 path to upload the snapshot to (e.g., s3://bucket/folder/), Default doesn't upload. It will only upload the compressed snapshot if --compress is set") - snapshotCreateCmd.Flags().BoolVar(&compressSnapshot, "compress", false, "Compress the snapshot and save as well as the original snapshot with .tar.gz extension") -} diff --git a/cmd/snapshot-restore.go b/cmd/snapshot-restore.go deleted file mode 100644 index 62b4f8a7..00000000 --- a/cmd/snapshot-restore.go +++ /dev/null @@ -1,103 +0,0 @@ -package cmd - -import ( - "context" - "fmt" - "os" - "os/exec" - "time" - - "github.com/Layr-Labs/sidecar/internal/logger" - "github.com/spf13/cobra" -) - -var snapshotRestoreFromPath string -var cleanRestore bool - -var snapshotRestoreCmd = &cobra.Command{ - Use: "snapshot-restore", - Short: "Restore from a snapshot", - RunE: func(cmd *cobra.Command, args []string) error { - // Initialize logger (in a real scenario, you'd likely pass in a config) - l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: true}) // or false depending on your need - - // Retrieve database connection details from global flags - dbHost, _ := cmd.Flags().GetString("database.host") - dbName, _ := cmd.Flags().GetString("database.db_name") - dbUser, _ := cmd.Flags().GetString("database.user") - dbPassword, _ := cmd.Flags().GetString("database.password") - dbPort, _ := cmd.Flags().GetInt("database.port") - - // Log database connection details at info level - l.Sugar().Infow("Database connection details", - "host", dbHost, - "name", dbName, - "user", dbUser, - "port", dbPort, - ) - - // Log the snapshot path - l.Sugar().Infow("Snapshot path", "path", snapshotRestoreFromPath) - - // Set the password in an environment variable for pg_restore - if err := os.Setenv("PGPASSWORD", dbPassword); err != nil { - l.Sugar().Fatalw("Failed to set environment variable", "error", err) - return err - } - - // Construct the pg_restore command with a context for timeout - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) - defer cancel() - - pgRestoreCmd := exec.CommandContext(ctx, - "pg_restore", - "--host", dbHost, - "--port", fmt.Sprintf("%d", dbPort), - "--username", dbUser, - "--dbname", dbName, - "--no-owner", - ) - - // Add the clean option if specified - if cleanRestore { - pgRestoreCmd.Args = append(pgRestoreCmd.Args, "--clean") - } - - // Add the snapshot path - pgRestoreCmd.Args = append(pgRestoreCmd.Args, snapshotRestoreFromPath) - - // Direct command output to stdout/stderr - pgRestoreCmd.Stdout = os.Stdout - pgRestoreCmd.Stderr = os.Stderr - - l.Sugar().Infow("Running pg_restore command", - "command", pgRestoreCmd.String(), - ) - - // Run the pg_restore command - if err := pgRestoreCmd.Run(); err != nil { - if ctx.Err() == context.DeadlineExceeded { - l.Sugar().Errorw("pg_restore command timed out", - "error", err, - "command", pgRestoreCmd.String(), - ) - return fmt.Errorf("pg_restore command timed out: %w", err) - } - l.Sugar().Errorw("Failed to run pg_restore", - "error", err, - "command", pgRestoreCmd.String(), - ) - return fmt.Errorf("failed to run pg_restore: %w", err) - } - - l.Sugar().Infow("Successfully restored from snapshot") - return nil - }, -} - -func init() { - rootCmd.AddCommand(snapshotRestoreCmd) - snapshotRestoreCmd.Flags().StringVarP(&snapshotRestoreFromPath, "snapshot-restore-from-path", "p", "", "Path to snapshot file (required)") - snapshotRestoreCmd.Flags().BoolVar(&cleanRestore, "clean", false, "Clean restore (default false)") - snapshotRestoreCmd.MarkFlagRequired("snapshot-restore-from-path") -}