Support dbStats (#232)

Closes #162.
FerretDB · Jan 6, 2022 · 3d359b9 · 3d359b9
1 parent b8683f3
commit 3d359b9
Show file tree

Hide file tree

Showing 4 changed files with 167 additions and 0 deletions.
diff --git a/internal/handlers/handler.go b/internal/handlers/handler.go
@@ -160,6 +160,8 @@ func (h *Handler) handleOpMsg(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg
 		return h.shared.MsgCollStats(ctx, msg)
 	case "create":
 		return h.shared.MsgCreate(ctx, msg)
+	case "dbstats":
+		return h.shared.MsgDBStats(ctx, msg)
 	case "drop":
 		return h.shared.MsgDrop(ctx, msg)
 	case "dropdatabase":

diff --git a/internal/handlers/handler_test.go b/internal/handlers/handler_test.go
@@ -503,6 +503,66 @@ func TestReadOnlyHandlers(t *testing.T) {
 			),
 		},
 
+		"DBStats": {
+			req: types.MustMakeDocument(
+				"dbstats", int32(1),
+			),
+			reqSetDB: true,
+			resp: types.MustMakeDocument(
+				"db", "monila",
+				"collections", int32(14),
+				"views", int32(0),
+				"objects", int32(30224),
+				"avgObjSize", 437.7342509264161,
+				"dataSize", 1.323008e+07,
+				"indexes", int32(0),
+				"indexSize", float64(0),
+				"totalSize", 1.3615104e+07,
+				"scaleFactor", float64(1),
+				"ok", float64(1),
+			),
+			compareFunc: func(t testing.TB, req types.Document, actual, expected types.CompositeType) {
+				db, err := req.Get("$db")
+				require.NoError(t, err)
+				if db.(string) == "monila" {
+					testutil.CompareAndSetByPathNum(t, expected, actual, 20, "avgObjSize")
+					testutil.CompareAndSetByPathNum(t, expected, actual, 400_000, "dataSize")
+					testutil.CompareAndSetByPathNum(t, expected, actual, 400_000, "totalSize")
+					assert.Equal(t, expected, actual)
+				}
+			},
+		},
+		"DBStatsWithScale": {
+			req: types.MustMakeDocument(
+				"dbstats", int32(1),
+				"scale", float64(1_000),
+			),
+			reqSetDB: true,
+			resp: types.MustMakeDocument(
+				"db", "monila",
+				"collections", int32(14),
+				"views", int32(0),
+				"objects", int32(30224),
+				"avgObjSize", 437.7342509264161,
+				"dataSize", 13_230.08,
+				"indexes", int32(0),
+				"indexSize", float64(0),
+				"totalSize", 13_615.104,
+				"scaleFactor", float64(1_000),
+				"ok", float64(1),
+			),
+			compareFunc: func(t testing.TB, req types.Document, actual, expected types.CompositeType) {
+				db, err := req.Get("$db")
+				require.NoError(t, err)
+				if db.(string) == "monila" {
+					testutil.CompareAndSetByPathNum(t, expected, actual, 20, "avgObjSize")
+					testutil.CompareAndSetByPathNum(t, expected, actual, 400, "dataSize")
+					testutil.CompareAndSetByPathNum(t, expected, actual, 400, "totalSize")
+					assert.Equal(t, expected, actual)
+				}
+			},
+		},
+
 		"FindProjectionActorsFirstAndLastName": {
 			req: types.MustMakeDocument(
 				"find", "actor",

diff --git a/internal/handlers/shared/msg_dbstats.go b/internal/handlers/shared/msg_dbstats.go
@@ -0,0 +1,65 @@
+// 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 shared
+
+import (
+	"context"
+
+	"github.com/FerretDB/FerretDB/internal/types"
+	"github.com/FerretDB/FerretDB/internal/util/lazyerrors"
+	"github.com/FerretDB/FerretDB/internal/wire"
+)
+
+func (h *Handler) MsgDBStats(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, error) {
+	document, err := msg.Document()
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	m := document.Map()
+	db := m["$db"].(string)
+	scale, ok := m["scale"].(float64)
+	if !ok {
+		scale = 1
+	}
+
+	stats, err := h.pgPool.DBStats(ctx, db)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	var reply wire.OpMsg
+	err = reply.SetSections(wire.OpMsgSection{
+		Documents: []types.Document{types.MustMakeDocument(
+			"db", db,
+			"collections", stats.CountTables,
+			// TODO https://github.com/FerretDB/FerretDB/issues/176
+			"views", int32(0),
+			"objects", stats.CountRows,
+			"avgObjSize", float64(stats.SizeSchema)/float64(stats.CountRows),
+			"dataSize", float64(stats.SizeSchema)/scale,
+			"indexes", stats.CountIndexes,
+			"indexSize", float64(stats.SizeIndexes)/scale,
+			"totalSize", float64(stats.SizeTotal)/scale,
+			"scaleFactor", scale,
+			"ok", float64(1),
+		)},
+	})
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return &reply, nil
+}
diff --git a/internal/pg/pool.go b/internal/pg/pool.go
@@ -58,6 +58,17 @@ type TableStats struct {
 	Rows        int64
 }
 
+// DBStats describes some statistics for a database.
+type DBStats struct {
+	Name         string
+	CountTables  int32
+	CountRows    int32
+	SizeTotal    int64
+	SizeIndexes  int64
+	SizeSchema   int64
+	CountIndexes int32
+}
+
 // NewPool returns a pgxpool, a concurrency-safe connection pool for pgx.
 func NewPool(connString string, logger *zap.Logger, lazy bool) (*Pool, error) {
 	config, err := pgxpool.ParseConfig(connString)
@@ -301,3 +312,32 @@ func (pgPool *Pool) TableStats(ctx context.Context, db, table string) (*TableSta
 
 	return res, nil
 }
+
+// DBStats returns a set of statistics for a database.
+func (pgPool *Pool) DBStats(ctx context.Context, db string) (*DBStats, error) {
+	res := new(DBStats)
+	sql := `
+    SELECT COUNT(distinct t.table_name)                                                             AS CountTables,
+           COALESCE(SUM(s.n_live_tup), 0)                                                           AS CountRows,
+           COALESCE(SUM(pg_total_relation_size('"'||t.table_schema||'"."'||t.table_name||'"')), 0)  AS SizeTotal,
+           COALESCE(SUM(pg_indexes_size('"'||t.table_schema||'"."'||t.table_name||'"')), 0)         AS SizeIndexes,
+           COALESCE(SUM(pg_relation_size('"'||t.table_schema||'"."'||t.table_name||'"')), 0)        AS SizeSchema,
+           COUNT(distinct i.indexname)                                                              AS CountIndexes
+      FROM information_schema.tables AS t
+      LEFT OUTER
+      JOIN pg_stat_user_tables       AS s ON s.schemaname = t.table_schema
+                                         AND s.relname = t.table_name
+      LEFT OUTER
+      JOIN pg_indexes                AS i ON i.schemaname = t.table_schema
+                                         AND i.tablename = t.table_name
+     WHERE t.table_schema = $1`
+
+	res.Name = db
+	err := pgPool.QueryRow(ctx, sql, db).
+		Scan(&res.CountTables, &res.CountRows, &res.SizeTotal, &res.SizeIndexes, &res.SizeSchema, &res.CountIndexes)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return res, nil
+}