initial impl of collStats for SQLite

FerretDB · AlekSi · Sep 4, 2023 · Aug 31, 2023 · Aug 31, 2023 · Sep 1, 2023
commit 37ba896a8844e810e5be868a37d4a77204627317
@@ -23,53 +23,45 @@ import (
 
 	"github.com/FerretDB/FerretDB/integration/setup"
 	"github.com/FerretDB/FerretDB/integration/shareddata"
-	"github.com/FerretDB/FerretDB/internal/util/testutil/testtb"
 )
 
 func TestAggregateCompatCollStats(tt *testing.T) {
 	tt.Parallel()
 
 	for name, tc := range map[string]struct {
-		skip           string                   // skip test for all handlers, must have issue number mentioned
-		collStats      bson.D                   // required
-		resultType     compatTestCaseResultType // defaults to nonEmptyResult
-		failsForSQLite string                   // non-empty value expects test to fail for SQLite backend
+		skip       string                   // skip test for all handlers, must have issue number mentioned
+		collStats  bson.D                   // required
+		resultType compatTestCaseResultType // defaults to nonEmptyResult
 	}{
 		"NilCollStats": {
 			collStats:  nil,
 			resultType: emptyResult,
 		},
 		"EmptyCollStats": {
-			collStats:      bson.D{},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{},
 		},
 		"Count": {
-			collStats:      bson.D{{"count", bson.D{}}},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{{"count", bson.D{}}},
 		},
 		"StorageStats": {
-			collStats:      bson.D{{"storageStats", bson.D{}}},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{{"storageStats", bson.D{}}},
 		},
 		"StorageStatsWithScale": {
-			collStats:      bson.D{{"storageStats", bson.D{{"scale", 1000}}}},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{{"storageStats", bson.D{{"scale", 1000}}}},
 		},
 		"StorageStatsNegativeScale": {
 			collStats:  bson.D{{"storageStats", bson.D{{"scale", -1000}}}},
 			resultType: emptyResult,
 		},
 		"StorageStatsFloatScale": {
-			collStats:      bson.D{{"storageStats", bson.D{{"scale", 42.42}}}},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{{"storageStats", bson.D{{"scale", 42.42}}}},
 		},
 		"StorageStatsInvalidScale": {
 			collStats:  bson.D{{"storageStats", bson.D{{"scale", "invalid"}}}},
 			resultType: emptyResult,
 		},
 		"CountAndStorageStats": {
-			collStats:      bson.D{{"count", bson.D{}}, {"storageStats", bson.D{}}},
-			failsForSQLite: "https://github.com/FerretDB/FerretDB/issues/3259",
+			collStats: bson.D{{"count", bson.D{}}, {"storageStats", bson.D{}}},
 		},
 	} {
 		name, tc := name, tc
@@ -92,13 +84,8 @@ func TestAggregateCompatCollStats(tt *testing.T) {
 			for i := range targetCollections {
 				targetCollection := targetCollections[i]
 				compatCollection := compatCollections[i]
-				tt.Run(targetCollection.Name(), func(tt *testing.T) {
-					tt.Helper()
-
-					var t testtb.TB = tt
-					if tc.failsForSQLite != "" {
-						t = setup.FailsForSQLite(tt, tc.failsForSQLite)
-					}
+				tt.Run(targetCollection.Name(), func(t *testing.T) {
+					t.Helper()
 
 					command := bson.A{bson.D{{"$collStats", tc.collStats}}}
 
@@ -145,11 +132,6 @@ func TestAggregateCompatCollStats(tt *testing.T) {
 				})
 			}
 
-			// TODO https://github.com/FerretDB/FerretDB/issues/3259
-			if setup.IsSQLite(tt) {
-				return
-			}
-
 			switch tc.resultType {
 			case nonEmptyResult:
 				assert.True(tt, nonEmptyResults, "expected non-empty results (some documents should be modified)")

@@ -62,11 +62,10 @@ func TestCommandsAdministrationCompatCollStatsWithScale(t *testing.T) {
 	} {
 		name, tc := name, tc
 
-		t.Run(name, func(tt *testing.T) {
-			tt.Helper()
+		t.Run(name, func(t *testing.T) {
+			t.Helper()
 
-			tt.Parallel()
-			t := setup.FailsForSQLite(tt, "https://github.com/FerretDB/FerretDB/issues/3259")
+			t.Parallel()
 
 			var targetRes bson.D
 			targetCommand := bson.D{{"collStats", targetCollection.Name()}, {"scale", tc.scale}}

@@ -675,11 +675,9 @@ func TestCommandsAdministrationBuildInfoFerretdbExtensions(t *testing.T) {
 	assert.NotEmpty(t, aggregationStagesArray)
 }
 
-func TestCommandsAdministrationCollStatsEmpty(tt *testing.T) {
-	tt.Parallel()
-	ctx, collection := setup.Setup(tt)
-
-	t := setup.FailsForSQLite(tt, "https://github.com/FerretDB/FerretDB/issues/3259")
+func TestCommandsAdministrationCollStatsEmpty(t *testing.T) {
+	t.Parallel()
+	ctx, collection := setup.Setup(t)
 
 	var actual bson.D
 	command := bson.D{{"collStats", collection.Name()}}
@@ -699,10 +697,9 @@ func TestCommandsAdministrationCollStatsEmpty(tt *testing.T) {
 	assert.Equal(t, float64(1), must.NotFail(doc.Get("ok")))
 }
 
-func TestCommandsAdministrationCollStats(tt *testing.T) {
-	tt.Parallel()
+func TestCommandsAdministrationCollStats(t *testing.T) {
+	t.Parallel()
 
-	t := setup.FailsForSQLite(tt, "https://github.com/FerretDB/FerretDB/issues/3259")
 	ctx, collection := setup.Setup(t, shareddata.DocumentsStrings)
 
 	var actual bson.D
@@ -732,10 +729,8 @@ func TestCommandsAdministrationCollStats(tt *testing.T) {
 	assert.InDelta(t, 32_000, must.NotFail(doc.Get("totalSize")), 30_000)
 }
 
-func TestCommandsAdministrationCollStatsWithScale(tt *testing.T) {
-	tt.Parallel()
-
-	t := setup.FailsForSQLite(tt, "https://github.com/FerretDB/FerretDB/issues/3259")
+func TestCommandsAdministrationCollStatsWithScale(t *testing.T) {
+	t.Parallel()
 
 	ctx, collection := setup.Setup(t, shareddata.DocumentsStrings)
 

@@ -187,7 +187,9 @@ func (dbc *databaseContract) RenameCollection(ctx context.Context, params *Renam
 }
 
 // StatsParams represents the parameters of Database.Stats method.
-type StatsParams struct{}
+type StatsParams struct {
+	Collection string
+}
 
 // StatsResult represents the results of Database.Stats method.
 //
@@ -202,6 +204,8 @@ type StatsResult struct {
 }
 
 // Stats returns statistics about the database.
+// If collection is specified in *StatsParams.Collection, it returns statistics of only that collection.
+// Otherwise, it returns statistics of the entire database.
 //
 // Database may not exist; that's not an error.
 func (dbc *databaseContract) Stats(ctx context.Context, params *StatsParams) (*StatsResult, error) {

@@ -16,6 +16,8 @@ package sqlite
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"github.com/FerretDB/FerretDB/internal/backends"
 	"github.com/FerretDB/FerretDB/internal/backends/sqlite/metadata"
@@ -102,8 +104,106 @@ func (db *database) RenameCollection(ctx context.Context, params *backends.Renam
 }
 
 // Stats implements backends.Database interface.
+//
+// If the database does not exist, it returns *backends.DBStatsResult filled with zeros for all the fields.
 func (db *database) Stats(ctx context.Context, params *backends.StatsParams) (*backends.StatsResult, error) {
-	panic("not implemented")
+	stats := new(backends.StatsResult)
+
+	d := db.r.DatabaseGetExisting(ctx, db.name)
+	if d == nil {
+		return stats, nil
+	}
+
+	var list []*metadata.Collection
+	var err error
+
+	singleCollectionStats := params.Collection != ""
+
+	if singleCollectionStats {
+		c := db.r.CollectionGet(ctx, db.name, params.Collection)
+		if c == nil {
+			return stats, nil
+		}
+
+		list = append(list, c)
+	} else {
+		if list, err = db.r.CollectionList(ctx, db.name); err != nil {
+			return nil, lazyerrors.Error(err)
+		}
+	}
+
+	stats.CountCollections = int64(len(list))
+
+	// Call ANALYZE to update statistics of tables and indexes,
+	// see https://www.sqlite.org/lang_analyze.html.
+	q := `ANALYZE`
+	if _, err = d.ExecContext(ctx, q); err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	placeholders := make([]string, len(list))
+	args := make([]any, len(list))
+
+	for i, c := range list {
+		placeholders[i] = "?"
+		args[i] = c.TableName
+	}
+
+	// Use number of cells to approximate total row count,
+	// see https://www.sqlite.org/fileformat.html.
+	q = fmt.Sprintf(`
+		SELECT
+		    SUM(pgsize) AS SizeTables,
+		    SUM(ncell)  AS CountCells
+		FROM dbstat
+		WHERE name IN (%s) AND aggregate = TRUE`,
+		strings.Join(placeholders, ", "),
+	)
+
+	if err = d.QueryRowContext(ctx, q, args...).Scan(
+		&stats.SizeCollections,
+		&stats.CountObjects,
+	); err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	// Use sqlite_schema table to get indexes of each tables,
+	// see https://www.sqlite.org/schematab.html.
+	q = fmt.Sprintf(`
+		SELECT
+			COUNT(s.name)             AS CountIndexes,
+			COALESCE(SUM(d.pgsize),0) AS SizeIndexes
+		FROM sqlite_schema AS s
+			LEFT JOIN dbstat AS d ON d.name = s.tbl_name
+		WHERE s.type = 'index' AND s.tbl_name IN (%s)`,
+		strings.Join(placeholders, ", "),
+	)
+
+	if err = d.QueryRowContext(ctx, q, args...).Scan(
+		&stats.CountIndexes,
+		&stats.SizeIndexes,
+	); err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	if singleCollectionStats {
+		stats.SizeTotal = stats.SizeCollections + stats.SizeIndexes
+		return stats, nil
+	}
+
+	// Total size is the disk space used by the database,
+	// see https://www.sqlite.org/dbstat.html.
+	q = `
+		SELECT
+			SUM(pgsize)
+		FROM dbstat WHERE aggregate = TRUE`
+
+	err = d.QueryRowContext(ctx, q).Scan(&stats.SizeTotal)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return stats, nil
 }
 
 // check interfaces

@@ -0,0 +1,83 @@
+// Copyright 2021 FerretDB Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sqlite
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/FerretDB/FerretDB/internal/backends"
+	"github.com/FerretDB/FerretDB/internal/backends/sqlite/metadata"
+	"github.com/FerretDB/FerretDB/internal/util/testutil"
+)
+
+func TestStats(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Ctx(t)
+
+	r, err := metadata.NewRegistry("file:./?mode=memory", testutil.Logger(t))
+	require.NoError(t, err)
+	t.Cleanup(r.Close)
+
+	dbName := t.Name()
+	db := newDatabase(r, dbName)
+
+	t.Cleanup(func() {
+		db.Close()
+	})
+
+	t.Run("NonExistingDatabase", func(t *testing.T) {
+		var res *backends.StatsResult
+		res, err = db.Stats(ctx, new(backends.StatsParams))
+		require.NoError(t, err)
+		require.Equal(t, new(backends.StatsResult), res)
+	})
+
+	collectionOne := "collectionOne"
+	err = db.CreateCollection(ctx, &backends.CreateCollectionParams{Name: collectionOne})
+	require.NoError(t, err)
+	require.NotNil(t, db)
+
+	collectionTwo := "collectionTwo"
+	err = db.CreateCollection(ctx, &backends.CreateCollectionParams{Name: collectionTwo})
+	require.NoError(t, err)
+	require.NotNil(t, db)
+
+	t.Cleanup(func() {
+		r.DatabaseDrop(ctx, dbName)
+	})
+
+	var dbStatsRes *backends.StatsResult
+	t.Run("EmptyCollection", func(t *testing.T) {
+		dbStatsRes, err = db.Stats(ctx, new(backends.StatsParams))
+		require.NoError(t, err)
+		require.NotZero(t, dbStatsRes.SizeTotal)
+		require.NotZero(t, dbStatsRes.CountCollections)
+		require.NotZero(t, dbStatsRes.SizeCollections)
+		require.Zero(t, dbStatsRes.CountObjects)
+	})
+
+	t.Run("CollectionOne", func(t *testing.T) {
+		res, err := db.Stats(ctx, &backends.StatsParams{Collection: collectionOne})
+		require.NoError(t, err)
+		require.NotZero(t, res.SizeTotal)
+		require.Less(t, res.SizeTotal, dbStatsRes.SizeTotal)
+		require.Equal(t, res.CountCollections, int64(1))
+		require.NotZero(t, res.SizeCollections)
+		require.Less(t, res.SizeCollections, dbStatsRes.SizeCollections)
+		require.Zero(t, res.CountObjects)
+	})
+}
@@ -93,7 +93,7 @@ func (h *Handler) MsgCollStats(ctx context.Context, msg *wire.OpMsg) (*wire.OpMs
 	}
 
 	pairs = append(pairs,
-		"storageSize", stats.SizeTotal/scale,
+		"storageSize", stats.SizeCollection/scale,
 		"nindexes", stats.CountIndexes,
 		"totalIndexSize", stats.SizeIndexes/scale,
 		"totalSize", stats.SizeTotal/scale,