From fbfd5c5f4255856ce02c523b41dd47c89c76bfac Mon Sep 17 00:00:00 2001 From: dailz Date: Fri, 10 Apr 2026 08:39:30 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E5=92=8C=E5=AD=98=E5=82=A8=E5=B1=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - model: JobTemplate、SubmitJobRequest、JobHistoryQuery 等模型定义 - store: NewGormDB MySQL 连接池，使用 zap 日志替代 GORM 默认日志 - store: TemplateStore CRUD 操作，支持 GORM AutoMigrate - NewGormDB 接受 gormLevel 参数，由上层传入配置值 - 完整 TDD 测试覆盖 Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- internal/model/cluster.go | 23 ++ internal/model/job.go | 47 ++++ internal/model/template.go | 45 ++++ .../001_create_job_templates.down.sql | 1 + .../001_create_job_templates.up.sql | 14 ++ internal/store/mysql.go | 44 ++++ internal/store/mysql_test.go | 14 ++ internal/store/template_store.go | 113 ++++++++++ internal/store/template_store_test.go | 205 ++++++++++++++++++ 9 files changed, 506 insertions(+) create mode 100644 internal/model/cluster.go create mode 100644 internal/model/job.go create mode 100644 internal/model/template.go create mode 100644 internal/store/migrations/001_create_job_templates.down.sql create mode 100644 internal/store/migrations/001_create_job_templates.up.sql create mode 100644 internal/store/mysql.go create mode 100644 internal/store/mysql_test.go create mode 100644 internal/store/template_store.go create mode 100644 internal/store/template_store_test.go diff --git a/internal/model/cluster.go b/internal/model/cluster.go new file mode 100644 index 0000000..f2214f2 --- /dev/null +++ b/internal/model/cluster.go @@ -0,0 +1,23 @@ +package model + +// NodeResponse is the simplified API response for a node. +type NodeResponse struct { + Name string `json:"name"` + State []string `json:"state"` + CPUs int32 `json:"cpus"` + RealMemory int64 `json:"real_memory"` + AllocMem int64 `json:"alloc_memory,omitempty"` + Arch string `json:"architecture,omitempty"` + OS string `json:"operating_system,omitempty"` +} + +// PartitionResponse is the simplified API response for a partition. +type PartitionResponse struct { + Name string `json:"name"` + State []string `json:"state"` + Nodes string `json:"nodes,omitempty"` + TotalCPUs int32 `json:"total_cpus,omitempty"` + TotalNodes int32 `json:"total_nodes,omitempty"` + MaxTime string `json:"max_time,omitempty"` + Default bool `json:"default,omitempty"` +} diff --git a/internal/model/job.go b/internal/model/job.go new file mode 100644 index 0000000..f42e633 --- /dev/null +++ b/internal/model/job.go @@ -0,0 +1,47 @@ +package model + +// SubmitJobRequest is the API request for submitting a job. +type SubmitJobRequest struct { + Script string `json:"script" binding:"required"` + Partition string `json:"partition,omitempty"` + QOS string `json:"qos,omitempty"` + CPUs int32 `json:"cpus,omitempty"` + Memory string `json:"memory,omitempty"` + TimeLimit string `json:"time_limit,omitempty"` + JobName string `json:"job_name,omitempty"` + Environment map[string]string `json:"environment,omitempty"` +} + +// JobResponse is the simplified API response for a job. +type JobResponse struct { + JobID int32 `json:"job_id"` + Name string `json:"name"` + State []string `json:"job_state"` + Partition string `json:"partition"` + SubmitTime *int64 `json:"submit_time,omitempty"` + StartTime *int64 `json:"start_time,omitempty"` + EndTime *int64 `json:"end_time,omitempty"` + ExitCode *int32 `json:"exit_code,omitempty"` + Nodes string `json:"nodes,omitempty"` +} + +// JobListResponse is the paginated response for job listings. +type JobListResponse struct { + Jobs []JobResponse `json:"jobs"` + Total int `json:"total"` + Page int `json:"page"` + PageSize int `json:"page_size"` +} + +// JobHistoryQuery contains query parameters for job history. +type JobHistoryQuery struct { + Users string `form:"users" json:"users,omitempty"` + StartTime string `form:"start_time" json:"start_time,omitempty"` + EndTime string `form:"end_time" json:"end_time,omitempty"` + Account string `form:"account" json:"account,omitempty"` + Partition string `form:"partition" json:"partition,omitempty"` + State string `form:"state" json:"state,omitempty"` + JobName string `form:"job_name" json:"job_name,omitempty"` + Page int `form:"page,default=1" json:"page,omitempty"` + PageSize int `form:"page_size,default=20" json:"page_size,omitempty"` +} diff --git a/internal/model/template.go b/internal/model/template.go new file mode 100644 index 0000000..5e8d2e5 --- /dev/null +++ b/internal/model/template.go @@ -0,0 +1,45 @@ +package model + +import "time" + +// JobTemplate represents a saved job template. +type JobTemplate struct { + ID int64 `json:"id" gorm:"primaryKey;autoIncrement"` + Name string `json:"name" gorm:"uniqueIndex;size:255;not null"` + Description string `json:"description,omitempty" gorm:"type:text"` + Script string `json:"script" gorm:"type:text;not null"` + Partition string `json:"partition,omitempty" gorm:"size:255"` + QOS string `json:"qos,omitempty" gorm:"column:qos;size:255"` + CPUs int `json:"cpus,omitempty" gorm:"column:cpus"` + Memory string `json:"memory,omitempty" gorm:"size:50"` + TimeLimit string `json:"time_limit,omitempty" gorm:"column:time_limit;size:50"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// TableName specifies the database table name for GORM. +func (JobTemplate) TableName() string { return "job_templates" } + +// CreateTemplateRequest is the API request for creating a template. +type CreateTemplateRequest struct { + Name string `json:"name" binding:"required"` + Description string `json:"description,omitempty"` + Script string `json:"script" binding:"required"` + Partition string `json:"partition,omitempty"` + QOS string `json:"qos,omitempty"` + CPUs int `json:"cpus,omitempty"` + Memory string `json:"memory,omitempty"` + TimeLimit string `json:"time_limit,omitempty"` +} + +// UpdateTemplateRequest is the API request for updating a template. +type UpdateTemplateRequest struct { + Name string `json:"name,omitempty"` + Description string `json:"description,omitempty"` + Script string `json:"script,omitempty"` + Partition string `json:"partition,omitempty"` + QOS string `json:"qos,omitempty"` + CPUs int `json:"cpus,omitempty"` + Memory string `json:"memory,omitempty"` + TimeLimit string `json:"time_limit,omitempty"` +} diff --git a/internal/store/migrations/001_create_job_templates.down.sql b/internal/store/migrations/001_create_job_templates.down.sql new file mode 100644 index 0000000..9626210 --- /dev/null +++ b/internal/store/migrations/001_create_job_templates.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS job_templates; diff --git a/internal/store/migrations/001_create_job_templates.up.sql b/internal/store/migrations/001_create_job_templates.up.sql new file mode 100644 index 0000000..128ff1d --- /dev/null +++ b/internal/store/migrations/001_create_job_templates.up.sql @@ -0,0 +1,14 @@ +CREATE TABLE IF NOT EXISTS job_templates ( + id BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY, + name VARCHAR(255) NOT NULL, + description TEXT, + script TEXT NOT NULL, + partition VARCHAR(255), + qos VARCHAR(255), + cpus INT UNSIGNED, + memory VARCHAR(50), + time_limit VARCHAR(50), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + UNIQUE KEY idx_name (name) +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4; diff --git a/internal/store/mysql.go b/internal/store/mysql.go new file mode 100644 index 0000000..2fe19f8 --- /dev/null +++ b/internal/store/mysql.go @@ -0,0 +1,44 @@ +package store + +import ( + "fmt" + "time" + + "gcy_hpc_server/internal/logger" + "gcy_hpc_server/internal/model" + + "go.uber.org/zap" + "gorm.io/driver/mysql" + "gorm.io/gorm" +) + +// NewGormDB opens a GORM MySQL connection with sensible defaults. +func NewGormDB(dsn string, zapLogger *zap.Logger, gormLevel string) (*gorm.DB, error) { + gormCfg := &gorm.Config{ + Logger: logger.NewGormLogger(zapLogger, gormLevel), + } + db, err := gorm.Open(mysql.Open(dsn), gormCfg) + if err != nil { + return nil, fmt.Errorf("failed to open gorm mysql: %w", err) + } + + sqlDB, err := db.DB() + if err != nil { + return nil, fmt.Errorf("failed to get underlying sql.DB: %w", err) + } + + sqlDB.SetMaxOpenConns(25) + sqlDB.SetMaxIdleConns(5) + sqlDB.SetConnMaxLifetime(5 * time.Minute) + + if err := sqlDB.Ping(); err != nil { + return nil, fmt.Errorf("failed to ping mysql: %w", err) + } + + return db, nil +} + +// AutoMigrate runs GORM auto-migration for all models. +func AutoMigrate(db *gorm.DB) error { + return db.AutoMigrate(&model.JobTemplate{}) +} diff --git a/internal/store/mysql_test.go b/internal/store/mysql_test.go new file mode 100644 index 0000000..5a05e91 --- /dev/null +++ b/internal/store/mysql_test.go @@ -0,0 +1,14 @@ +package store + +import ( + "testing" + + "go.uber.org/zap" +) + +func TestNewGormDBInvalidDSN(t *testing.T) { + _, err := NewGormDB("invalid:dsn@tcp(localhost:99999)/nonexistent?parseTime=true", zap.NewNop(), "warn") + if err == nil { + t.Fatal("expected error for invalid DSN, got nil") + } +} diff --git a/internal/store/template_store.go b/internal/store/template_store.go new file mode 100644 index 0000000..8d196ff --- /dev/null +++ b/internal/store/template_store.go @@ -0,0 +1,113 @@ +package store + +import ( + "context" + "errors" + + "gorm.io/gorm" + + "gcy_hpc_server/internal/model" +) + +// TemplateStore provides CRUD operations for job templates via GORM. +type TemplateStore struct { + db *gorm.DB +} + +// NewTemplateStore creates a new TemplateStore. +func NewTemplateStore(db *gorm.DB) *TemplateStore { + return &TemplateStore{db: db} +} + +// List returns a paginated list of job templates and the total count. +func (s *TemplateStore) List(ctx context.Context, page, pageSize int) ([]model.JobTemplate, int, error) { + var templates []model.JobTemplate + var total int64 + + if err := s.db.WithContext(ctx).Model(&model.JobTemplate{}).Count(&total).Error; err != nil { + return nil, 0, err + } + + offset := (page - 1) * pageSize + if err := s.db.WithContext(ctx).Order("id DESC").Limit(pageSize).Offset(offset).Find(&templates).Error; err != nil { + return nil, 0, err + } + + return templates, int(total), nil +} + +// GetByID returns a single job template by ID. Returns nil, nil when not found. +func (s *TemplateStore) GetByID(ctx context.Context, id int64) (*model.JobTemplate, error) { + var t model.JobTemplate + err := s.db.WithContext(ctx).First(&t, id).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, err + } + return &t, nil +} + +// Create inserts a new job template and returns the generated ID. +func (s *TemplateStore) Create(ctx context.Context, req *model.CreateTemplateRequest) (int64, error) { + t := &model.JobTemplate{ + Name: req.Name, + Description: req.Description, + Script: req.Script, + Partition: req.Partition, + QOS: req.QOS, + CPUs: req.CPUs, + Memory: req.Memory, + TimeLimit: req.TimeLimit, + } + if err := s.db.WithContext(ctx).Create(t).Error; err != nil { + return 0, err + } + return t.ID, nil +} + +// Update modifies an existing job template. Only non-empty/non-zero fields are updated. +func (s *TemplateStore) Update(ctx context.Context, id int64, req *model.UpdateTemplateRequest) error { + updates := map[string]interface{}{} + if req.Name != "" { + updates["name"] = req.Name + } + if req.Description != "" { + updates["description"] = req.Description + } + if req.Script != "" { + updates["script"] = req.Script + } + if req.Partition != "" { + updates["partition"] = req.Partition + } + if req.QOS != "" { + updates["qos"] = req.QOS + } + if req.CPUs > 0 { + updates["cpus"] = req.CPUs + } + if req.Memory != "" { + updates["memory"] = req.Memory + } + if req.TimeLimit != "" { + updates["time_limit"] = req.TimeLimit + } + + if len(updates) == 0 { + return nil // nothing to update + } + + result := s.db.WithContext(ctx).Model(&model.JobTemplate{}).Where("id = ?", id).Updates(updates) + return result.Error +} + +// Delete removes a job template by ID. Idempotent — returns nil even if the row doesn't exist. +func (s *TemplateStore) Delete(ctx context.Context, id int64) error { + result := s.db.WithContext(ctx).Delete(&model.JobTemplate{}, id) + if result.Error != nil { + return result.Error + } + return nil +} diff --git a/internal/store/template_store_test.go b/internal/store/template_store_test.go new file mode 100644 index 0000000..d0f9703 --- /dev/null +++ b/internal/store/template_store_test.go @@ -0,0 +1,205 @@ +package store + +import ( + "context" + "testing" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "gcy_hpc_server/internal/model" +) + +func newTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + Logger: logger.Default.LogMode(logger.Silent), + }) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&model.JobTemplate{}); err != nil { + t.Fatalf("auto migrate: %v", err) + } + return db +} + +func TestTemplateStore_List(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + s.Create(context.Background(), &model.CreateTemplateRequest{Name: "job-1", Script: "echo 1"}) + s.Create(context.Background(), &model.CreateTemplateRequest{Name: "job-2", Script: "echo 2"}) + + templates, total, err := s.List(context.Background(), 1, 10) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if total != 2 { + t.Errorf("total = %d, want 2", total) + } + if len(templates) != 2 { + t.Fatalf("len(templates) = %d, want 2", len(templates)) + } + // DESC order, so job-2 is first + if templates[0].Name != "job-2" { + t.Errorf("templates[0].Name = %q, want %q", templates[0].Name, "job-2") + } +} + +func TestTemplateStore_List_Page2(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + for i := 0; i < 15; i++ { + s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "job-" + string(rune('A'+i)), Script: "echo", + }) + } + + templates, total, err := s.List(context.Background(), 2, 10) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if total != 15 { + t.Errorf("total = %d, want 15", total) + } + if len(templates) != 5 { + t.Fatalf("len(templates) = %d, want 5", len(templates)) + } +} + +func TestTemplateStore_GetByID(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + id, _ := s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "test-job", Script: "echo hi", Partition: "batch", QOS: "normal", CPUs: 2, Memory: "4G", + }) + + tpl, err := s.GetByID(context.Background(), id) + if err != nil { + t.Fatalf("GetByID() error = %v", err) + } + if tpl == nil { + t.Fatal("GetByID() returned nil") + } + if tpl.Name != "test-job" { + t.Errorf("Name = %q, want %q", tpl.Name, "test-job") + } + if tpl.CPUs != 2 { + t.Errorf("CPUs = %d, want 2", tpl.CPUs) + } +} + +func TestTemplateStore_GetByID_NotFound(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + tpl, err := s.GetByID(context.Background(), 999) + if err != nil { + t.Fatalf("GetByID() error = %v, want nil", err) + } + if tpl != nil { + t.Fatal("GetByID() should return nil for not found") + } +} + +func TestTemplateStore_Create(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + id, err := s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "new-job", Script: "echo", Partition: "gpu", + }) + if err != nil { + t.Fatalf("Create() error = %v", err) + } + if id == 0 { + t.Fatal("Create() returned id=0") + } +} + +func TestTemplateStore_Update(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + id, _ := s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "old", Script: "echo", + }) + + err := s.Update(context.Background(), id, &model.UpdateTemplateRequest{ + Name: "updated", + Script: "echo new", + CPUs: 8, + }) + if err != nil { + t.Fatalf("Update() error = %v", err) + } + + tpl, _ := s.GetByID(context.Background(), id) + if tpl.Name != "updated" { + t.Errorf("Name = %q, want %q", tpl.Name, "updated") + } + if tpl.CPUs != 8 { + t.Errorf("CPUs = %d, want 8", tpl.CPUs) + } +} + +func TestTemplateStore_Update_Partial(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + id, _ := s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "original", Script: "echo orig", Partition: "batch", + }) + + err := s.Update(context.Background(), id, &model.UpdateTemplateRequest{ + Name: "renamed", + }) + if err != nil { + t.Fatalf("Update() error = %v", err) + } + + tpl, _ := s.GetByID(context.Background(), id) + if tpl.Name != "renamed" { + t.Errorf("Name = %q, want %q", tpl.Name, "renamed") + } + // Script and Partition should be unchanged + if tpl.Script != "echo orig" { + t.Errorf("Script = %q, want %q", tpl.Script, "echo orig") + } + if tpl.Partition != "batch" { + t.Errorf("Partition = %q, want %q", tpl.Partition, "batch") + } +} + +func TestTemplateStore_Delete(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + id, _ := s.Create(context.Background(), &model.CreateTemplateRequest{ + Name: "to-delete", Script: "echo", + }) + + err := s.Delete(context.Background(), id) + if err != nil { + t.Fatalf("Delete() error = %v", err) + } + + tpl, _ := s.GetByID(context.Background(), id) + if tpl != nil { + t.Fatal("Delete() did not remove the record") + } +} + +func TestTemplateStore_Delete_NotFound(t *testing.T) { + db := newTestDB(t) + s := NewTemplateStore(db) + + err := s.Delete(context.Background(), 999) + if err != nil { + t.Fatalf("Delete() should not error for non-existent record, got: %v", err) + } +}