From f523c6eaf08c8b26c261a625c65d49942e900ce7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Magnus=20=C3=85hall?= <magnus.ahall@gibon.se>
Date: Tue, 23 Apr 2024 20:17:43 +0200
Subject: [PATCH] Initial commit

---
 .gitignore |  2 ++
 flag.go    | 42 ++++++++++++++++++++++++
 go.mod     |  3 ++
 main.go    | 93 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 140 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 flag.go
 create mode 100644 go.mod
 create mode 100644 main.go

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..3e6f252
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+pgsplit
+*.sql
diff --git a/flag.go b/flag.go
new file mode 100644
index 0000000..acbc6ac
--- /dev/null
+++ b/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)
+	}
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..1e57484
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module pgsplit
+
+go 1.21.5
diff --git a/main.go b/main.go
new file mode 100644
index 0000000..b383e34
--- /dev/null
+++ b/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()
+}