db调用代码修改 etcd修改

Tiktok-Lite · Aug 29, 2023 · 8b8099d · 8b8099d
1 parent dc49632
commit 8b8099d
Show file tree

Hide file tree

Showing 6 changed files with 86 additions and 20 deletions.
diff --git a/cmd/api/rpc/comment.go b/cmd/api/rpc/comment.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/comment"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/comment/commentservice"
+	"github.com/Tiktok-Lite/kotkit/pkg/etcd"
+	"github.com/Tiktok-Lite/kotkit/pkg/helper/constant"
 	"github.com/cloudwego/kitex/client"
 	"github.com/spf13/viper"
 )
@@ -14,20 +16,41 @@ var (
 )
 
 func InitComment(config *viper.Viper) {
+	r, err := etcd.Resolver()
+	if err != nil {
+		logger.Errorf("Error occurs when creating etcd resolver: %v", err)
+		panic(err)
+	}
+
 	commentServiceName := config.GetString("server.name")
-	commentServiceAddr := fmt.Sprintf("%s:%d", config.GetString("server.host"), config.GetInt("server.port"))
 
-	c, err := commentservice.NewClient(commentServiceName, client.WithHostPorts(commentServiceAddr))
+	c, err := commentservice.NewClient(commentServiceName, client.WithResolver(r))
+
 	if err != nil {
+		logger.Errorf("Error occurs when creating login client: %v", err)
 		panic(err)
 	}
-
 	commentClient = c
 }
 
 func CommentList(ctx context.Context, req *comment.DouyinCommentListRequest) (*comment.DouyinCommentListResponse, error) {
-	return commentClient.CommentList(ctx, req)
+	resp, err := commentClient.CommentList(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	if resp.StatusCode == constant.StatusErrorCode {
+		return resp, fmt.Errorf(resp.StatusMsg)
+	}
+	return resp, nil
 }
 func CommentAction(ctx context.Context, req *comment.DouyinCommentActionRequest) (*comment.DouyinCommentActionResponse, error) {
-	return commentClient.CommentAction(ctx, req)
+	resp, err := commentClient.CommentAction(ctx, req)
+	if err != nil {
+		return resp, err
+	}
+	if resp.StatusCode == constant.StatusErrorCode {
+		return resp, fmt.Errorf(resp.StatusMsg)
+	}
+
+	return resp, nil
 }
diff --git a/cmd/api/rpc/message.go b/cmd/api/rpc/message.go
@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/message"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/message/messageservice"
+	"github.com/Tiktok-Lite/kotkit/pkg/etcd"
+	"github.com/Tiktok-Lite/kotkit/pkg/helper/constant"
 	"github.com/cloudwego/kitex/client"
 	"github.com/spf13/viper"
 )
@@ -14,20 +16,41 @@ var (
 )
 
 func InitMessage(config *viper.Viper) {
+	r, err := etcd.Resolver()
+	if err != nil {
+		logger.Errorf("Error occurs when creating etcd resolver: %v", err)
+		panic(err)
+	}
+
 	messageServiceName := config.GetString("server.name")
-	messageServiceAddr := fmt.Sprintf("%s:%d", config.GetString("server.host"), config.GetInt("server.port"))
 
-	c, err := messageservice.NewClient(messageServiceName, client.WithHostPorts(messageServiceAddr))
+	c, err := messageservice.NewClient(messageServiceName, client.WithResolver(r))
+
 	if err != nil {
+		logger.Errorf("Error occurs when creating login client: %v", err)
 		panic(err)
 	}
-
 	messageClient = c
 }
 
 func MessageList(ctx context.Context, req *message.DouyinMessageChatRequest) (*message.DouyinMessageChatResponse, error) {
-	return messageClient.MessageChat(ctx, req)
+	resp, err := messageClient.MessageChat(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	if resp.StatusCode == constant.StatusErrorCode {
+		return resp, fmt.Errorf(resp.StatusMsg)
+	}
+	return resp, nil
 }
 func MessageAction(ctx context.Context, req *message.DouyinMessageActionRequest) (*message.DouyinMessageActionResponse, error) {
-	return messageClient.MessageAction(ctx, req)
+	resp, err := messageClient.MessageAction(ctx, req)
+	if err != nil {
+		return resp, err
+	}
+	if resp.StatusCode == constant.StatusErrorCode {
+		return resp, fmt.Errorf(resp.StatusMsg)
+	}
+
+	return resp, nil
 }
diff --git a/cmd/comment/main.go b/cmd/comment/main.go
@@ -4,9 +4,11 @@ import (
 	"fmt"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/comment/commentservice"
 	"github.com/Tiktok-Lite/kotkit/pkg/conf"
+	"github.com/Tiktok-Lite/kotkit/pkg/etcd"
 	"github.com/Tiktok-Lite/kotkit/pkg/helper/constant"
 	"github.com/Tiktok-Lite/kotkit/pkg/helper/jwt"
 	"github.com/Tiktok-Lite/kotkit/pkg/log"
+	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/server"
 	"net"
 )
@@ -21,20 +23,28 @@ var (
 	Jwt           *jwt.JWT
 )
 
-func Init() {
+func init() {
 	Jwt = jwt.NewJWT([]byte(signingKey))
 }
 
 func main() {
+	r, err := etcd.Registry()
+	if err != nil {
+		logger.Errorf("Error occurs when creating etcd registry: %v", err)
+		panic(err)
+	}
+
 	addr, err := net.ResolveTCPAddr("tcp", serviceAddr)
 	if err != nil {
 		logger.Errorf("Error occurs when resolving comment service address: %v", err)
 		panic(err)
 	}
 	svr := commentservice.NewServer(new(CommentServiceImpl),
+		server.WithServerBasicInfo(&rpcinfo.EndpointBasicInfo{ServiceName: serviceName}),
 		server.WithServiceAddr(addr),
+		server.WithRegistry(r),
 	)
-	Init()
+
 	err = svr.Run()
 
 	if err != nil {

diff --git a/cmd/message/main.go b/cmd/message/main.go
@@ -4,9 +4,11 @@ import (
 	"fmt"
 	"github.com/Tiktok-Lite/kotkit/kitex_gen/message/messageservice"
 	"github.com/Tiktok-Lite/kotkit/pkg/conf"
+	"github.com/Tiktok-Lite/kotkit/pkg/etcd"
 	"github.com/Tiktok-Lite/kotkit/pkg/helper/constant"
 	"github.com/Tiktok-Lite/kotkit/pkg/helper/jwt"
 	"github.com/Tiktok-Lite/kotkit/pkg/log"
+	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/server"
 	"net"
 )
@@ -21,25 +23,33 @@ var (
 	Jwt           *jwt.JWT
 )
 
-func Init() {
+func init() {
 	Jwt = jwt.NewJWT([]byte(signingKey))
 }
 
 func main() {
+	r, err := etcd.Registry()
+	if err != nil {
+		logger.Errorf("Error occurs when creating etcd registry: %v", err)
+		panic(err)
+	}
+
 	addr, err := net.ResolveTCPAddr("tcp", serviceAddr)
 	if err != nil {
 		logger.Errorf("Error occurs when resolving message service address: %v", err)
 		panic(err)
 	}
 	svr := messageservice.NewServer(new(MessageServiceImpl),
+		server.WithServerBasicInfo(&rpcinfo.EndpointBasicInfo{ServiceName: serviceName}),
 		server.WithServiceAddr(addr),
+		server.WithRegistry(r),
 	)
-	Init()
+
 	err = svr.Run()
 
 	if err != nil {
 		logger.Errorf("Error occurs when running message service server: %v", err)
 		panic(err)
 	}
-	logger.Infof("message service server start successfully at %s", serviceAddr)
+	logger.Infof("Message service server start successfully at %s", serviceAddr)
 }
diff --git a/internal/db/comment.go b/internal/db/comment.go
@@ -7,15 +7,15 @@ import (
 )
 
 func AddComment(comment *model.Comment) error {
-	if err := db.Save(comment).Error; err != nil {
+	if err := DB().Save(comment).Error; err != nil {
 		return errors.New("failed to add comment")
 	}
 	return nil
 }
 
 func QueryCommentByVideoID(videoId int64) ([]*model.Comment, error) {
 	var comments []*model.Comment
-	if err := db.Debug().Where("video_id", videoId).
+	if err := DB().Debug().Where("video_id", videoId).
 		Order("created_at desc").Find(&comments).Error; err != nil {
 		logger.Errorf("failed to query comments from databse: %v", err)
 		return nil, err
@@ -24,7 +24,7 @@ func QueryCommentByVideoID(videoId int64) ([]*model.Comment, error) {
 }
 
 func DeleteCommentById(commentId int64) error {
-	if err := db.Delete(&model.Comment{}, commentId).Error; err != nil {
+	if err := DB().Delete(&model.Comment{}, commentId).Error; err != nil {
 		return errors.New("failed to delete comment")
 	}
 	return nil

diff --git a/internal/db/message.go b/internal/db/message.go
@@ -7,14 +7,14 @@ import (
 )
 
 func SendMessage(message *model.Message) error {
-	if err := db.Save(message).Error; err != nil {
+	if err := DB().Save(message).Error; err != nil {
 		return errors.New("failed to send message")
 	}
 	return nil
 }
 func QueryMessageList(userId int64, toUserId int64) ([]*model.Message, error) {
 	var messages []*model.Message
-	if err := db.Debug().Where("from_user_id", userId).Where("to_user_id", toUserId).
+	if err := DB().Debug().Where("from_user_id", userId).Where("to_user_id", toUserId).
 		Order("created_at desc").Find(&messages).Error; err != nil {
 		logger.Errorf("failed to query messages from databse: %v", err)
 		return nil, err