feat: manually mirror opencoze's code from bytedance

Change-Id: I09a73aadda978ad9511264a756b2ce51f5761adf

backend/domain/shortcutcmd/internal/dal/dao.go

@@ -0,0 +1,204 @@
+/*
+ * Copyright 2025 coze-dev Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package dal
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"gorm.io/gorm"
+
+	"github.com/coze-dev/coze-studio/backend/domain/shortcutcmd/entity"
+	"github.com/coze-dev/coze-studio/backend/domain/shortcutcmd/internal/dal/model"
+	"github.com/coze-dev/coze-studio/backend/domain/shortcutcmd/internal/dal/query"
+	"github.com/coze-dev/coze-studio/backend/infra/contract/idgen"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/conv"
+	"github.com/coze-dev/coze-studio/backend/pkg/lang/slices"
+	"github.com/coze-dev/coze-studio/backend/pkg/logs"
+)
+
+type ShortCutCmdDAO struct {
+	db    *gorm.DB
+	query *query.Query
+	idgen idgen.IDGenerator
+}
+
+func NewShortCutCmdDAO(db *gorm.DB, idgen idgen.IDGenerator) *ShortCutCmdDAO {
+	return &ShortCutCmdDAO{
+		db:    db,
+		query: query.Use(db),
+		idgen: idgen,
+	}
+}
+func (dao *ShortCutCmdDAO) Create(ctx context.Context, shortcut *entity.ShortcutCmd) (*entity.ShortcutCmd, error) {
+	createPO, err := dao.buildCreatePO(ctx, shortcut)
+	if err != nil {
+		return nil, err
+	}
+	createErr := dao.query.ShortcutCommand.WithContext(ctx).Debug().Create(createPO)
+	logs.CtxInfof(ctx, "ShortcutCommand %v, err:%v", conv.DebugJsonToStr(createPO), err)
+	if createErr != nil {
+		return nil, createErr
+	}
+	return createPO, nil
+}
+
+func (dao *ShortCutCmdDAO) buildCreatePO(ctx context.Context, shortcut *entity.ShortcutCmd) (*model.ShortcutCommand, error) {
+	cmdID, err := dao.idgen.GenID(ctx)
+
+	if err != nil {
+		return nil, err
+	}
+
+	po := &model.ShortcutCommand{
+		ObjectID:        shortcut.ObjectID,
+		CommandID:       cmdID,
+		CommandName:     shortcut.CommandName,
+		ShortcutCommand: shortcut.ShortcutCommand,
+		Description:     shortcut.Description,
+		SendType:        shortcut.SendType,
+		ToolType:        shortcut.ToolType,
+		WorkFlowID:      shortcut.WorkFlowID,
+		PluginID:        shortcut.PluginID,
+		PluginToolName:  shortcut.PluginToolName,
+		TemplateQuery:   shortcut.TemplateQuery,
+		Components:      shortcut.Components,
+		CardSchema:      shortcut.CardSchema,
+		ToolInfo:        shortcut.ToolInfo,
+		Status:          1,
+		CreatorID:       shortcut.CreatorID,
+		CreatedAt:       time.Now().UnixMilli(),
+		UpdatedAt:       time.Now().UnixMilli(),
+		AgentID:         shortcut.AgentID,
+		ShortcutIcon:    shortcut.ShortcutIcon,
+		PluginToolID:    shortcut.PluginToolID,
+	}
+	return po, nil
+}
+
+func (dao *ShortCutCmdDAO) Update(ctx context.Context, shortcut *entity.ShortcutCmd) (*entity.ShortcutCmd, error) {
+	updatePO := dao.buildUpdatePO(ctx, shortcut)
+	_, updateErr := dao.query.ShortcutCommand.WithContext(ctx).Debug().Where(dao.query.ShortcutCommand.CommandID.Eq(shortcut.CommandID)).Updates(updatePO)
+
+	if updateErr != nil {
+		return nil, updateErr
+	}
+	return updatePO, nil
+}
+
+func (dao *ShortCutCmdDAO) buildUpdatePO(ctx context.Context, shortcut *entity.ShortcutCmd) *model.ShortcutCommand {
+
+	po := shortcut
+	po.UpdatedAt = time.Now().UnixMilli()
+
+	return po
+}
+
+func (dao *ShortCutCmdDAO) List(ctx context.Context, lm *entity.ListMeta) ([]*entity.ShortcutCmd, error) {
+	if len(lm.CommandIDs) == 0 {
+		return nil, nil
+	}
+
+	do := dao.query.ShortcutCommand.WithContext(ctx).Where(dao.query.ShortcutCommand.ObjectID.Eq(lm.ObjectID)).Debug().
+		Where(dao.query.ShortcutCommand.CommandID.In(lm.CommandIDs...)).
+		Where(dao.query.ShortcutCommand.IsOnline.Eq(lm.IsOnline))
+
+	poList, err := do.Find()
+
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, err
+	}
+	return slices.Transform(poList, func(po *model.ShortcutCommand) *entity.ShortcutCmd {
+		return po
+	}), nil
+}
+
+func (dao *ShortCutCmdDAO) GetByCmdID(ctx context.Context, cmdID int64, isOnline int32) (*entity.ShortcutCmd, error) {
+
+	po, err := dao.query.ShortcutCommand.WithContext(ctx).Where(dao.query.ShortcutCommand.CommandID.Eq(cmdID)).
+		Where(dao.query.ShortcutCommand.IsOnline.Eq(isOnline)).Debug().
+		First()
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return nil, err
+	}
+	return po, nil
+}
+
+func (dao *ShortCutCmdDAO) PublishCMDs(ctx context.Context, objID int64, cmdIDs []int64) error {
+	if len(cmdIDs) == 0 {
+		return nil
+	}
+	draftCmds, err := dao.query.ShortcutCommand.WithContext(ctx).Debug().
+		Where(dao.query.ShortcutCommand.ObjectID.Eq(objID)).
+		Where(dao.query.ShortcutCommand.CommandID.In(cmdIDs...)).
+		Where(dao.query.ShortcutCommand.IsOnline.Eq(0)).
+		Find()
+	if err != nil {
+		return err
+	}
+
+	tx := dao.query.Begin()
+	defer func() {
+		if tx.Error != nil {
+			rbErr := tx.Rollback()
+			if rbErr != nil {
+				logs.CtxErrorf(ctx, "rollback failed, err:%v", rbErr)
+			}
+		}
+		cErr := tx.Commit()
+		if cErr != nil {
+			logs.CtxErrorf(ctx, "commit failed, err:%v", cErr)
+		}
+	}()
+
+	onlineCmds, err := dao.query.ShortcutCommand.WithContext(ctx).Debug().
+		Where(dao.query.ShortcutCommand.ObjectID.Eq(objID)).
+		Where(dao.query.ShortcutCommand.CommandID.In(cmdIDs...)).
+		Where(dao.query.ShortcutCommand.IsOnline.Eq(1)).Find()
+
+	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
+		return err
+	}
+	onelineCmdMap := make(map[int64]*model.ShortcutCommand)
+	for _, one := range onlineCmds {
+		onelineCmdMap[one.CommandID] = one
+	}
+
+	for _, item := range draftCmds {
+		item.IsOnline = 1
+		item.UpdatedAt = time.Now().UnixMilli()
+		item.ID = 0
+		var opErr error
+		if _, ok := onelineCmdMap[item.CommandID]; !ok {
+			opErr = tx.ShortcutCommand.WithContext(ctx).Debug().Create(item)
+		} else {
+			opErr = tx.ShortcutCommand.WithContext(ctx).Debug().
+				Where(dao.query.ShortcutCommand.ObjectID.Eq(item.ObjectID)).
+				Where(dao.query.ShortcutCommand.CommandID.Eq(item.CommandID)).
+				Where(dao.query.ShortcutCommand.IsOnline.Eq(item.IsOnline)).
+				Save(item)
+		}
+
+		logs.CtxInfof(ctx, "publish cmd %v, err:%v", conv.DebugJsonToStr(item), opErr)
+		if opErr != nil {
+			return opErr
+		}
+	}
+
+	return nil
+}