browserSpider/program/spider/cli/counter/main.go

// 数据统计
package main

import (
	"encoding/csv"
	"flag"
	"os"
	"sort"
	"strconv"

	"database/sql"
	"log"

	_ "github.com/mattn/go-sqlite3"
)

type (
	Folder struct {
		Id               int
		Name             string
		ParentName       string
		ParentParentName string
		ParentId         int
		Count            int
	}
	Cities []*Folder
)

var (
	dbfile  = flag.String("db", "/Users/taozhang/Downloads/cloudreve.db", "数据文件")
	outfile = flag.String("of", "./结果.csv", "结果输出文件")
	db      *sql.DB
)

// init
func init() {
	flag.Parse()
	var err error
	db, err = sql.Open("sqlite3", *dbfile)
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	defer db.Close()
	// 执行查询，加载目录
	rows, err := db.Query("SELECT id, name,parent_id FROM folders")
	if err != nil {
		log.Fatal(err)
	}
	allCities := map[int]*Folder{}
	// 先拿到所有目录
	for rows.Next() {
		var id, _pid int
		var parentId sql.NullInt32
		var name string
		err := rows.Scan(&id, &name, &parentId)
		if err != nil {
			log.Fatal(err)
		}
		_pid = int(parentId.Int32)
		allCities[id] = &Folder{
			Id:       id,
			Name:     name,
			ParentId: _pid,
		}
	}
	rows.Close()
	//TODO 需要筛选的目录，只筛选政府类的市，但是要补上父级，爷爷级节点名称，便于统计
	selectCities := map[int]*Folder{}
	for _, v := range allCities {
		if v.Name == "政府" {
			if pv, ok := allCities[v.ParentId]; ok {
				//给pv也补上父级关系
				if ppv, ok := allCities[pv.ParentId]; ok {
					pv.ParentName = ppv.Name
					//补上爷爷级节点
					if pppv, ok := allCities[ppv.ParentId]; ok {
						pv.ParentParentName = pppv.Name
					}
				}
				selectCities[v.Id] = pv
			}
		}
	}
	//TODO 计算所有目录的文件数量
	rows, err = db.Query("select folder_id,count(id) as fs from files group by folder_id;")
	if err != nil {
		log.Fatal(err)
	}
	selectCitiesArray := make(Cities, 0, 0)
	for rows.Next() {
		var id int
		var filesCount int
		err := rows.Scan(&id, &filesCount)
		if err != nil {
			log.Fatal(err)
		}
		if v, ok := selectCities[id]; ok {
			v.Count = filesCount
			selectCitiesArray = append(selectCitiesArray, v)
		}
	}
	rows.Close()
	//TODO 正向排序
	sort.Slice(selectCitiesArray, func(i, j int) bool { return selectCitiesArray[i].Count < selectCitiesArray[j].Count })
	//TODO 结果输出到CSV文件
	fo, err := os.Create(*outfile)
	if err != nil {
		log.Fatal(err)
	}
	writer := csv.NewWriter(fo)
	writer.Write([]string{"年", "省", "城市", "组合筛选KEY", "文件数"})
	for _, v := range selectCitiesArray {
		writer.Write([]string{v.ParentParentName, v.ParentName, v.Name,
			v.ParentParentName + v.ParentName + v.Name, strconv.Itoa(v.Count)})
		log.Println(v.Id, v.ParentParentName, v.ParentName, v.Name, v.Count)
	}
	writer.Flush()
	fo.Close()
}