From e1a7bef9f373708cf740064a2b7cb3f69e515ca9 Mon Sep 17 00:00:00 2001
From: Duc Nguyen <minhduc140583@gmail.com>
Date: Sat, 5 Aug 2023 22:15:18 +0700
Subject: [PATCH] Refactor mongo health check

---
 elasticsearch/health_checker.go | 45 ---------------------------------
 mongo/health_checker.go         | 15 +++++------
 mongo/server_info_checker.go    |  7 ++---
 3 files changed, 11 insertions(+), 56 deletions(-)
 delete mode 100644 elasticsearch/health_checker.go

diff --git a/elasticsearch/health_checker.go b/elasticsearch/health_checker.go
deleted file mode 100644
index ef551ee..0000000
--- a/elasticsearch/health_checker.go
+++ /dev/null
@@ -1,45 +0,0 @@
-package elasticsearch
-
-import (
-	"context"
-	"github.com/elastic/go-elasticsearch"
-)
-
-type HealthChecker struct {
-	client *elasticsearch.Client
-	name   string
-}
-
-func NewHealthChecker(client *elasticsearch.Client, options ...string) *HealthChecker {
-	var name string
-	if len(options) > 0 && len(options[0]) > 0 {
-		name = options[0]
-	} else {
-		name = "elasticsearch"
-	}
-	return &HealthChecker{client, name}
-}
-
-func (e *HealthChecker) Name() string {
-	return e.name
-}
-
-func (e *HealthChecker) Check(ctx context.Context) (map[string]interface{}, error) {
-	res := make(map[string]interface{})
-	_, err := e.client.Ping()
-	if err != nil {
-		return res, err
-	}
-	return res, nil
-}
-
-func (e *HealthChecker) Build(ctx context.Context, data map[string]interface{}, err error) map[string]interface{} {
-	if err == nil {
-		return data
-	}
-	if data == nil {
-		data = make(map[string]interface{}, 0)
-	}
-	data["error"] = err.Error()
-	return data
-}
diff --git a/mongo/health_checker.go b/mongo/health_checker.go
index 3236b92..df8df5f 100644
--- a/mongo/health_checker.go
+++ b/mongo/health_checker.go
@@ -4,33 +4,33 @@ import (
 	"context"
 	"fmt"
 	"go.mongodb.org/mongo-driver/mongo"
-	"go.mongodb.org/mongo-driver/x/bsonx"
+	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"time"
 )
 
 type HealthChecker struct {
-	db      *mongo.Database
+	client  *mongo.Client
 	name    string
 	timeout time.Duration
 }
 
-func NewMongoHealthChecker(db *mongo.Database, name string, timeouts ...time.Duration) *HealthChecker {
+func NewMongoHealthChecker(client *mongo.Client, name string, timeouts ...time.Duration) *HealthChecker {
 	var timeout time.Duration
 	if len(timeouts) >= 1 {
 		timeout = timeouts[0]
 	} else {
 		timeout = 4 * time.Second
 	}
-	return &HealthChecker{db: db, name: name, timeout: timeout}
+	return &HealthChecker{client: client, name: name, timeout: timeout}
 }
-func NewHealthChecker(db *mongo.Database, options ...string) *HealthChecker {
+func NewHealthChecker(client *mongo.Client, options ...string) *HealthChecker {
 	var name string
 	if len(options) >= 1 && len(options[0]) > 0 {
 		name = options[0]
 	} else {
 		name = "mongo"
 	}
-	return NewMongoHealthChecker(db, name, 4 * time.Second)
+	return NewMongoHealthChecker(client, name, 4*time.Second)
 }
 
 func (s *HealthChecker) Name() string {
@@ -45,10 +45,9 @@ func (s *HealthChecker) Check(ctx context.Context) (map[string]interface{}, erro
 	defer cancel()
 
 	res := make(map[string]interface{})
-	info := make(map[string]interface{})
 	checkerChan := make(chan error)
 	go func() {
-		checkerChan <- s.db.RunCommand(ctx, bsonx.Doc{{"ping", bsonx.Int32(1)}}).Decode(&info)
+		checkerChan <- s.client.Ping(ctx, readpref.Primary())
 	}()
 	select {
 	case err := <-checkerChan:
diff --git a/mongo/server_info_checker.go b/mongo/server_info_checker.go
index 336ae7d..b6bc830 100644
--- a/mongo/server_info_checker.go
+++ b/mongo/server_info_checker.go
@@ -3,8 +3,8 @@ package mongo
 import (
 	"context"
 	"fmt"
+	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
-	"go.mongodb.org/mongo-driver/x/bsonx"
 	"time"
 )
 
@@ -24,7 +24,7 @@ func NewServerInfoChecker(db *mongo.Database, options ...string) *ServerInfoChec
 	} else {
 		name = "mongo"
 	}
-	return NewServerInfoCheckerWithTimeout(db, name, 4 * time.Second)
+	return NewServerInfoCheckerWithTimeout(db, name, 4*time.Second)
 }
 
 func (s *ServerInfoChecker) Name() string {
@@ -41,8 +41,9 @@ func (s *ServerInfoChecker) Check(ctx context.Context) (map[string]interface{},
 	res := make(map[string]interface{})
 	info := make(map[string]interface{})
 	checkerChan := make(chan error)
+	command := bson.D{{ "serverStatus", 1 }}
 	go func() {
-		checkerChan <- s.db.RunCommand(ctx, bsonx.Doc{{"serverStatus", bsonx.Int32(1)}}).Decode(&info)
+		checkerChan <- s.db.RunCommand(ctx, command).Decode(&info)
 	}()
 	select {
 	case err := <-checkerChan: