package actions import ( "fmt" "jhman/model" "gorm.io/gorm" ) const batchSize = 1000 func FindEdgesByCity(tx *gorm.DB, cityId int) ([]model.Edge, error) { var edges []model.Edge err := tx.Find(&edges, "city_id = ?", cityId).Error if err != nil { return nil, fmt.Errorf("failed to find edges: %w", err) } return edges, nil } func SliceActiveEdges(tx *gorm.DB, cityId int, afterEdgeId int, count int) ([]model.Edge, error) { var edges []model.Edge err := tx.Limit(count).Find(&edges, "id > ? and city_id = ? and active = 1", afterEdgeId, cityId).Error if err != nil { return nil, fmt.Errorf("failed to find edges with offset: %w", err) } return edges, nil } func SaveEdges(tx *gorm.DB, edges []model.Edge) error { if len(edges) == 0 { return nil } // 手动区分需要创建和更新的节点,以避免 on conflict 更新的 id 占用问题 macAddrs := make([]string, len(edges)) for i, e := range edges { macAddrs[i] = e.Macaddr } var existingEdges []model.Edge if err := tx.Select("id", "macaddr").Where("macaddr IN ?", macAddrs).Find(&existingEdges).Error; err != nil { return fmt.Errorf("failed to find existing edges: %w", err) } existingEdgesMap := make(map[string]model.Edge) for _, e := range existingEdges { existingEdgesMap[e.Macaddr] = e } var toCreate, toUpdate []model.Edge for _, edge := range edges { if existing, ok := existingEdgesMap[edge.Macaddr]; ok { edge.Id = existing.Id toUpdate = append(toUpdate, edge) } else { toCreate = append(toCreate, edge) } } // 批量创建 if len(toCreate) > 0 { if err := tx.CreateInBatches(toCreate, batchSize).Error; err != nil { return fmt.Errorf("failed to create edge: %w", err) } } // 批量更新(使用临时表避免性能问题) if len(toUpdate) > 0 { // 使用临时表进行批量更新 tempTableName := "temp_edge_update" if err := tx.Exec(fmt.Sprintf("CREATE TEMPORARY TABLE %s LIKE edge", tempTableName)).Error; err != nil { return fmt.Errorf("failed to create temporary table: %w", err) } // 将待更新数据插入临时表 if err := tx.Table(tempTableName).CreateInBatches(toUpdate, batchSize).Error; err != nil { return fmt.Errorf("failed to insert into temporary table: %w", err) } // 使用 JOIN 更新主表 updateQuery := fmt.Sprintf(` UPDATE edge e JOIN %s te ON e.id = te.id SET e.macaddr = te.macaddr, e.public = te.public, e.isp = te.isp, e.single = te.single, e.sole = te.sole, e.arch = te.arch, e.online = te.online, e.city_id = te.city_id, e.active = te.active `, tempTableName) if err := tx.Exec(updateQuery).Error; err != nil { return fmt.Errorf("failed to batch update edges from temporary table: %w", err) } // 删除临时表 if err := tx.Exec(fmt.Sprintf("DROP TEMPORARY TABLE %s", tempTableName)).Error; err != nil { fmt.Printf("warning: failed to drop temporary table %s: %v\n", tempTableName, err) } } return nil } func DisableEdgesByMacs(tx *gorm.DB, macs []string) error { if len(macs) == 0 { return nil } // 分批处理MAC地址列表 for i := 0; i < len(macs); i += batchSize { end := min(i+batchSize, len(macs)) batch := macs[i:end] err := tx.Model(&model.Edge{}).Where("macaddr IN ?", batch).Update("active", false).Error if err != nil { return fmt.Errorf("failed to disable edges batch %d-%d: %w", i, end-1, err) } } return nil }