fix:multiple JSON return errors and optimize program structure"

Author: taolx0

internal/dms/biz/cloudbeaver.go

@@ -14,16 +14,14 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/99designs/gqlgen/graphql"
+	"github.com/99designs/gqlgen/graphql/executor"
 	"github.com/actiontech/dms/internal/dms/pkg/constant"
 	"github.com/actiontech/dms/internal/pkg/cloudbeaver"
 	"github.com/actiontech/dms/internal/pkg/cloudbeaver/model"
 	"github.com/actiontech/dms/internal/pkg/cloudbeaver/resolver"
-
 	"github.com/actiontech/dms/pkg/dms-common/api/jwt"
 	"github.com/actiontech/dms/pkg/dms-common/pkg/aes"
-
-	"github.com/99designs/gqlgen/graphql"
-	"github.com/99designs/gqlgen/graphql/executor"
 	utilLog "github.com/actiontech/dms/pkg/dms-common/pkg/log"
 	"github.com/labstack/echo/v4"
 )
@@ -334,19 +332,13 @@ func (cu *CloudbeaverUsecase) GraphQLDistributor() echo.MiddlewareFunc {
 				ctx := graphql.StartOperationTrace(context.Background())
 
 				var dbService *DBService
-				var taskID *string
 				if params.OperationName == "asyncSqlExecuteQuery" {
 					dbService, err = cu.getDbService(c.Request().Context(), params)
 					if err != nil {
 						cu.log.Error(err)
 						return err
 					}
 
-					taskID, err = cu.SaveCbOpLog(c, dbService, params, next)
-					if err != nil {
-						return err
-					}
-
 					if !cu.isEnableSQLAudit(dbService) {
 						cloudbeaverResBuf := new(bytes.Buffer)
 						mw := io.MultiWriter(c.Response().Writer, cloudbeaverResBuf)
@@ -482,9 +474,14 @@ func (cu *CloudbeaverUsecase) GraphQLDistributor() echo.MiddlewareFunc {
 					cloudbeaverNext = func(c echo.Context) ([]byte, error) {
 						resp, ok = c.Get(cloudbeaver.AuditResultKey).(cloudbeaver.AuditResults)
 						if ok {
-							cu.UpdateCbOp(ctx, resp, taskID)
 							if !resp.IsSuccess {
+								cu.SaveCbOperationLogWithoutNext(c, dbService, params, resp)
 								return nil, c.JSON(http.StatusOK, convertToResp(resp))
+							} else {
+								err = cu.SaveCbOpLog(c, dbService, params, resp, next)
+								if err != nil {
+									return nil, nil
+								}
 							}
 						}
 
@@ -510,9 +507,16 @@ func (cu *CloudbeaverUsecase) GraphQLDistributor() echo.MiddlewareFunc {
 				} else {
 					cloudbeaverNext = func(c echo.Context) ([]byte, error) {
 						resp, ok = c.Get(cloudbeaver.AuditResultKey).(cloudbeaver.AuditResults)
-						cu.UpdateCbOp(ctx, resp, taskID)
-						if !resp.IsSuccess {
-							return nil, c.JSON(http.StatusOK, convertToResp(resp))
+						if ok {
+							if !resp.IsSuccess {
+								cu.SaveCbOperationLogWithoutNext(c, dbService, params, resp)
+								return nil, c.JSON(http.StatusOK, convertToResp(resp))
+							} else {
+								err = cu.SaveCbOpLog(c, dbService, params, resp, next)
+								if err != nil {
+									return nil, nil
+								}
+							}
 						}
 
 						resWrite = &responseProcessWriter{tmp: &bytes.Buffer{}, ResponseWriter: c.Response().Writer}

internal/dms/biz/cloudbeaver_ce.go

@@ -15,18 +15,18 @@ func (cu *CloudbeaverUsecase) ResetDbServiceByAuth(ctx context.Context, activeDB
 	return activeDBServices, nil
 }
 
-func (cu *CloudbeaverUsecase) UpdateCbOp(ctx context.Context, resp cloudbeaver.AuditResults, taskID *string) {
-	return
-}
-
 func (cu *CloudbeaverUsecase) UpdateCbOpResult(c echo.Context, cloudbeaverResBuf *bytes.Buffer, params *graphql.RawParams, ctx context.Context) error {
 	return nil
 }
 
-func (cu *CloudbeaverUsecase) SaveCbOpLog(c echo.Context, dbService *DBService, params *graphql.RawParams, next echo.HandlerFunc) (*string, error) {
-	return nil, nil
+func (cu *CloudbeaverUsecase) SaveCbOpLog(c echo.Context, dbService *DBService, params *graphql.RawParams, resp cloudbeaver.AuditResults, next echo.HandlerFunc) error {
+	return nil
 }
 
 func (cu *CloudbeaverUsecase) SaveUiOp(c echo.Context, buf *bytes.Buffer, params *graphql.RawParams) error {
 	return nil
 }
+
+func (cu *CloudbeaverUsecase) SaveCbOperationLogWithoutNext(c echo.Context, dbService *DBService, params *graphql.RawParams, resp cloudbeaver.AuditResults) {
+	return
+}