pgsplit/main.go

package main

import (
	// Standard
	"bufio"
	"fmt"
	"os"
	"regexp"
	"slices"
)

const VERSION = "v5"
const MAXLINE = 1048576

type Db struct {
	Name      string
	Completed bool
	file      *os.File
}

func NewDb(name string) (db Db) {
	db.Name = name
	return
}

func (db *Db) Open() (err error) {
	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]
			db = NewDb(dbName)
			if (flagIncludeAll && slices.Contains(flagExcept, dbName)) || (flagExcludeAll && !slices.Contains(flagExcept, dbName)) {
				fmt.Printf("Skipping   %s\n", dbName)
				continue
			}

			fmt.Printf("Extracting %s\n", db.Name)
			if flagList {
				continue
			}
			if err = db.Open(); err != nil {
				fmt.Printf("%s: %s\n", dbName, err)
				os.Exit(1)
			}
		}

		if !flagList && 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()
}