From 6e5db089c8740fec0be8895c4b85a16b34883992 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <shlomi-noach@github.com>
Date: Mon, 29 Aug 2016 09:58:31 +0200
Subject: [PATCH] supporting onRowCountComplete hook

---
 go/logic/migrator.go | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/go/logic/migrator.go b/go/logic/migrator.go
index e767455..c44d196 100644
--- a/go/logic/migrator.go
+++ b/go/logic/migrator.go
@@ -379,13 +379,24 @@ func (this *Migrator) countTableRows() (err error) {
 		log.Debugf("Noop operation; not really counting table rows")
 		return nil
 	}
+
+	countRowsFunc := func() error {
+		if err := this.inspector.CountTableRows(); err != nil {
+			return err
+		}
+		if err := this.hooksExecutor.onRowCountComplete(); err != nil {
+			return err
+		}
+		return nil
+	}
+
 	if this.migrationContext.ConcurrentCountTableRows {
-		go this.inspector.CountTableRows()
 		log.Infof("As instructed, counting rows in the background; meanwhile I will use an estimated count, and will update it later on")
+		go countRowsFunc()
 		// and we ignore errors, because this turns to be a background job
 		return nil
 	}
-	return this.inspector.CountTableRows()
+	return countRowsFunc()
 }
 
 // Migrate executes the complete migration logic. This is *the* major gh-ost function.