Add early stats for the per-domain level

Fixes https://github.com/turt2live/matrix-media-repo/issues/169

api/custom/usage.go

+package custom
+
+import (
+	"net/http"
+
+	"github.com/gorilla/mux"
+	"github.com/sirupsen/logrus"
+	"github.com/turt2live/matrix-media-repo/api"
+	"github.com/turt2live/matrix-media-repo/storage"
+)
+
+type UsageInfo struct {
+	Total      int64 `json:"total"`
+	Media      int64 `json:"media"`
+	Thumbnails int64 `json:"thumbnails"`
+}
+
+type DomainUsageResponse struct {
+	NumUsers  int       `json:"user_count"`
+	UserIDs   []string  `json:"known_user_ids"`
+	RawBytes  UsageInfo `json:"raw_bytes"`
+	RawCounts UsageInfo `json:"raw_counts"`
+}
+
+func GetDomainUsage(r *http.Request, log *logrus.Entry, user api.UserInfo) interface{} {
+	// TODO: Auth check on host (allow local admin to query their own domain)
+
+	params := mux.Vars(r)
+
+	serverName := params["serverName"]
+
+	log = log.WithFields(logrus.Fields{
+		"serverName": serverName,
+	})
+
+	db := storage.GetDatabase().GetMetadataStore(r.Context(), log)
+
+	userIds, err := db.GetUsersForServer(serverName)
+	if err != nil {
+		log.Error(err)
+		return api.InternalServerError("Failed to get users belonging to server")
+	}
+
+	mediaBytes, thumbBytes, err := db.GetByteUsageForServer(serverName)
+	if err != nil {
+		log.Error(err)
+		return api.InternalServerError("Failed to get byte usage for server")
+	}
+
+	mediaCount, thumbCount, err := db.GetCountUsageForServer(serverName)
+	if err != nil {
+		log.Error(err)
+		return api.InternalServerError("Failed to get count usage for server")
+	}
+
+	return &DomainUsageResponse{
+		NumUsers: len(userIds),
+		UserIDs:  userIds,
+		RawBytes: UsageInfo{
+			Total:      mediaBytes + thumbBytes,
+			Media:      mediaBytes,
+			Thumbnails: thumbBytes,
+		},
+		RawCounts: UsageInfo{
+			Total:      mediaCount + thumbCount,
+			Media:      mediaCount,
+			Thumbnails: thumbCount,
+		},
+	}
+}

api/webserver/webserver.go

@@ -42,6 +42,7 @@ func Init() {
 	dsTransferHandler := handler{api.RepoAdminRoute(custom.MigrateBetweenDatastores), "datastore_transfer", counter, false}
 	fedTestHandler := handler{api.RepoAdminRoute(custom.GetFederationInfo), "federation_test", counter, false}
 	healthzHandler := handler{api.AccessTokenOptionalRoute(custom.GetHealthz), "healthz", counter, true}
+	domainUsageHandler := handler{api.AccessTokenOptionalRoute(custom.GetDomainUsage), "domain_usage", counter, false}
 
 	routes := make(map[string]route)
 	versions := []string{"r0", "v1", "unstable"} // r0 is typically clients and v1 is typically servers. v1 is deprecated.
@@ -64,6 +65,7 @@ func Init() {
 		routes["/_matrix/media/"+version+"/admin/datastores"] = route{"GET", datastoreListHandler}
 		routes["/_matrix/media/"+version+"/admin/datastores/{sourceDsId:[^/]+}/transfer_to/{targetDsId:[^/]+}"] = route{"POST", dsTransferHandler}
 		routes["/_matrix/media/"+version+"/admin/federation/test/{serverName:[a-zA-Z0-9.:\\-_]+}"] = route{"GET", fedTestHandler}
+		routes["/_matrix/media/"+version+"/admin/usage/{serverName:[a-zA-Z0-9.:\\-_]+}"] = route{"GET", domainUsageHandler}
 
 		// Routes that we should handle but aren't in the media namespace (synapse compat)
 		routes["/_matrix/client/"+version+"/admin/purge_media_cache"] = route{"POST", purgeHandler}

migrations/9_origin_indexes_down.sql

+DROP INDEX IF EXISTS idx_origin_media;
+DROP INDEX IF EXISTS idx_origin_thumbnails;
+DROP INDEX IF EXISTS idx_origin_user_id_media;

migrations/9_origin_indexes_up.sql

+CREATE INDEX IF NOT EXISTS idx_origin_media ON media(origin);
+CREATE INDEX IF NOT EXISTS idx_origin_thumbnails ON thumbnails(origin);
+CREATE INDEX IF NOT EXISTS idx_origin_user_id_media ON media(origin, user_id);

storage/stores/metadata_store.go

@@ -18,6 +18,9 @@ const selectMediaLastAccessedBeforeInDatastore = "SELECT m.sha256_hash, m.size_b
 const selectThumbnailsLastAccessedBeforeInDatastore = "SELECT m.sha256_hash, m.size_bytes, m.datastore_id, m.location, m.creation_ts, a.last_access_ts FROM thumbnails AS m JOIN last_access AS a ON m.sha256_hash = a.sha256_hash WHERE a.last_access_ts < $1 AND m.datastore_id = $2"
 const changeDatastoreOfMediaHash = "UPDATE media SET datastore_id = $1, location = $2 WHERE sha256_hash = $3"
 const changeDatastoreOfThumbnailHash = "UPDATE thumbnails SET datastore_id = $1, location = $2 WHERE sha256_hash = $3"
+const selectUploadCountsForServer = "SELECT COALESCE((SELECT COUNT(origin) FROM media WHERE origin = $1), 0) AS media, COALESCE((SELECT COUNT(origin) FROM thumbnails WHERE origin = $1), 0) AS thumbnails"
+const selectUploadSizesForServer = "SELECT COALESCE((SELECT SUM(size_bytes) FROM media WHERE origin = $1), 0) AS media, COALESCE((SELECT SUM(size_bytes) FROM thumbnails WHERE origin = $1), 0) AS thumbnails"
+const selectUsersForServer = "SELECT DISTINCT user_id FROM media WHERE origin = $1 AND user_id IS NOT NULL AND LENGTH(user_id) > 0"
 
 type metadataStoreStatements struct {
 	upsertLastAccessed                            *sql.Stmt
@@ -26,6 +29,9 @@ type metadataStoreStatements struct {
 	selectThumbnailsLastAccessedBeforeInDatastore *sql.Stmt
 	changeDatastoreOfMediaHash                    *sql.Stmt
 	changeDatastoreOfThumbnailHash                *sql.Stmt
+	selectUploadCountsForServer                   *sql.Stmt
+	selectUploadSizesForServer                    *sql.Stmt
+	selectUsersForServer                          *sql.Stmt
 }
 
 type MetadataStoreFactory struct {
@@ -64,6 +70,15 @@ func InitMetadataStore(sqlDb *sql.DB) (*MetadataStoreFactory, error) {
 	if store.stmts.changeDatastoreOfThumbnailHash, err = store.sqlDb.Prepare(changeDatastoreOfThumbnailHash); err != nil {
 		return nil, err
 	}
+	if store.stmts.selectUsersForServer, err = store.sqlDb.Prepare(selectUsersForServer); err != nil {
+		return nil, err
+	}
+	if store.stmts.selectUploadSizesForServer, err = store.sqlDb.Prepare(selectUploadSizesForServer); err != nil {
+		return nil, err
+	}
+	if store.stmts.selectUploadCountsForServer, err = store.sqlDb.Prepare(selectUploadCountsForServer); err != nil {
+		return nil, err
+	}
 
 	return &store, nil
 }
@@ -151,3 +166,48 @@ func (s *MetadataStore) GetOldThumbnailsInDatastore(datastoreId string, beforeTs
 
 	return results, nil
 }
+
+func (s *MetadataStore) GetUsersForServer(serverName string) ([]string, error) {
+	rows, err := s.statements.selectUsersForServer.QueryContext(s.ctx, serverName)
+	if err != nil {
+		return nil, err
+	}
+
+	results := make([]string, 0)
+	for rows.Next() {
+		v := ""
+		err = rows.Scan(&v)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, v)
+	}
+
+	return results, nil
+}
+
+func (s *MetadataStore) GetByteUsageForServer(serverName string) (int64, int64, error) {
+	row := s.statements.selectUploadSizesForServer.QueryRowContext(s.ctx, serverName)
+
+	media := int64(0)
+	thumbs := int64(0)
+	err := row.Scan(&media, &thumbs)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return media, thumbs, nil
+}
+
+func (s *MetadataStore) GetCountUsageForServer(serverName string) (int64, int64, error) {
+	row := s.statements.selectUploadCountsForServer.QueryRowContext(s.ctx, serverName)
+
+	media := int64(0)
+	thumbs := int64(0)
+	err := row.Scan(&media, &thumbs)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	return media, thumbs, nil
+}