Skip to content

Commit

Permalink
fix(dbm-services): 权限克隆代码优化 TencentBlueKing#7048
Browse files Browse the repository at this point in the history
  • Loading branch information
ymakedaq committed Sep 20, 2024
1 parent 9a4817d commit ad0677f
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 17 deletions.
7 changes: 4 additions & 3 deletions dbm-services/mysql/db-priv/handler/clone_instance_priv.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package handler

import (
"encoding/json"
"io"
"io/ioutil"
"log/slog"
"strings"
Expand Down Expand Up @@ -43,15 +44,15 @@ func (m *PrivService) CloneInstancePriv(c *gin.Context) {
var input service.CloneInstancePrivPara
ticket := strings.TrimPrefix(c.FullPath(), "/priv/")

body, err := ioutil.ReadAll(c.Request.Body)
body, err := io.ReadAll(c.Request.Body)
if err != nil {
slog.Error("msg", err)
slog.Error("msg", slog.Any("err", err))
SendResponse(c, errno.ErrBind, err)
return
}

if err = json.Unmarshal(body, &input); err != nil {
slog.Error("msg", err)
slog.Error("msg", slog.Any("err", err))
SendResponse(c, errno.ErrBind, err)
return
}
Expand Down
8 changes: 6 additions & 2 deletions dbm-services/mysql/db-priv/service/clone_instance_priv.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@ import (
"strings"
"time"

"github.com/samber/lo"

"dbm-services/common/go-pubpkg/errno"
)

Expand Down Expand Up @@ -59,8 +61,10 @@ func (m *CloneInstancePrivPara) CloneInstancePriv(jsonPara string, ticket string
}

// 此处单集群instanceType是single
if instanceType == machineTypeSingle || instanceType == machineTypeBackend ||
instanceType == machineTypeRemote || instanceType == machineTypeSpider {
// if instanceType == machineTypeSingle || instanceType == machineTypeBackend ||
// instanceType == machineTypeRemote || instanceType == machineTypeSpider
allowTypes := []string{machineTypeSingle, machineTypeBackend, machineTypeRemote, machineTypeSpider}
if lo.Contains(allowTypes, instanceType) {
userGrants, err := GetRemotePrivilege(m.Source.Address, "", *m.BkCloudId, instanceType, "", false)
if err != nil {
return err
Expand Down
31 changes: 19 additions & 12 deletions dbm-services/mysql/db-priv/service/clone_instance_priv_base_func.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package service

import (
"context"
"errors"
"fmt"
"log/slog"
"regexp"
Expand Down Expand Up @@ -102,7 +103,7 @@ func GetRemotePrivilege(address string, host string, bkCloudId int64, instanceTy
}()
var Grants []string
var err error
slog.Info("msg", "userHost", userHost)
// slog.Info("msg", "userHost", userHost)
err = GetUserGantSql(needShowCreateUser, userHost, address, &Grants, bkCloudId)
if err != nil {
errorChan <- err
Expand Down Expand Up @@ -586,11 +587,12 @@ func ImportMysqlPrivileges(userGrants []UserGrant, address string, bkCloudId int
mu sync.RWMutex
errs []string
}
var errMsg Err

wg := sync.WaitGroup{}
limit := rate.Every(time.Millisecond * 20) // QPS:50
burst := 50 // 桶容量 50
limiter := rate.NewLimiter(limit, burst)
errChan := make(chan error, 20)
for _, row := range userGrants {
errLimiter := limiter.Wait(context.Background())
if errLimiter != nil {
Expand All @@ -602,27 +604,32 @@ func ImportMysqlPrivileges(userGrants []UserGrant, address string, bkCloudId int
defer func() {
wg.Done()
}()
slog.Info("msg", "user@host", row.UserHost)
//slog.Info("msg", "user@host", row.UserHost)
queryRequest := QueryRequest{[]string{address}, row.Grants, true, 60, bkCloudId}
_, err := OneAddressExecuteSql(queryRequest)
if err != nil {
errMsg.mu.Lock()
errMsg.errs = append(errMsg.errs, err.Error())
errMsg.mu.Unlock()
errChan <- err
return
}
}(row)
}
wg.Wait()

go func() {
wg.Wait()
close(errChan)
}()
var errs []error
for err := range errChan {
errs = append(errs, err)
}

queryRequest := QueryRequest{[]string{address}, []string{flushPriv}, true, 60, bkCloudId}
_, err := OneAddressExecuteSql(queryRequest)
if err != nil {
errMsg.mu.Lock()
errMsg.errs = append(errMsg.errs, err.Error())
errMsg.mu.Unlock()
slog.Error("execute flush privilges failed", slog.Any("err message:", err))
}
if len(errMsg.errs) > 0 {
return fmt.Errorf(strings.Join(errMsg.errs, "\n"))
if len(errs) > 0 {
return errors.Join(errs...)
}
return nil
}
Expand Down

0 comments on commit ad0677f

Please sign in to comment.