Merge pull request #109 from TRON-US/feature/checkMultiDBRuntime

Feature/check multi db runtime

utils/grpc/runtime.go

@@ -42,7 +42,7 @@ func (s *RuntimeServer) CheckRuntime(ctx context.Context, req *sharedpb.SignedRu
 	}
 
 	//check runtime in shared
-	res, err := utils.CheckRuntime(ctx, req, connection)
+	res, err := utils.CheckDBConnection(ctx, req, connection)
 	if err != nil {
 		return nil, err
 	}

utils/grpc/setup_server.go

@@ -1,15 +1,17 @@
 package grpc
 
 import (
+	"context"
 	"fmt"
-	"net"
-
 	"github.com/tron-us/go-btfs-common/controller"
 	"github.com/tron-us/go-btfs-common/protos/escrow"
 	"github.com/tron-us/go-btfs-common/protos/guard"
 	"github.com/tron-us/go-btfs-common/protos/hub"
 	"github.com/tron-us/go-btfs-common/protos/shared"
 	"github.com/tron-us/go-btfs-common/protos/status"
+	"github.com/tron-us/go-btfs-common/utils"
+	"github.com/tron-us/go-common/v2/db"
+	"net"
 
 	"github.com/tron-us/go-common/v2/constant"
 	"github.com/tron-us/go-common/v2/log"
@@ -64,13 +66,29 @@ func (s *GrpcServer) GrpcServer(port string, dbURLs map[string]string, rdURL str
 
 	s.lis = lis
 
-	s.CreateServer(s.serverName, options...).
-		CreateHealthServer().
-		RegisterServer(server).
-		RegisterHealthServer().
-		WithReflection().
-		WithGracefulTermDetectAndExec().
-		AcceptConnection()
+	done := make(chan bool)
+
+	go func() {
+		s.CreateServer(s.serverName, options...).
+			CreateHealthServer().
+			RegisterServer(server).
+			RegisterHealthServer().
+			WithReflection().
+			WithGracefulTermDetectAndExec().
+			AcceptConnection()
+		done <- true
+	}()
+
+	ctx := context.Background()
+	req := new(shared.SignedRuntimeInfoRequest)
+	connection := db.ConnectionUrls{RdURL: rdURL, PgURL: dbURLs}
+
+	_, err = utils.CheckDBConnection(ctx, req, connection)
+	if err != nil {
+		log.Panic("Got error", zap.Error(err))
+	}
+
+	<-done
 
 	return s
 }

utils/runtime.go

@@ -17,52 +17,73 @@ import (
 const DBURLDNE = "DB URL does not exist !!"
 const RDURLDNE = "RD URL does not exist !!"
 
-func CheckRuntime(ctx context.Context, runtime *sharedpb.SignedRuntimeInfoRequest, connection db.ConnectionUrls) (*sharedpb.RuntimeInfoReport, error) {
+func CheckDBConnection(ctx context.Context, runtime *sharedpb.SignedRuntimeInfoRequest, connection db.ConnectionUrls) (*sharedpb.RuntimeInfoReport, error) {
 	// db runtime
 	report := new(sharedpb.RuntimeInfoReport)
 	report.Status = sharedpb.RuntimeInfoReport_RUNNING
-	report.DbStatusExtra = map[string]string{}
+	report.DbStatusExtra = make(map[string]string)
 
 	for key, url := range connection.PgURL {
+		// Assume the database connection is healthy
+		report.DbStatusExtra[key] = DBURLDNE
+
 		if url != "" {
+			// Log the connection string
+			pgdb := postgres.CreateTGPGDB(url)
+			log.Info("Postgres",
+				zap.String("name", key),
+				zap.String("user", pgdb.Options().User),
+				zap.String("host", pgdb.Options().Addr),
+				zap.String("db", pgdb.Options().Database),
+			)
+
 			// Check postgres dbWrite
-			PGDBWrite := postgres.CreateTGPGDB(url)
-			if err := PGDBWrite.Ping(); err != nil {
+			if err := pgdb.Ping(); err != nil {
 				report.DbStatusExtra[key] = constant.DBWriteConnectionError
 				report.Status = sharedpb.RuntimeInfoReport_SICK
 				log.Error(constant.DBWriteConnectionError, zap.Error(err))
 			}
-			// Check postgres dbRead
-			PGDBRead := postgres.CreateTGPGDB(url)
-			if err := PGDBRead.Ping(); err != nil {
-				report.DbStatusExtra[key] = constant.DBReadConnectionError
-				report.Status = sharedpb.RuntimeInfoReport_SICK
-				log.Error(constant.DBReadConnectionError, zap.Error(err))
-			}
-			// Assume the database connection is healthy
+
+			// Set the database connection is healthy
 			report.DbStatusExtra[key] = constant.DBConnectionHealthy
-		} else {
-			report.DbStatusExtra[key] = DBURLDNE
 		}
+
+		// Log status
+		log.Info(key + ":" + report.DbStatusExtra[key])
 	}
 
+	// Assume the redis connection is not present
+	report.RdStatusExtra = RDURLDNE
+
 	// Check redis environment variable
 	if connection.RdURL != "" {
+		// Parse redis url
 		opts, errParse := redis.ParseRedisURL(connection.RdURL)
 		if errParse != nil {
 			log.Error(constant.RDURLParseError, zap.Error(errParse))
 		}
+
+		// Log connection string
+		log.Info("Redis URL",
+			zap.String("host", opts.Addr),
+			zap.String("db", string(opts.DB)),
+		)
+
+		// Check redis connection
 		errConn := redis.CheckRedisConnection(redis.NewRedisConn(opts))
 		if errConn != nil {
 			report.RdStatusExtra = constant.RDConnectionError
 			report.Status = sharedpb.RuntimeInfoReport_SICK
 			log.Error(constant.RDConnectionError, zap.Error(errConn))
 		}
-		// Assume the redis connection is healthy
+
+		// Set redis connection to healthy
 		report.RdStatusExtra = constant.RDConnectionHealthy
-	} else {
-		report.RdStatusExtra = RDURLDNE
 	}
+
+	// Log status
+	log.Info(report.RdStatusExtra)
+
 	// Remaining fields will be populated by the calling service
 	// Reserve: only pass fatal error to higher level
 	return report, nil

utils/runtime_test.go

@@ -23,7 +23,7 @@ func init() {
 	}
 }
 
-func TestCheckRuntimeDB(t *testing.T) {
+func TestCheckDBConnection(t *testing.T) {
 
 	//setup connection (postgres) object
 	pgConMaps := map[string]string{"DB_URL_STATUS": config.DbStatusURL, "DB_URL_GUARD": config.DbGuardURL}
@@ -34,7 +34,7 @@ func TestCheckRuntimeDB(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 	shared := new(sharedpb.SignedRuntimeInfoRequest)
-	runtimeInfoReportPass, err := CheckRuntime(ctx, shared, connection)
+	runtimeInfoReportPass, err := CheckDBConnection(ctx, shared, connection)
 	assert.Nil(t, err, zap.Error(err))
 	assert.True(t, strings.Contains(runtimeInfoReportPass.DbStatusExtra["DB_URL_STATUS"], constant.DBConnectionHealthy), "connection not successful")
 	assert.True(t, strings.Contains(runtimeInfoReportPass.DbStatusExtra["DB_URL_GUARD"], constant.DBConnectionHealthy), "connection not successful")
@@ -45,12 +45,12 @@ func TestCheckRuntimeDB(t *testing.T) {
 		PgURL: pgEmptyConMaps,
 		RdURL: "",
 	}
-	runtimeInfoReportFail, err := CheckRuntime(ctx, shared, emptyConnection)
+	runtimeInfoReportFail, err := CheckDBConnection(ctx, shared, emptyConnection)
 	assert.Nil(t, err, zap.Error(err))
 	assert.True(t, strings.Contains(runtimeInfoReportFail.DbStatusExtra["DB_URL_STATUS"], "DB URL does not exist !!"), "DB URL does not exist !!")
 	assert.True(t, strings.Contains(runtimeInfoReportFail.DbStatusExtra["DB_URL_GUARD"], "DB URL does not exist !!"), "DB URL does not exist !!")
 }
-func TestCheckRuntimeRD(t *testing.T) {
+func TestCheckDBConnectionRD(t *testing.T) {
 	const RDURLDNE = "RD URL does not exist !!"
 	//setup connection (redis) object
 	var connection = db.ConnectionUrls{
@@ -60,12 +60,12 @@ func TestCheckRuntimeRD(t *testing.T) {
 	shared := new(sharedpb.SignedRuntimeInfoRequest)
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
-	runtimeInfoReportPass, err := CheckRuntime(ctx, shared, connection)
+	runtimeInfoReportPass, err := CheckDBConnection(ctx, shared, connection)
 	assert.Nil(t, err, zap.Error(err))
 	assert.True(t, strings.Contains(runtimeInfoReportPass.RdStatusExtra, constant.RDConnectionHealthy), "Redis is not running")
 	//disable connection string
 	connection.RdURL = ""
-	runtimeInfoReportFail, err := CheckRuntime(ctx, shared, connection)
+	runtimeInfoReportFail, err := CheckDBConnection(ctx, shared, connection)
 	assert.Nil(t, err, zap.Error(err))
 	assert.True(t, strings.Contains(runtimeInfoReportFail.RdStatusExtra, RDURLDNE), "Redis connection is still provided, error!")
 }