我的第一个golang command脚本

场景

维护数据库的树形结构的pid集合

具体代码

/**
  @author: Administrator
  @since: 2022/9/29
  @desc: 处理category的 脚本暂时不写
**/

package main

import (
    "context"
    "database/sql"
    "fmt"
    "strconv"
    "strings"
    "time"

    "github.com/Masterminds/squirrel"
    "github.com/zeromicro/go-zero/core/stores/builder"
    "github.com/zeromicro/go-zero/core/stores/sqlx"
)

var (
    categoryFieldNames = builder.RawFieldNames(&Category{})
    categoryRows       = strings.Join(categoryFieldNames, ",")
)

type (
    Fixer struct {
        ctx        context.Context
        sqlConn    sqlx.SqlConn
        pidListMap map[int]string
        pidList    []int
    }

    Category struct {
        Id         int          `db:"id"`
        Pid        int          `db:"pid"` // 上级id
        Type       string       `db:"type"`
        Name       string       `db:"name"` // 类型名称
        Icon       string       `db:"icon"` // icon 路径
        Sort       int64        `db:"sort"` // 排序
        Ext1       string       `db:"ext1"`
        Ext2       string       `db:"ext2"`
        Ext3       string       `db:"ext3"`
        PidList    string       `db:"pid_list"`
        CreateTime time.Time    `db:"create_time"`
        UpdateTime time.Time    `db:"update_time"`
        DeleteTime sql.NullTime `db:"delete_time"`
    }
)

func newFixer() *Fixer {
    sqlConn := sqlx.NewMysql("iot:STfbeXiRp4Pp2xYZ@tcp(175.128.51.74:1111)/iot?charset=utf8mb4&parseTime=true&loc=Asia%2FShanghai")
    pidlist := make([]int, 0)
    pidListMap := make(map[int]string, 0)
    return &Fixer{sqlConn: sqlConn, ctx: context.Background(), pidListMap: pidListMap, pidList: pidlist}
}

func main() {
    fixer := newFixer()
    fixer.fixPidList()
}

// 重置pidList 保留上一轮
func (f *Fixer) resetPidList(list []Category) {

    f.pidListMap = make(map[int]string, 0)
    f.pidList = make([]int, 0)

    for _, category := range list {
        f.pidListMap[category.Id] = category.PidList
        f.pidList = append(f.pidList, category.Id)
    }
}

func (f *Fixer) fixPidList() error {

    var list []Category

    if len(f.pidListMap) == 0 {
        sqlBuilderList := squirrel.Select(categoryRows).From("category").Where(" pid = 0")
        listQuery, listvalues, _ := sqlBuilderList.OrderBy("sort DESC").ToSql()
        // 查询列表
        f.sqlConn.QueryRowsCtx(f.ctx, &list, listQuery, listvalues...)
        // 装载 id  和 k-v
        f.resetPidList(list)
        return f.fixPidList()
    }

    sqlBuilderList := squirrel.Select(categoryRows).From("category").Where(squirrel.Eq{"pid": f.pidList})
    listQuery, listvalues, _ := sqlBuilderList.OrderBy("sort DESC").ToSql()
    // 查询列表
    f.sqlConn.QueryRowsCtx(f.ctx, &list, listQuery, listvalues...)

    if len(list) <= 0 {
        fmt.Sprintf("全部处理完毕")
        return nil
    }
    // 处理数据
    for i, category := range list {

        if category.Pid == 0 {
            continue
        }

        pidList := fmt.Sprintf("%s,%s", f.pidListMap[category.Pid], strconv.Itoa(category.Pid))

        list[i].PidList = pidList

        UpdateQuery, Updatevalues, _ := squirrel.Update("category").Set("pid_list", pidList).Where("id = ?", category.Id).ToSql()
        _, err := f.sqlConn.ExecCtx(f.ctx, UpdateQuery, Updatevalues...)
        if err != nil {
            fmt.Println(err)
            panic(fmt.Sprintf("修改 %s 出现错误 sql %s ", category.Id, UpdateQuery))
        }
    }

    f.resetPidList(list)
    return f.fixPidList()
}