create pg-fill

cmd/pg-fill/main.go

@@ -0,0 +1,77 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"pg/src/with"
+
+	_ "github.com/lib/pq"
+)
+
+func main() {
+	if err := with.Context(run); err != nil {
+		panic(err)
+	}
+}
+
+func run(ctx context.Context) error {
+	fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
+	c := fs.String("c", "postgresql://pulsegres:pulsegres@localhost:15432", "conn string")
+	n := fs.Int("n", 12, "mb to generate")
+	if err := fs.Parse(os.Args[1:]); err != nil {
+		panic(err)
+	}
+
+	return with.PSQL(ctx, *c, func(pg *sql.DB) error {
+		report := func() error {
+			var n2 uint64
+			if row := pg.QueryRowContext(ctx, `SELECT pg_total_relation_size('fill_with_data')/1024/1024`); row.Err() != nil {
+				return row.Err()
+			} else if err := row.Scan(&n2); err != nil {
+				return err
+			}
+			log.Printf("filled %vMiB of requested %vMiB", n2, *n)
+			return nil
+		}
+
+		go with.Every(ctx, 5*time.Second, func() {
+			report()
+		})
+
+		if _, err := pg.ExecContext(ctx, `
+			DROP TABLE IF EXISTS fill_with_data;
+			CREATE TABLE IF NOT EXISTS fill_with_data(x BIGINT, y DOUBLE PRECISION, z DOUBLE PRECISION);
+		`); err != nil {
+			return fmt.Errorf("failed initial table setup: %w", err)
+		}
+
+		// https://gist.github.com/ololobus/5b25c432f208d7eb31051a5f238dffff
+		// 2e6=1GB, so 2e6/8=12MB
+		for i := 0; i < (*n)/12; i++ {
+			for ctx.Err() == nil {
+				if _, err := pg.ExecContext(ctx, `
+			   	INSERT INTO fill_with_data (x, y, z)
+			   	SELECT ROUND(RANDOM()), RANDOM(), RANDOM()
+			   	FROM generate_series(1, 2e6/8)
+			   `); err != nil {
+					return fmt.Errorf("failed lo_from_bytea: %w", err)
+				} else {
+					break
+				}
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				default:
+				}
+			}
+		}
+
+		return report()
+	})
+}