Merge pull request #8589 from dolthub/taylor/stats

Fix dolt_statistics table for multiple schemas for doltgres

go/go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/dolthub/fslock v0.0.3
 	github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718
 	github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81
-	github.com/dolthub/vitess v0.0.0-20241126205153-88972ec5fe52
+	github.com/dolthub/vitess v0.0.0-20241126223332-cd8f828f26ac
 	github.com/dustin/go-humanize v1.0.1
 	github.com/fatih/color v1.13.0
 	github.com/flynn-archive/go-shlex v0.0.0-20150515145356-3f9db97f8568
@@ -57,7 +57,7 @@ require (
 	github.com/cespare/xxhash/v2 v2.2.0
 	github.com/creasty/defaults v1.6.0
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
-	github.com/dolthub/go-mysql-server v0.18.2-0.20241126213642-de3d54a398ac
+	github.com/dolthub/go-mysql-server v0.18.2-0.20241127000145-a1809677932e
 	github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63
 	github.com/dolthub/swiss v0.1.0
 	github.com/goccy/go-json v0.10.2

go/go.sum

@@ -183,8 +183,8 @@ github.com/dolthub/fslock v0.0.3 h1:iLMpUIvJKMKm92+N1fmHVdxJP5NdyDK5bK7z7Ba2s2U=
 github.com/dolthub/fslock v0.0.3/go.mod h1:QWql+P17oAAMLnL4HGB5tiovtDuAjdDTPbuqx7bYfa0=
 github.com/dolthub/go-icu-regex v0.0.0-20240916130659-0118adc6b662 h1:aC17hZD6iwzBwwfO5M+3oBT5E5gGRiQPdn+vzpDXqIA=
 github.com/dolthub/go-icu-regex v0.0.0-20240916130659-0118adc6b662/go.mod h1:KPUcpx070QOfJK1gNe0zx4pA5sicIK1GMikIGLKC168=
-github.com/dolthub/go-mysql-server v0.18.2-0.20241126213642-de3d54a398ac h1:tit+TpLRkR918++5HAFLvbfBQklUsn7+XCYJOCZCZ+Q=
-github.com/dolthub/go-mysql-server v0.18.2-0.20241126213642-de3d54a398ac/go.mod h1:grCMUuyP/ZaQh8JDhr/AdFc/+fc9h71uislzsV4k3UM=
+github.com/dolthub/go-mysql-server v0.18.2-0.20241127000145-a1809677932e h1:2oysRPgywCpyW/h4m6DxYUpUsdZ+JkJXvC51TVr4sUE=
+github.com/dolthub/go-mysql-server v0.18.2-0.20241127000145-a1809677932e/go.mod h1:QdaXQKE8XFwM4P1yN14m2eydx4V2xyuqpQp4tmNoXzQ=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63 h1:OAsXLAPL4du6tfbBgK0xXHZkOlos63RdKYS3Sgw/dfI=
 github.com/dolthub/gozstd v0.0.0-20240423170813-23a2903bca63/go.mod h1:lV7lUeuDhH5thVGDCKXbatwKy2KW80L4rMT46n+Y2/Q=
 github.com/dolthub/ishell v0.0.0-20240701202509-2b217167d718 h1:lT7hE5k+0nkBdj/1UOSFwjWpNxf+LCApbRHgnCA17XE=
@@ -197,8 +197,8 @@ github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 h1:7/v8q9X
 github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY=
 github.com/dolthub/swiss v0.1.0 h1:EaGQct3AqeP/MjASHLiH6i4TAmgbG/c4rA6a1bzCOPc=
 github.com/dolthub/swiss v0.1.0/go.mod h1:BeucyB08Vb1G9tumVN3Vp/pyY4AMUnr9p7Rz7wJ7kAQ=
-github.com/dolthub/vitess v0.0.0-20241126205153-88972ec5fe52 h1:RH0eygj4DLPQ6fvJCBBk4ZOq3PpeIQU/1ot6Fye9WCU=
-github.com/dolthub/vitess v0.0.0-20241126205153-88972ec5fe52/go.mod h1:alcJgfdyIhFaAiYyEmuDCFSLCzedz3KCaIclLoCUtJg=
+github.com/dolthub/vitess v0.0.0-20241126223332-cd8f828f26ac h1:A0U/OdIqdCkAV0by7MVBbnSyZBsa94ZjIZxx7PhjBW4=
+github.com/dolthub/vitess v0.0.0-20241126223332-cd8f828f26ac/go.mod h1:alcJgfdyIhFaAiYyEmuDCFSLCzedz3KCaIclLoCUtJg=
 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
 github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
 github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=

go/libraries/doltcore/sqle/database.go

@@ -675,7 +675,26 @@ func (db Database) getTableInsensitive(ctx *sql.Context, head *doltdb.Commit, ds
 			}
 		}
 	case doltdb.StatisticsTableName:
-		dt, found = dtables.NewStatisticsTable(ctx, db.Name(), db.ddb, asOf), true
+		if resolve.UseSearchPath && db.schemaName == "" {
+			schemaName, err := resolve.FirstExistingSchemaOnSearchPath(ctx, root)
+			if err != nil {
+				return nil, false, err
+			}
+			db.schemaName = schemaName
+		}
+
+		var tables []string
+		var err error
+		branch, ok := asOf.(string)
+		if ok && branch != "" {
+			tables, err = db.GetTableNamesAsOf(ctx, branch)
+		} else {
+			tables, err = db.GetTableNames(ctx)
+		}
+		if err != nil {
+			return nil, false, err
+		}
+		dt, found = dtables.NewStatisticsTable(ctx, db.Name(), db.schemaName, branch, tables), true
 	case doltdb.ProceduresTableName:
 		found = true
 		backingTable, _, err := db.getTable(ctx, root, doltdb.ProceduresTableName)
@@ -856,7 +875,14 @@ func (db Database) GetTableNamesAsOf(ctx *sql.Context, time interface{}) ([]stri
 		return nil, nil
 	}
 
-	tblNames, err := db.getAllTableNames(ctx, root, false)
+	showSystemTablesVar, err := ctx.GetSessionVariable(ctx, dsess.ShowSystemTables)
+	if err != nil {
+		return nil, err
+	}
+
+	showSystemTables := showSystemTablesVar.(int8) == 1
+
+	tblNames, err := db.getAllTableNames(ctx, root, showSystemTables)
 	if err != nil {
 		return nil, err
 	}

go/libraries/doltcore/sqle/dtables/statistics_table.go

@@ -15,8 +15,6 @@
 package dtables
 
 import (
-	"fmt"
-
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/stats"
 
@@ -28,21 +26,18 @@ import (
 
 // StatisticsTable is a sql.Table implementation that implements a system table which shows the dolt commit log
 type StatisticsTable struct {
-	dbName string
-	branch string
-	ddb    *doltdb.DoltDB
+	dbName     string
+	schemaName string
+	branch     string
+	tableNames []string
 }
 
 var _ sql.Table = (*StatisticsTable)(nil)
 var _ sql.StatisticsTable = (*StatisticsTable)(nil)
 
 // NewStatisticsTable creates a StatisticsTable
-func NewStatisticsTable(_ *sql.Context, dbName string, ddb *doltdb.DoltDB, asOf interface{}) sql.Table {
-	ret := &StatisticsTable{dbName: dbName, ddb: ddb}
-	if branch, ok := asOf.(string); ok {
-		ret.branch = branch
-	}
-	return ret
+func NewStatisticsTable(_ *sql.Context, dbName, schemaName, branch string, tableNames []string) sql.Table {
+	return &StatisticsTable{dbName: dbName, schemaName: schemaName, branch: branch, tableNames: tableNames}
 }
 
 // DataLength implements sql.StatisticsTable
@@ -73,28 +68,16 @@ func (st *StatisticsTable) DataLength(ctx *sql.Context) (uint64, error) {
 }
 
 type BranchStatsProvider interface {
-	GetTableDoltStats(ctx *sql.Context, branch, db, table string) ([]sql.Statistic, error)
+	GetTableDoltStats(ctx *sql.Context, branch, db, schema, table string) ([]sql.Statistic, error)
 }
 
 // RowCount implements sql.StatisticsTable
 func (st *StatisticsTable) RowCount(ctx *sql.Context) (uint64, bool, error) {
 	dSess := dsess.DSessFromSess(ctx.Session)
-	prov := dSess.Provider()
-
-	sqlDb, err := prov.Database(ctx, st.dbName)
-	if err != nil {
-		return 0, false, err
-	}
-
-	tables, err := sqlDb.GetTableNames(ctx)
-	if err != nil {
-		return 0, false, err
-	}
-
 	var cnt int
-	for _, table := range tables {
+	for _, table := range st.tableNames {
 		// only Dolt-specific provider has branch support
-		dbStats, err := dSess.StatsProvider().(BranchStatsProvider).GetTableDoltStats(ctx, st.branch, st.dbName, table)
+		dbStats, err := dSess.StatsProvider().(BranchStatsProvider).GetTableDoltStats(ctx, st.branch, st.dbName, st.schemaName, table)
 		if err != nil {
 
 		}
@@ -136,28 +119,10 @@ func (st *StatisticsTable) Partitions(*sql.Context) (sql.PartitionIter, error) {
 // PartitionRows is a sql.Table interface function that gets a row iterator for a partition
 func (st *StatisticsTable) PartitionRows(ctx *sql.Context, _ sql.Partition) (sql.RowIter, error) {
 	dSess := dsess.DSessFromSess(ctx.Session)
-	prov := dSess.Provider()
-
-	var sqlDb sql.Database
-	var err error
-	if st.branch != "" {
-		sqlDb, err = prov.Database(ctx, fmt.Sprintf("%s/%s", st.dbName, st.branch))
-	} else {
-		sqlDb, err = prov.Database(ctx, st.dbName)
-	}
-	if err != nil {
-		return nil, err
-	}
-
-	tables, err := sqlDb.GetTableNames(ctx)
-	if err != nil {
-		return nil, err
-	}
-
 	statsPro := dSess.StatsProvider().(BranchStatsProvider)
 	var dStats []sql.Statistic
-	for _, table := range tables {
-		dbStats, err := statsPro.GetTableDoltStats(ctx, st.branch, st.dbName, table)
+	for _, table := range st.tableNames {
+		dbStats, err := statsPro.GetTableDoltStats(ctx, st.branch, st.dbName, st.schemaName, table)
 		if err != nil {
 			return nil, err
 		}

go/libraries/doltcore/sqle/statsnoms/iter.go

@@ -32,7 +32,7 @@ import (
 
 var ErrIncompatibleVersion = errors.New("client stats version mismatch")
 
-func NewStatsIter(ctx *sql.Context, m prolly.Map) (*statsIter, error) {
+func NewStatsIter(ctx *sql.Context, schemaName string, m prolly.Map) (*statsIter, error) {
 	iter, err := m.IterAll(ctx)
 	if err != nil {
 		return nil, err
@@ -43,11 +43,12 @@ func NewStatsIter(ctx *sql.Context, m prolly.Map) (*statsIter, error) {
 	ns := m.NodeStore()
 
 	return &statsIter{
-		iter:  iter,
-		kb:    keyBuilder,
-		vb:    valueBuilder,
-		ns:    ns,
-		planb: planbuilder.New(ctx, nil, nil, nil),
+		iter:       iter,
+		kb:         keyBuilder,
+		vb:         valueBuilder,
+		ns:         ns,
+		schemaName: schemaName,
+		planb:      planbuilder.New(ctx, nil, nil, nil),
 	}, nil
 }
 
@@ -61,6 +62,7 @@ type statsIter struct {
 	ns           tree.NodeStore
 	planb        *planbuilder.Builder
 	currentQual  string
+	schemaName   string
 	currentTypes []sql.Type
 }
 
@@ -118,7 +120,7 @@ func (s *statsIter) Next(ctx *sql.Context) (sql.Row, error) {
 		typs[i] = strings.TrimSpace(t)
 	}
 
-	qual := sql.NewStatQualifier(dbName, tableName, indexName)
+	qual := sql.NewStatQualifier(dbName, s.schemaName, tableName, indexName)
 	if curQual := qual.String(); !strings.EqualFold(curQual, s.currentQual) {
 		s.currentQual = curQual
 		s.currentTypes, err = parseTypeStrings(typs)

go/libraries/doltcore/sqle/statsnoms/load.go

@@ -39,7 +39,8 @@ import (
 
 func loadStats(ctx *sql.Context, db dsess.SqlDatabase, m prolly.Map) (map[sql.StatQualifier]*statspro.DoltStats, error) {
 	qualToStats := make(map[sql.StatQualifier]*statspro.DoltStats)
-	iter, err := NewStatsIter(ctx, m)
+	schemaName := db.SchemaName()
+	iter, err := NewStatsIter(ctx, schemaName, m)
 	if err != nil {
 		return nil, err
 	}
@@ -72,7 +73,7 @@ func loadStats(ctx *sql.Context, db dsess.SqlDatabase, m prolly.Map) (map[sql.St
 			typs[i] = strings.TrimSpace(t)
 		}
 
-		qual := sql.NewStatQualifier(dbName, tableName, indexName)
+		qual := sql.NewStatQualifier(dbName, schemaName, tableName, indexName)
 		if currentStat.Statistic.Qual.String() != qual.String() {
 			if !currentStat.Statistic.Qual.Empty() {
 				currentStat.Statistic.LowerBnd, currentStat.Tb, err = loadLowerBound(ctx, db, currentStat.Statistic.Qual, len(currentStat.Columns()))

go/libraries/doltcore/sqle/statspro/analyze.go

@@ -100,6 +100,10 @@ func (p *Provider) RefreshTableStatsWithBranch(ctx *sql.Context, table sql.Table
 
 	tableName := strings.ToLower(table.Name())
 	dbName := strings.ToLower(db)
+	var schemaName string
+	if schTab, ok := table.(sql.DatabaseSchemaTable); ok {
+		schemaName = strings.ToLower(schTab.DatabaseSchema().SchemaName())
+	}
 
 	iat, ok := table.(sql.IndexAddressableTable)
 	if !ok {
@@ -146,7 +150,7 @@ func (p *Provider) RefreshTableStatsWithBranch(ctx *sql.Context, table sql.Table
 		ctx.GetLogger().Debugf("statistics refresh: detected table schema change: %s,%s/%s", dbName, table, branch)
 		statDb.SetSchemaHash(branch, tableName, schHash)
 
-		stats, err := p.GetTableDoltStats(ctx, branch, dbName, tableName)
+		stats, err := p.GetTableDoltStats(ctx, branch, dbName, schemaName, tableName)
 		if err != nil {
 			return err
 		}
@@ -163,7 +167,7 @@ func (p *Provider) RefreshTableStatsWithBranch(ctx *sql.Context, table sql.Table
 			cols[i] = strings.TrimPrefix(strings.ToLower(c), tablePrefix)
 		}
 
-		qual := sql.NewStatQualifier(db, table.Name(), strings.ToLower(idx.ID()))
+		qual := sql.NewStatQualifier(db, schemaName, table.Name(), strings.ToLower(idx.ID()))
 		curStat, ok := statDb.GetStat(branch, qual)
 		if !ok {
 			curStat = NewDoltStats()

go/libraries/doltcore/sqle/statspro/auto_refresh.go

@@ -163,13 +163,18 @@ func (p *Provider) checkRefresh(ctx *sql.Context, sqlDb sql.Database, dbName, br
 			return err
 		}
 
+		var schemaName string
+		if schTab, ok := sqlTable.(sql.DatabaseSchemaTable); ok {
+			schemaName = strings.ToLower(schTab.DatabaseSchema().SchemaName())
+		}
+
 		if oldSchHash := statDb.GetSchemaHash(branch, table); oldSchHash.IsEmpty() {
 			statDb.SetSchemaHash(branch, table, schHash)
 		} else if oldSchHash != schHash {
 			ctx.GetLogger().Debugf("statistics refresh: detected table schema change: %s,%s/%s", dbName, table, branch)
 			statDb.SetSchemaHash(branch, table, schHash)
 
-			stats, err := p.GetTableDoltStats(ctx, branch, dbName, table)
+			stats, err := p.GetTableDoltStats(ctx, branch, dbName, schemaName, table)
 			if err != nil {
 				return err
 			}
@@ -191,7 +196,7 @@ func (p *Provider) checkRefresh(ctx *sql.Context, sqlDb sql.Database, dbName, br
 		// collect indexes and ranges to be updated
 		var idxMetas []indexMeta
 		for _, index := range indexes {
-			qual := sql.NewStatQualifier(dbName, table, strings.ToLower(index.ID()))
+			qual := sql.NewStatQualifier(dbName, schemaName, table, strings.ToLower(index.ID()))
 			qualExists[qual] = true
 			curStat, ok := statDb.GetStat(branch, qual)
 			if !ok {

go/libraries/doltcore/sqle/statspro/stats_provider.go

@@ -169,11 +169,15 @@ func (p *Provider) GetTableStats(ctx *sql.Context, db string, table sql.Table) (
 		return nil, nil
 	}
 
-	// TODO: schema name
-	return p.GetTableDoltStats(ctx, branch, db, table.Name())
+	var schemaName string
+	if schTab, ok := table.(sql.DatabaseSchemaTable); ok {
+		schemaName = strings.ToLower(schTab.DatabaseSchema().SchemaName())
+	}
+
+	return p.GetTableDoltStats(ctx, branch, db, schemaName, table.Name())
 }
 
-func (p *Provider) GetTableDoltStats(ctx *sql.Context, branch, db, table string) ([]sql.Statistic, error) {
+func (p *Provider) GetTableDoltStats(ctx *sql.Context, branch, db, schema, table string) ([]sql.Statistic, error) {
 	statDb, ok := p.getStatDb(db)
 	if !ok || statDb == nil {
 		return nil, nil
@@ -190,7 +194,7 @@ func (p *Provider) GetTableDoltStats(ctx *sql.Context, branch, db, table string)
 
 	var ret []sql.Statistic
 	for _, qual := range statDb.ListStatQuals(branch) {
-		if strings.EqualFold(db, qual.Database) && strings.EqualFold(table, qual.Tab) {
+		if strings.EqualFold(db, qual.Database) && strings.EqualFold(schema, qual.Sch) && strings.EqualFold(table, qual.Tab) {
 			stat, _ := statDb.GetStat(branch, qual)
 			ret = append(ret, stat)
 		}
@@ -333,8 +337,12 @@ func (p *Provider) RowCount(ctx *sql.Context, db string, table sql.Table) (uint6
 		return 0, err
 	}
 
-	// TODO: schema name
-	priStats, ok := statDb.GetStat(branch, sql.NewStatQualifier(db, table.Name(), "primary"))
+	var schemaName string
+	if schTab, ok := table.(sql.DatabaseSchemaTable); ok {
+		schemaName = strings.ToLower(schTab.DatabaseSchema().SchemaName())
+	}
+
+	priStats, ok := statDb.GetStat(branch, sql.NewStatQualifier(db, schemaName, table.Name(), "primary"))
 	if !ok {
 		return 0, nil
 	}
@@ -354,8 +362,12 @@ func (p *Provider) DataLength(ctx *sql.Context, db string, table sql.Table) (uin
 		return 0, err
 	}
 
-	// TODO: schema name
-	priStats, ok := statDb.GetStat(branch, sql.NewStatQualifier(db, table.Name(), "primary"))
+	var schemaName string
+	if schTab, ok := table.(sql.DatabaseSchemaTable); ok {
+		schemaName = strings.ToLower(schTab.DatabaseSchema().SchemaName())
+	}
+
+	priStats, ok := statDb.GetStat(branch, sql.NewStatQualifier(db, schemaName, table.Name(), "primary"))
 	if !ok {
 		return 0, nil
 	}
@@ -404,7 +416,7 @@ func (p *Provider) Prune(ctx *sql.Context) error {
 					}
 					defer p.UnlockTable(branch, dbName, t)
 
-					tableStats, err := p.GetTableDoltStats(ctx, branch, dbName, t)
+					tableStats, err := p.GetTableDoltStats(ctx, branch, dbName, sqlDb.SchemaName(), t)
 					if err != nil {
 						return err
 					}