From 1b2cc3866f13e296d57f73f35403465cf13e5bc5 Mon Sep 17 00:00:00 2001 From: Andrey Ivanov Date: Sun, 31 Jan 2021 09:53:32 +0300 Subject: [PATCH] =?UTF-8?q?=D0=A5=D0=BE=D1=87=D1=83=20=D0=BF=D0=B5=D1=80?= =?UTF-8?q?=D1=81=D0=B8=D1=81=D1=82=D0=B5=D0=BD=D1=82=D0=BD=D1=83=D1=8E=20?= =?UTF-8?q?=D0=91=D0=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pkg/dataset/generator.go | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/pkg/dataset/generator.go b/pkg/dataset/generator.go index 50c6bbb..511e59a 100644 --- a/pkg/dataset/generator.go +++ b/pkg/dataset/generator.go @@ -43,12 +43,16 @@ func NewPerson() (p Person) { } func FillDB(db *sql.DB, lim int) { - log.Println("Try to drop table USERS...") - if _, err := db.Exec(`DELETE FROM users`); err != nil { - log.Fatalf("can't exec query: %s", err.Error()) + var uCount int + if err := db.QueryRow(`SELECT COUNT(*) FROM users`).Scan(&uCount); err != nil { + log.Fatalf("can't get total of user profiles from DB: ") } - log.Printf("Try to generate %d rows and fill the DB...", lim) - for i := 1; i < lim; i++ { + uCount = lim - uCount + if uCount <= 0 { + log.Printf("Ok. We have more users then %s.", lim) + } + log.Printf("Try to generate %d rows and fill the DB...", uCount) + for i := 1; i < uCount; i++ { if i%100 == 0 { log.Printf("Successfully inserted %d rows", i) } @@ -63,7 +67,7 @@ func FillDB(db *sql.DB, lim int) { p.City, strings.Join(p.Interests, ","), ); err != nil { - log.Fatalf("can't insert row in DB. Inserted %d rows of %d: %s", i, lim, err.Error()) + log.Fatalf("can't insert row in DB. Inserted %d rows of %d: %s", i, uCount, err.Error()) } } log.Println("Table USERS filled successfully")