fix plan tests

dolthub · Feb 11, 2025 · fe72f62 · fe72f62
1 parent 6df2999
commit fe72f62
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 19 deletions.
diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_harness.go b/go/libraries/doltcore/sqle/enginetest/dolt_harness.go
@@ -34,7 +34,6 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/branch_control"
 	"github.com/dolthub/dolt/go/libraries/doltcore/dtestutils"
 	"github.com/dolthub/dolt/go/libraries/doltcore/env"
-	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
 	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/kvexec"
@@ -259,10 +258,6 @@ func (d *DoltHarness) NewEngine(t *testing.T) (enginetest.QueryEngine, error) {
 		}
 		statsPro := statspro.NewStatsCoord(doltProvider, ctxGen, sqlCtx.Session.GetLogger().Logger, bThreads, d.multiRepoEnv.GetEnv(d.multiRepoEnv.GetFirstDatabase()))
 		statsPro.SetTimers(int64(1*time.Nanosecond), int64(1*time.Second), int64(1*time.Second))
-		err = statsPro.Restart(ctx)
-		if err != nil {
-			return nil, err
-		}
 		d.statsPro = statsPro
 
 		e, err := enginetest.NewEngine(t, d, d.provider, d.setupData, d.statsPro)
@@ -296,20 +291,19 @@ func (d *DoltHarness) NewEngine(t *testing.T) (enginetest.QueryEngine, error) {
 		e = e.WithBackgroundThreads(bThreads)
 
 		if d.configureStats {
-			dSess := dsess.DSessFromSess(sqlCtx.Session)
-			dbCache := dSess.DatabaseCache(sqlCtx)
-			dsessDbs := make([]dsess.SqlDatabase, len(dbs))
-			for i, dbName := range dbs {
-				dsessDbs[i], _ = dbCache.GetCachedRevisionDb(fmt.Sprintf("%s/main", dbName), dbName)
-				fs, err := doltProvider.FileSystemForDatabase(dsessDbs[i].AliasedName())
-				if err != nil {
-					return nil, err
+			var dsessDbs []dsess.SqlDatabase
+			for _, db := range databases {
+				if sqlDb, ok := db.(dsess.SqlDatabase); ok {
+					dsessDbs = append(dsessDbs, sqlDb)
 				}
-				done, err := statsPro.Add(sqlCtx, dsessDbs[i], ref.NewBranchRef("main"), fs, false)
-				if err != nil {
-					return nil, err
-				}
-				<-done
+			}
+			if err := statsPro.Init(ctx, dsessDbs, false); err != nil {
+				return nil, err
+			}
+
+			err = statsPro.Restart(ctx)
+			if err != nil {
+				return nil, err
 			}
 
 			statsOnlyQueries := filterStatsOnlyQueries(d.setupData)

diff --git a/go/libraries/doltcore/sqle/statspro/provider.go b/go/libraries/doltcore/sqle/statspro/provider.go
@@ -310,7 +310,7 @@ func (sc *StatsCoord) Init(ctx context.Context, dbs []dsess.SqlDatabase, keepSto
 
 	eg := errgroup.Group{}
 	for _, db := range dbs {
-		if db, ok := db.(*sqle.Database); ok { // exclude read replica dbs
+		if db, ok := db.(sqle.Database); ok { // exclude read replica dbs
 			br, err := db.DbData().Ddb.GetBranches(ctx)
 			if err != nil {
 				return err