Initial commit

.gitignore

@@ -0,0 +1,2 @@
+pgsplit
+*.sql

flag.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	// Stnadard
+	"flag"
+	"fmt"
+	"os"
+)
+
+type arrayFlags []string
+
+func (i *arrayFlags) String() string {
+	return ""
+}
+
+func (i *arrayFlags) Set(value string) error {
+	*i = append(*i, value)
+	return nil
+}
+
+var (
+	flagIncludeAll bool
+	flagExcludeAll bool
+	flagExcept     arrayFlags
+)
+
+func init() {
+	flag.BoolVar(&flagIncludeAll, "include-all", false, "Split out all databases, with provided exceptions.")
+	flag.BoolVar(&flagExcludeAll, "exclude-all", false, "Split out no databases, with provided exceptions.")
+	flag.Var(&flagExcept, "except", "Exception to rule, can be used multiple times")
+	flag.Parse()
+
+	if flagIncludeAll && flagExcludeAll {
+		fmt.Println("-include-all and -exclude-all are mutually exclusive")
+		os.Exit(1)
+	}
+
+	if flagExcludeAll && len(flagExcept) == 0 {
+		fmt.Println("All databases are excluded and pgsplit is pointless.")
+		os.Exit(1)
+	}
+}

go.mod

@@ -0,0 +1,3 @@
+module pgsplit
+
+go 1.21.5

main.go

@@ -0,0 +1,93 @@
+package main
+
+import (
+	// Standard
+	"bufio"
+	"fmt"
+	"os"
+	"regexp"
+	"slices"
+)
+
+const VERSION = "v1"
+const MAXLINE = 1048576
+
+type Db struct {
+	Name      string
+	Completed bool
+	file      *os.File
+}
+
+func NewDb(name string) (db Db, err error) {
+	db.Name = name
+	fmt.Printf("Database %s\n", db.Name)
+
+	db.file, err = os.OpenFile(
+		db.Name+".sql",
+		os.O_CREATE|os.O_WRONLY|os.O_TRUNC,
+		0600,
+	)
+	return
+}
+
+func (db *Db) Close() {
+	if db.IsOpen() {
+		db.file.Close()
+		db.file = nil
+	}
+}
+
+func (db *Db) Write(s string) (err error) {
+	_, err = db.file.WriteString(s + "\n")
+	return
+}
+
+func (db *Db) IsOpen() bool {
+	return db.file != nil
+}
+
+func main() {
+	var prevDb Db
+	var db Db
+	var err error
+
+	dbStart := regexp.MustCompile(`^-- Database "([^"]+)" dump$`)
+	dbDone := regexp.MustCompile(`^-- PostgreSQL database dump complete$`)
+
+	scanner := bufio.NewScanner(os.Stdin)
+	buf := make([]byte, MAXLINE)
+	scanner.Buffer(buf, MAXLINE)
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		if dbMatch := dbStart.FindStringSubmatch(line); len(dbMatch) > 1 {
+			prevDb.Close()
+
+			dbName := dbMatch[1]
+			if (flagIncludeAll && slices.Contains(flagExcept, dbName)) || (flagExcludeAll && !slices.Contains(flagExcept, dbName)) {
+				fmt.Printf("Skipping %s\n", dbName)
+				continue
+			}
+
+			if db, err = NewDb(dbName); err != nil {
+				fmt.Printf("%s: %s\n", dbName, err)
+				os.Exit(1)
+			}
+		}
+
+		if db.IsOpen() {
+			db.Write(line)
+		}
+
+		if dbDone.MatchString(line) {
+			db.Completed = true
+			db.Close()
+			prevDb = db
+		}
+	}
+
+	if db.Name != "" && !db.Completed {
+		fmt.Printf("\n!!! Dump for '%s' is not complete !!!\n", db.Name)
+	}
+	db.Close()
+}