87 lines
2.6 KiB
Go
87 lines
2.6 KiB
Go
|
package repository
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/google/uuid"
|
||
|
|
"gorm.io/gorm"
|
||
|
|
"gorm.io/gorm/clause"
|
||
|
|
)
|
||
|
|
|
||
|
|
// MCPToolApprovalRepository implements interfaces.MCPToolApprovalRepository.
|
||
|
|
type MCPToolApprovalRepository struct {
|
||
|
|
db *gorm.DB
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewMCPToolApprovalRepository creates a repository backed by GORM.
|
||
|
|
func NewMCPToolApprovalRepository(db *gorm.DB) interfaces.MCPToolApprovalRepository {
|
||
|
|
return &MCPToolApprovalRepository{db: db}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ListByService returns all stored approval rows for an MCP service (may be empty).
|
||
|
|
func (r *MCPToolApprovalRepository) ListByService(ctx context.Context, tenantID uint64, serviceID string) ([]*types.MCPToolApproval, error) {
|
||
|
|
var rows []*types.MCPToolApproval
|
||
|
|
err := r.db.WithContext(ctx).
|
||
|
|
Where("tenant_id = ? AND service_id = ?", tenantID, serviceID).
|
||
|
|
Order("tool_name ASC").
|
||
|
|
Find(&rows).Error
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("list mcp tool approvals: %w", err)
|
||
|
|
}
|
||
|
|
return rows, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// IsRequired returns true when a row exists with require_approval = true.
|
||
|
|
func (r *MCPToolApprovalRepository) IsRequired(ctx context.Context, tenantID uint64, serviceID, toolName string) (bool, error) {
|
||
|
|
var row types.MCPToolApproval
|
||
|
|
err := r.db.WithContext(ctx).
|
||
|
|
Select("require_approval").
|
||
|
|
Where("tenant_id = ? AND service_id = ? AND tool_name = ?", tenantID, serviceID, toolName).
|
||
|
|
First(&row).Error
|
||
|
|
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
|
|
return false, nil
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
return false, fmt.Errorf("get mcp tool approval: %w", err)
|
||
|
|
}
|
||
|
|
return row.RequireApproval, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Upsert creates or updates the approval flag for a tool atomically.
|
||
|
|
// Uses ON CONFLICT against the (tenant_id, service_id, tool_name) unique index
|
||
|
|
// so concurrent writers don't race the prior SELECT-then-INSERT path into
|
||
|
|
// duplicate-key 500s.
|
||
|
|
func (r *MCPToolApprovalRepository) Upsert(ctx context.Context, row *types.MCPToolApproval) error {
|
||
|
|
if row == nil {
|
||
|
|
return errors.New("row is nil")
|
||
|
|
}
|
||
|
|
if row.ID == "" {
|
||
|
|
row.ID = uuid.New().String()
|
||
|
|
}
|
||
|
|
now := time.Now()
|
||
|
|
row.UpdatedAt = now
|
||
|
|
if row.CreatedAt.IsZero() {
|
||
|
|
row.CreatedAt = now
|
||
|
|
}
|
||
|
|
err := r.db.WithContext(ctx).Clauses(clause.OnConflict{
|
||
|
|
Columns: []clause.Column{
|
||
|
|
{Name: "tenant_id"},
|
||
|
|
{Name: "service_id"},
|
||
|
|
{Name: "tool_name"},
|
||
|
|
},
|
||
|
|
DoUpdates: clause.Assignments(map[string]interface{}{
|
||
|
|
"require_approval": row.RequireApproval,
|
||
|
|
"updated_at": now,
|
||
|
|
}),
|
||
|
|
}).Create(row).Error
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("upsert mcp tool approval: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|