coze-studio/backend/domain/shortcutcmd/internal/dal/dao.go

/*
 * 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
}