GetTupleStatistics callback

arenadata · May 24, 2024 · 9812a2c · 9812a2c
1 parent 14c9a7e
commit 9812a2c
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 26 deletions.
diff --git a/backup/queries_statistics.go b/backup/queries_statistics.go
@@ -145,7 +145,7 @@ type TupleStatistic struct {
 	RelTuples float64
 }
 
-func getTupleStatisticsRows(connectionPool *dbconn.DBConn, tables []Table) (*sqlx.Rows, error) {
+func GetTupleStatistics(connectionPool *dbconn.DBConn, tables []Table, processRow func(tupleStat *TupleStatistic)) {
 	tablenames := make([]string, 0)
 	for _, table := range tables {
 		tablenames = append(tablenames, table.FQN())
@@ -163,19 +163,13 @@ func getTupleStatisticsRows(connectionPool *dbconn.DBConn, tables []Table) (*sql
 	ORDER BY n.nspname, c.relname`,
 		SchemaFilterClause("n"), utils.SliceToQuotedString(tablenames))
 
-	return connectionPool.Query(query)
-}
-
-func GetTupleStatistics(connectionPool *dbconn.DBConn, tables []Table) map[uint32]TupleStatistic {
-	results, err := getTupleStatisticsRows(connectionPool, tables)
+	rows, err := connectionPool.Query(query)
 	gplog.FatalOnError(err)
-	stats := make(map[uint32]TupleStatistic)
-	for results.Next() {
+	for rows.Next() {
 		var tupleStat TupleStatistic
-		err = results.StructScan(&tupleStat)
+		err = rows.StructScan(&tupleStat)
 		gplog.FatalOnError(err)
-		stats[tupleStat.Oid] = tupleStat
+		processRow(&tupleStat)
 	}
-	gplog.FatalOnError(results.Err())
-	return stats
+	gplog.FatalOnError(rows.Err())
 }
diff --git a/backup/wrappers.go b/backup/wrappers.go
@@ -777,15 +777,9 @@ func backupTableStatistics(statisticsFile *utils.FileWithByteCount, tables []Tab
 	for _, table := range tables {
 		tablesMap[table.Oid] = table
 	}
-	tupleStats, err := getTupleStatisticsRows(connectionPool, tables)
-	gplog.FatalOnError(err)
-	for tupleStats.Next() {
-		var tupleStat TupleStatistic
-		err = tupleStats.StructScan(&tupleStat)
-		gplog.FatalOnError(err)
-		printTupleStatisticsStatementForTable(statisticsFile, globalTOC, tablesMap[tupleStat.Oid], tupleStat)
-	}
-	gplog.FatalOnError(tupleStats.Err())
+	GetTupleStatistics(connectionPool, tables, func(tupleStat *TupleStatistic) {
+		printTupleStatisticsStatementForTable(statisticsFile, globalTOC, tablesMap[tupleStat.Oid], *tupleStat)
+	})
 	attStats, err := getAttributeStatisticsRows(connectionPool, tables)
 	gplog.FatalOnError(err)
 	for attStats.Next() {

diff --git a/integration/statistics_create_test.go b/integration/statistics_create_test.go
@@ -31,7 +31,10 @@ var _ = Describe("backup integration tests", func() {
 			tables[0].Oid = oldTableOid
 
 			beforeAttStats := backup.GetAttributeStatistics(connectionPool, tables)
-			beforeTupleStats := backup.GetTupleStatistics(connectionPool, tables)
+			beforeTupleStats := make(map[uint32]backup.TupleStatistic)
+			backup.GetTupleStatistics(connectionPool, tables, func(tupleStat *backup.TupleStatistic) {
+				beforeTupleStats[tupleStat.Oid] = *tupleStat
+			})
 			beforeTupleStat := beforeTupleStats[oldTableOid]
 
 			// Drop and recreate the table to clear the statistics
@@ -45,7 +48,10 @@ var _ = Describe("backup integration tests", func() {
 			newTableOid := testutils.OidFromObjectName(connectionPool, "public", "foo", backup.TYPE_RELATION)
 			tables[0].Oid = newTableOid
 			afterAttStats := backup.GetAttributeStatistics(connectionPool, tables)
-			afterTupleStats := backup.GetTupleStatistics(connectionPool, tables)
+			afterTupleStats := make(map[uint32]backup.TupleStatistic)
+			backup.GetTupleStatistics(connectionPool, tables, func(tupleStat *backup.TupleStatistic) {
+				afterTupleStats[tupleStat.Oid] = *tupleStat
+			})
 			afterTupleStat := afterTupleStats[newTableOid]
 
 			oldAtts := beforeAttStats[oldTableOid]
@@ -75,7 +81,10 @@ var _ = Describe("backup integration tests", func() {
 			tables[0].Oid = oldTableOid
 
 			beforeAttStats := backup.GetAttributeStatistics(connectionPool, tables)
-			beforeTupleStats := backup.GetTupleStatistics(connectionPool, tables)
+			beforeTupleStats := make(map[uint32]backup.TupleStatistic)
+			backup.GetTupleStatistics(connectionPool, tables, func(tupleStat *backup.TupleStatistic) {
+				beforeTupleStats[tupleStat.Oid] = *tupleStat
+			})
 			beforeTupleStat := beforeTupleStats[oldTableOid]
 
 			// Drop and recreate the table to clear the statistics
@@ -89,7 +98,10 @@ var _ = Describe("backup integration tests", func() {
 			newTableOid := testutils.OidFromObjectName(connectionPool, "public", "foo''\"''''bar", backup.TYPE_RELATION)
 			tables[0].Oid = newTableOid
 			afterAttStats := backup.GetAttributeStatistics(connectionPool, tables)
-			afterTupleStats := backup.GetTupleStatistics(connectionPool, tables)
+			afterTupleStats := make(map[uint32]backup.TupleStatistic)
+			backup.GetTupleStatistics(connectionPool, tables, func(tupleStat *backup.TupleStatistic) {
+				afterTupleStats[tupleStat.Oid] = *tupleStat
+			})
 			afterTupleStat := afterTupleStats[newTableOid]
 
 			oldAtts := beforeAttStats[oldTableOid]

diff --git a/integration/statistics_queries_test.go b/integration/statistics_queries_test.go
@@ -83,7 +83,10 @@ var _ = Describe("backup integration tests", func() {
 	})
 	Describe("GetTupleStatistics", func() {
 		It("returns tuple statistics for a table", func() {
-			tupleStats := backup.GetTupleStatistics(connectionPool, tables)
+			tupleStats := make(map[uint32]backup.TupleStatistic)
+			backup.GetTupleStatistics(connectionPool, tables, func(tupleStat *backup.TupleStatistic) {
+				tupleStats[tupleStat.Oid] = *tupleStat
+			})
 			Expect(tupleStats).To(HaveLen(1))
 			tableTupleStats := tupleStats[tableOid]