chore(devtools): ods db operations (#6661)

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: jmelahman

tools/ods/cmd/cherry-pick.go

@@ -1,7 +1,6 @@
 package cmd
 
 import (
-	"bufio"
 	"fmt"
 	"os"
 	"os/exec"
@@ -10,6 +9,8 @@ import (
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
+
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/prompt"
 )
 
 // CherryPickOptions holds options for the cherry-pick command
@@ -110,7 +111,7 @@ func runCherryPick(cmd *cobra.Command, args []string, opts *CherryPickOptions) {
 
 		// Prompt user for confirmation
 		if !opts.Yes {
-			if !promptYesNo(fmt.Sprintf("Auto-detected release version: %s. Continue? (yes/no): ", version)) {
+			if !prompt.Confirm(fmt.Sprintf("Auto-detected release version: %s. Continue? (yes/no): ", version)) {
 				log.Info("If you want to cherry-pick to a different release, use the --release flag. Exiting...")
 				return
 			}
@@ -223,26 +224,6 @@ func cherryPickToRelease(commitSHAs []string, branchSuffix, version, prTitle str
 	return prURL, nil
 }
 
-// promptYesNo prompts the user with a yes/no question and returns true for yes, false for no
-func promptYesNo(prompt string) bool {
-	reader := bufio.NewReader(os.Stdin)
-	for {
-		fmt.Print(prompt)
-		response, err := reader.ReadString('\n')
-		if err != nil {
-			log.Fatalf("Failed to read input: %v", err)
-		}
-		response = strings.TrimSpace(strings.ToLower(response))
-		if response == "yes" || response == "y" || response == "" {
-			return true
-		}
-		if response == "no" || response == "n" {
-			return false
-		}
-		fmt.Println("Please enter 'yes' or 'no'")
-	}
-}
-
 // getCurrentBranch returns the name of the current git branch
 func getCurrentBranch() (string, error) {
 	cmd := exec.Command("git", "branch", "--show-current")

tools/ods/cmd/db.go

@@ -0,0 +1,28 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+)
+
+// NewDBCommand creates the parent db command.
+func NewDBCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "db",
+		Short: "Database administration commands",
+		Long: `Database administration commands for managing PostgreSQL and Alembic migrations.
+
+Commands include dropping/recreating databases, creating and restoring snapshots,
+and managing Alembic migrations (upgrade, downgrade, current, history).`,
+	}
+
+	// Add subcommands
+	cmd.AddCommand(NewDBDropCommand())
+	cmd.AddCommand(NewDBDumpCommand())
+	cmd.AddCommand(NewDBRestoreCommand())
+	cmd.AddCommand(NewDBUpgradeCommand())
+	cmd.AddCommand(NewDBDowngradeCommand())
+	cmd.AddCommand(NewDBCurrentCommand())
+	cmd.AddCommand(NewDBHistoryCommand())
+
+	return cmd
+}

tools/ods/cmd/db_drop.go

@@ -0,0 +1,142 @@
+package cmd
+
+import (
+	"fmt"
+	"regexp"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/docker"
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/postgres"
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/prompt"
+)
+
+// validIdentifier matches valid PostgreSQL identifiers (letters, digits, underscores, starting with letter/underscore)
+var validIdentifier = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
+
+// DBDropOptions holds options for the db drop command.
+type DBDropOptions struct {
+	Yes    bool
+	Schema string
+}
+
+// NewDBDropCommand creates the db drop command.
+func NewDBDropCommand() *cobra.Command {
+	opts := &DBDropOptions{}
+
+	cmd := &cobra.Command{
+		Use:   "drop",
+		Short: "Drop and recreate the database",
+		Long: `Drop and recreate the PostgreSQL database.
+
+This command will:
+  1. Find the running PostgreSQL container
+  2. Drop all connections to the database
+  3. Drop the database (or schema if --schema is specified)
+  4. Recreate the database (or schema)
+
+WARNING: This is a destructive operation. All data will be lost.`,
+		Run: func(cmd *cobra.Command, args []string) {
+			runDBDrop(opts)
+		},
+	}
+
+	cmd.Flags().BoolVar(&opts.Yes, "yes", false, "Skip confirmation prompt")
+	cmd.Flags().StringVar(&opts.Schema, "schema", "", "Drop a specific schema instead of the entire database")
+
+	return cmd
+}
+
+func runDBDrop(opts *DBDropOptions) {
+	// Find PostgreSQL container
+	container, err := docker.FindPostgresContainer()
+	if err != nil {
+		log.Fatalf("Failed to find PostgreSQL container: %v", err)
+	}
+	log.Infof("Found PostgreSQL container: %s", container)
+
+	config := postgres.NewConfigFromEnv()
+
+	// Confirmation prompt
+	if !opts.Yes {
+		var msg string
+		if opts.Schema != "" {
+			msg = fmt.Sprintf("This will DROP the schema '%s' in database '%s'. All data will be lost. Continue? (yes/no): ",
+				opts.Schema, config.Database)
+		} else {
+			msg = fmt.Sprintf("This will DROP and RECREATE the database '%s'. All data will be lost. Continue? (yes/no): ",
+				config.Database)
+		}
+
+		if !prompt.Confirm(msg) {
+			log.Info("Aborted.")
+			return
+		}
+	}
+
+	env := config.Env()
+
+	if opts.Schema != "" {
+		// Validate schema name to prevent SQL injection
+		if !validIdentifier.MatchString(opts.Schema) {
+			log.Fatalf("Invalid schema name: %s", opts.Schema)
+		}
+
+		// Drop and recreate schema
+		log.Infof("Dropping schema: %s", opts.Schema)
+		dropSchemaSQL := fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE;", opts.Schema)
+		createSchemaSQL := fmt.Sprintf("CREATE SCHEMA %s;", opts.Schema)
+
+		args := append(config.PsqlArgs(), "-c", dropSchemaSQL)
+		if err := docker.ExecWithEnv(container, env, append([]string{"psql"}, args...)...); err != nil {
+			log.Fatalf("Failed to drop schema: %v", err)
+		}
+
+		args = append(config.PsqlArgs(), "-c", createSchemaSQL)
+		if err := docker.ExecWithEnv(container, env, append([]string{"psql"}, args...)...); err != nil {
+			log.Fatalf("Failed to create schema: %v", err)
+		}
+
+		log.Infof("Schema '%s' dropped and recreated successfully", opts.Schema)
+	} else {
+		// Drop and recreate entire database
+		log.Infof("Dropping database: %s", config.Database)
+
+		// Use template1 as maintenance database (can't drop a DB while connected to it)
+		maintenanceDB := "template1"
+
+		// Terminate existing connections
+		// Validate database name to prevent SQL injection
+		if !validIdentifier.MatchString(config.Database) {
+			log.Fatalf("Invalid database name: %s", config.Database)
+		}
+
+		// Terminate existing connections
+		terminateSQL := fmt.Sprintf(
+			"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%s' AND pid <> pg_backend_pid();",
+			config.Database)
+
+		args := []string{"psql", "-U", config.User, "-d", maintenanceDB, "-c", terminateSQL}
+		if err := docker.ExecWithEnv(container, env, args...); err != nil {
+			log.Warnf("Failed to terminate connections (this may be okay): %v", err)
+		}
+
+		// Drop database
+		dropSQL := fmt.Sprintf("DROP DATABASE IF EXISTS %s;", config.Database)
+		args = []string{"psql", "-U", config.User, "-d", maintenanceDB, "-c", dropSQL}
+		if err := docker.ExecWithEnv(container, env, args...); err != nil {
+			log.Fatalf("Failed to drop database: %v", err)
+		}
+
+		// Create database
+		createSQL := fmt.Sprintf("CREATE DATABASE %s;", config.Database)
+		args = []string{"psql", "-U", config.User, "-d", maintenanceDB, "-c", createSQL}
+		if err := docker.ExecWithEnv(container, env, args...); err != nil {
+			log.Fatalf("Failed to create database: %v", err)
+		}
+
+		log.Infof("Database '%s' dropped and recreated successfully", config.Database)
+		log.Info("Run 'ods db upgrade' to apply migrations")
+	}
+}

tools/ods/cmd/db_dump.go

@@ -0,0 +1,144 @@
+package cmd
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/docker"
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/paths"
+	"github.com/onyx-dot-app/onyx/tools/ods/internal/postgres"
+)
+
+// DBDumpOptions holds options for the db dump command.
+type DBDumpOptions struct {
+	Format string
+	Schema string
+	Output string
+}
+
+// NewDBDumpCommand creates the db dump command.
+func NewDBDumpCommand() *cobra.Command {
+	opts := &DBDumpOptions{}
+
+	cmd := &cobra.Command{
+		Use:   "dump [output-file]",
+		Short: "Create a database snapshot",
+		Long: `Create a database snapshot using pg_dump.
+
+The snapshot is saved to the specified output file, or to the default
+snapshots directory (~/.local/share/onyx-dev/snapshots/) if no file is specified.
+
+Examples:
+  ods db dump                           # Creates onyx_<timestamp>.dump in snapshots dir
+  ods db dump mybackup.dump             # Creates mybackup.dump in snapshots dir
+  ods db dump /path/to/backup.sql       # Creates backup.sql at specified path
+  ods db dump --format sql              # Creates SQL format instead of custom format`,
+		Args: cobra.MaximumNArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			if len(args) > 0 {
+				opts.Output = args[0]
+			}
+			runDBDump(opts)
+		},
+	}
+
+	cmd.Flags().StringVar(&opts.Format, "format", "custom", "Output format: 'custom' (pg_dump -Fc) or 'sql' (plain SQL)")
+	cmd.Flags().StringVar(&opts.Schema, "schema", "", "Dump only a specific schema")
+
+	return cmd
+}
+
+func runDBDump(opts *DBDumpOptions) {
+	// Find PostgreSQL container
+	container, err := docker.FindPostgresContainer()
+	if err != nil {
+		log.Fatalf("Failed to find PostgreSQL container: %v", err)
+	}
+	log.Infof("Found PostgreSQL container: %s", container)
+
+	config := postgres.NewConfigFromEnv()
+
+	// Determine output file path
+	outputPath := determineOutputPath(opts.Output, opts.Format)
+
+	// Ensure output directory exists
+	outputDir := filepath.Dir(outputPath)
+	if err := os.MkdirAll(outputDir, 0755); err != nil {
+		log.Fatalf("Failed to create output directory: %v", err)
+	}
+
+	log.Infof("Dumping database '%s' to: %s", config.Database, outputPath)
+
+	// Build pg_dump arguments
+	args := config.PgDumpArgs(opts.Format)
+	if opts.Schema != "" {
+		args = append(args, "-n", opts.Schema)
+	}
+
+	// Create a temporary file in the container
+	containerTmpFile := "/tmp/onyx_dump_tmp"
+	args = append(args, "-f", containerTmpFile)
+
+	// Run pg_dump in container
+	env := config.Env()
+	pgDumpArgs := append([]string{"pg_dump"}, args...)
+	if err := docker.ExecWithEnv(container, env, pgDumpArgs...); err != nil {
+		log.Fatalf("Failed to run pg_dump: %v", err)
+	}
+
+	// Copy the dump file from container to host
+	if err := docker.CopyFromContainer(container, containerTmpFile, outputPath); err != nil {
+		log.Fatalf("Failed to copy dump file: %v", err)
+	}
+
+	// Clean up temporary file in container
+	_ = docker.Exec(container, "rm", "-f", containerTmpFile)
+
+	// Get file size for info
+	if info, err := os.Stat(outputPath); err == nil {
+		log.Infof("Dump completed successfully (%s)", humanizeBytes(info.Size()))
+	} else {
+		log.Info("Dump completed successfully")
+	}
+}
+
+// determineOutputPath determines the output file path based on options.
+func determineOutputPath(output string, format string) string {
+	ext := ".dump"
+	if format == "sql" {
+		ext = ".sql"
+	}
+
+	if output == "" {
+		// Generate default filename with timestamp
+		timestamp := time.Now().Format("20060102_150405")
+		filename := fmt.Sprintf("onyx_%s%s", timestamp, ext)
+		return filepath.Join(paths.SnapshotsDir(), filename)
+	}
+
+	// Check if output is just a filename (no directory)
+	if filepath.Dir(output) == "." {
+		return filepath.Join(paths.SnapshotsDir(), output)
+	}
+
+	return output
+}
+
+// humanizeBytes converts bytes to a human-readable string.
+func humanizeBytes(bytes int64) string {
+	const unit = 1024
+	if bytes < unit {
+		return fmt.Sprintf("%d B", bytes)
+	}
+	div, exp := int64(unit), 0
+	for n := bytes / unit; n >= unit; n /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f %cB", float64(bytes)/float64(div), "KMGTPE"[exp])
+}