125 lines
3.4 KiB
Go
125 lines
3.4 KiB
Go
package actions
|
|
|
|
import (
|
|
"fmt"
|
|
"zzman/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
|
|
}
|