1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  //
  3  // Licensed under the Apache License, Version 2.0 (the "License");
  4  // you may not use this file except in compliance with the License.
  5  // You may obtain a copy of the License at
  6  //
  7  //     http://www.apache.org/licenses/LICENSE-2.0
  8  //
  9  // Unless required by applicable law or agreed to in writing, software
 10  // distributed under the License is distributed on an "AS IS" BASIS,
 11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  // See the License for the specific language governing permissions and
 13  // limitations under the License.
 14  //
 15  // Requirement: Any integration or derivative work must explicitly attribute
 16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  // documentation or user interface, as detailed in the NOTICE file.
 18  
 19  package database
 20  
 21  import (
 22  	"time"
 23  
 24  	"gorm.io/datatypes"
 25  	"gorm.io/gorm"
 26  )
 27  
 28  // Agent 表示一个注册的agent（GORM实现）
 29  type Agent struct {
 30  	ID           string         `gorm:"primaryKey;column:id" json:"agent_id"`
 31  	Hostname     string         `gorm:"column:hostname;not null" json:"hostname"`
 32  	IP           string         `gorm:"column:ip;not null" json:"ip"`
 33  	Version      string         `gorm:"column:version" json:"version"`
 34  	Capabilities datatypes.JSON `gorm:"column:capabilities" json:"capabilities"` // 存储为JSON
 35  	Meta         string         `gorm:"column:meta;not null" json:"meta"`
 36  	LastSeen     int64          `gorm:"column:last_seen;not null" json:"last_seen"`   // 时间戳毫秒级
 37  	CreatedAt    int64          `gorm:"column:created_at;not null" json:"created_at"` // 时间戳毫秒级
 38  	UpdatedAt    int64          `gorm:"column:updated_at;not null" json:"updated_at"` // 时间戳毫秒级
 39  	Online       bool           `gorm:"column:online;not null;default:false" json:"online"`
 40  }
 41  
 42  type AgentStore struct {
 43  	db *gorm.DB
 44  }
 45  
 46  // NewAgentStore 创建一个新的AgentStore实例
 47  func NewAgentStore(db *gorm.DB) *AgentStore {
 48  	return &AgentStore{db: db}
 49  }
 50  
 51  // Init 自动迁移agent表结构
 52  func (s *AgentStore) Init() error {
 53  	return s.db.AutoMigrate(&Agent{})
 54  }
 55  
 56  // Register 注册或更新agent信息
 57  func (s *AgentStore) Register(agent *Agent) error {
 58  	now := time.Now().UnixMilli()
 59  	agent.LastSeen = now
 60  	if agent.CreatedAt == 0 {
 61  		agent.CreatedAt = now
 62  	}
 63  	agent.UpdatedAt = now
 64  	return s.db.Save(agent).Error
 65  }
 66  
 67  // GetAgent 获取指定agent的信息
 68  func (s *AgentStore) GetAgent(id string) (*Agent, error) {
 69  	var agent Agent
 70  	err := s.db.First(&agent, "id = ?", id).Error
 71  	if err != nil {
 72  		return nil, err
 73  	}
 74  	return &agent, nil
 75  }
 76  
 77  // ListAgents 获取所有agent列表
 78  func (s *AgentStore) ListAgents() ([]*Agent, error) {
 79  	var agents []*Agent
 80  	err := s.db.Find(&agents).Error
 81  	if err != nil {
 82  		return nil, err
 83  	}
 84  	return agents, nil
 85  }
 86  
 87  // UpdateLastSeen 更新agent的最后在线时间
 88  func (s *AgentStore) UpdateLastSeen(id string) error {
 89  	now := time.Now().UnixMilli()
 90  	return s.db.Model(&Agent{}).Where("id = ?", id).Updates(map[string]interface{}{
 91  		"last_seen":  now,
 92  		"updated_at": now,
 93  	}).Error
 94  }
 95  
 96  // DeleteAgent 删除指定的agent
 97  func (s *AgentStore) DeleteAgent(id string) error {
 98  	return s.db.Delete(&Agent{}, "id = ?", id).Error
 99  }
100  
101  func (s *AgentStore) UpdateOnlineStatus(id string, online bool) error {
102  	return s.db.Model(&Agent{}).Where("id = ?", id).Update("online", online).Error
103  }