pg/cmd/pg-fill/main.go

74 lines
1.6 KiB
Go

package main
import (
"context"
"database/sql"
"flag"
"fmt"
"log"
"os"
"path"
"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 {
presently := func() int {
dbname := path.Base(*c)
var n2 int
if row := pg.QueryRowContext(ctx, `SELECT pg_database_size('`+dbname+`')/1024/1024`); row.Err() != nil {
return 0
} else if err := row.Scan(&n2); err != nil {
return 0
}
return n2
}
report := func() {
log.Printf("filled %vMiB of requested %vMiB", presently(), *n)
}
go with.Every(ctx, 5*time.Second, func() {
report()
})
if _, err := pg.ExecContext(ctx, `
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 ctx.Err() == nil && presently() < *n {
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)
}
}
report()
return nil
})
}