Support collStats command.

internal/handlers/handler.go

@@ -156,6 +156,8 @@ func (h *Handler) handleOpMsg(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg
 	switch cmd {
 	case "buildinfo":
 		return h.shared.MsgBuildInfo(ctx, msg)
+	case "collstats":
+		return h.shared.MsgCollStats(ctx, msg)
 	case "create":
 		return h.shared.MsgCreate(ctx, msg)
 	case "drop":

internal/handlers/handler_test.go

@@ -418,7 +418,7 @@ func TestReadOnlyHandlers(t *testing.T) {
 		req         types.Document
 		reqSetDB    bool
 		resp        types.Document
-		compareFunc func(t testing.TB, actual, expected types.CompositeType)
+		compareFunc func(t testing.TB, req types.Document, actual, expected types.CompositeType)
 	}
 
 	hostname, err := os.Hostname()
@@ -438,6 +438,29 @@ func TestReadOnlyHandlers(t *testing.T) {
 				"ok", float64(1),
 			),
 		},
+		"CollStats": {
+			req: types.MustMakeDocument(
+				"collstats", "film",
+			),
+			reqSetDB: true,
+			resp: types.MustMakeDocument(
+				"ns", "manila.film",
+				"count", int64(1_000),
+				"size", int64(704512),
+				"storageSize", int64(450560),
+				"totalIndexSize", int64(221184),
+				"totalSize", int64(704512),
+				"scaleFactor", int64(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) == "manila" {
+					assert.Equal(t, expected, actual)
+				}
+			},
+		},
 		"CountAllActors": {
 			req: types.MustMakeDocument(
 				"count", "actor",
@@ -501,7 +524,7 @@ func TestReadOnlyHandlers(t *testing.T) {
 				"readOnly", false,
 				"ok", float64(1),
 			),
-			compareFunc: func(t testing.TB, actual, expected types.CompositeType) {
+			compareFunc: func(t testing.TB, _ types.Document, actual, expected types.CompositeType) {
 				testutil.CompareAndSetByPathTime(t, expected, actual, time.Second, "localTime")
 				assert.Equal(t, expected, actual)
 			},
@@ -522,7 +545,7 @@ func TestReadOnlyHandlers(t *testing.T) {
 				"readOnly", false,
 				"ok", float64(1),
 			),
-			compareFunc: func(t testing.TB, actual, expected types.CompositeType) {
+			compareFunc: func(t testing.TB, _ types.Document, actual, expected types.CompositeType) {
 				testutil.CompareAndSetByPathTime(t, expected, actual, time.Second, "localTime")
 				assert.Equal(t, expected, actual)
 			},
@@ -546,7 +569,7 @@ func TestReadOnlyHandlers(t *testing.T) {
 				),
 				"ok", float64(1),
 			),
-			compareFunc: func(t testing.TB, actual, expected types.CompositeType) {
+			compareFunc: func(t testing.TB, _ types.Document, actual, expected types.CompositeType) {
 				testutil.CompareAndSetByPathTime(t, expected, actual, time.Second, "system", "currentTime")
 				assert.Equal(t, expected, actual)
 			},
@@ -580,7 +603,7 @@ func TestReadOnlyHandlers(t *testing.T) {
 					if tc.compareFunc == nil {
 						assert.Equal(t, tc.resp, actual)
 					} else {
-						tc.compareFunc(t, tc.resp, actual)
+						tc.compareFunc(t, tc.req, tc.resp, actual)
 					}
 				})
 			}

internal/handlers/shared/msg_collstats.go

@@ -0,0 +1,62 @@
+// 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"
+)
+
+// MsgCollStats returns a set of statistics for a collection.
+func (h *Handler) MsgCollStats(ctx context.Context, msg *wire.OpMsg) (*wire.OpMsg, error) {
+	document, err := msg.Document()
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	m := document.Map()
+	collection := m[document.Command()].(string)
+	db, ok := m["$db"].(string)
+	if !ok {
+		return nil, lazyerrors.New("no db")
+	}
+
+	stats, err := h.pgPool.TableStats(ctx, db, collection)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	var reply wire.OpMsg
+	err = reply.SetSections(wire.OpMsgSection{
+		Documents: []types.Document{types.MustMakeDocument(
+			"ns", db+"."+collection,
+			"count", stats.Rows,
+			"size", stats.SizeTotal,
+			"storageSize", stats.SizeTable,
+			"totalIndexSize", stats.SizeIndexes,
+			"totalSize", stats.SizeTotal,
+			"scaleFactor", int64(1),
+			"ok", float64(1),
+		)},
+	})
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return &reply, nil
+}

internal/pg/pool.go

@@ -48,6 +48,16 @@ type Pool struct {
 	*pgxpool.Pool
 }
 
+// TableStats describes some statistics for a table.
+type TableStats struct {
+	Table       string
+	TableType   string
+	SizeTotal   int64
+	SizeIndexes int64
+	SizeTable   int64
+	Rows        int64
+}
+
 // 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)
@@ -266,3 +276,28 @@ func (pgPool *Pool) DropTable(ctx context.Context, db, collection string) error
 
 	return err
 }
+
+// TableStats returns a set of statistics for a table.
+func (pgPool *Pool) TableStats(ctx context.Context, db, table string) (*TableStats, error) {
+	res := new(TableStats)
+	sql := `
+    SELECT table_name, table_type,
+           pg_total_relation_size('"'||t.table_schema||'"."'||t.table_name||'"'),
+           pg_indexes_size('"'||t.table_schema||'"."'||t.table_name||'"'),
+           pg_relation_size('"'||t.table_schema||'"."'||t.table_name||'"'),
+           COALESCE(s.n_live_tup, 0)
+      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
+     WHERE t.table_schema = $1
+       AND t.table_name = $2`
+
+	err := pgPool.QueryRow(ctx, sql, db, table).
+		Scan(&res.Table, &res.TableType, &res.SizeTotal, &res.SizeIndexes, &res.SizeTable, &res.Rows)
+	if err != nil {
+		return nil, lazyerrors.Error(err)
+	}
+
+	return res, nil
+}