Commit b69fc5fa authored by Andrei Mihu's avatar Andrei Mihu
Browse files

Faster row counts in console handlers.

parent c5aaecc3

main.go

+5 −3
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ 
package main



import (

	"context"

	"database/sql"

	"fmt"

	"math/rand"
@@ -215,8 +216,9 @@ func dbConnect(multiLogger *zap.Logger, config server.Config) (*sql.DB, string) 
	if err != nil {

		multiLogger.Fatal("Error connecting to database", zap.Error(err))

	}

	err = db.Ping()

	if err != nil {

	// Limit the time allowed to ping database and get version to 15 seconds total.

	ctx, _ := context.WithTimeout(context.Background(), 15*time.Second)

	if err = db.PingContext(ctx); err != nil {

		multiLogger.Fatal("Error pinging database", zap.Error(err))

	}
@@ -225,7 +227,7 @@ func dbConnect(multiLogger *zap.Logger, config server.Config) (*sql.DB, string) 
	db.SetMaxIdleConns(config.GetDatabase().MaxIdleConns)



	var dbVersion string

	if err := db.QueryRow("SELECT version()").Scan(&dbVersion); err != nil {

	if err = db.QueryRowContext(ctx, "SELECT version()").Scan(&dbVersion); err != nil {

		multiLogger.Fatal("Error querying database version", zap.Error(err))

	}

server/console_storage.go

+16 −3
Original line number Diff line number Diff line
@@ -215,9 +215,22 @@ func (s *ConsoleServer) WriteStorageObject(ctx context.Context, in *console.Writ 
}



func countStorage(ctx context.Context, logger *zap.Logger, db *sql.DB) int32 {

	var count int

	var count sql.NullInt64

	if err := db.QueryRowContext(ctx, "SELECT reltuples::BIGINT FROM pg_class WHERE relname = 'storage'").Scan(&count); err != nil {

		logger.Warn("Error counting storage objects.", zap.Error(err))

		if err == context.Canceled {

			// If the context was cancelled do not attempt the full count.

			return 0

		}

	}

	if count.Valid && count.Int64 != 0 {

		// Use fast count result.

		return int32(count.Int64)

	}



	// If the fast count failed, returned NULL, or returned 0 try a full count.

	if err := db.QueryRowContext(ctx, "SELECT count(collection) FROM storage").Scan(&count); err != nil {

		logger.Error("Error counting storage objects.", zap.Error(err))

		logger.Warn("Error counting storage objects.", zap.Error(err))

	}

	return int32(count)

	return int32(count.Int64)

}

server/console_user.go

+16 −3
Original line number Diff line number Diff line
@@ -220,9 +220,22 @@ func (s *ConsoleServer) ListUsers(ctx context.Context, in *console.ListUsersRequ 
}



func countUsers(ctx context.Context, logger *zap.Logger, db *sql.DB) int32 {

	var count int

	var count sql.NullInt64

	if err := db.QueryRowContext(ctx, "SELECT reltuples::BIGINT FROM pg_class WHERE relname = 'users'").Scan(&count); err != nil {

		logger.Warn("Error counting storage objects.", zap.Error(err))

		if err == context.Canceled {

			// If the context was cancelled do not attempt the full count.

			return 0

		}

	}

	if count.Valid && count.Int64 != 0 {

		// Use fast count result.

		return int32(count.Int64)

	}



	// If the fast count failed, returned NULL, or returned 0 try a full count.

	if err := db.QueryRowContext(ctx, "SELECT count(id) FROM users").Scan(&count); err != nil {

		logger.Error("Error counting users.", zap.Error(err))

		logger.Warn("Error counting storage objects.", zap.Error(err))

	}

	return int32(count)

	return int32(count.Int64)

}