// /* // 存储; // 暂时提供3种存储方式: // 1)控制台 // 2) MongoDb // 3) 文件 // */ package db // import ( // "bytes" // "context" // "encoding/json" // "fmt" // "io" // "log" // "net/http" // "time" // "github.com/bmaupin/go-epub" // "github.com/xuri/excelize/v2" // "go.mongodb.org/mongo-driver/mongo" // "go.mongodb.org/mongo-driver/mongo/options" // ) // type ( // //Storage // Storage interface { // //保存数据 // Save(siteName, spiderCode, siteChannelName, siteChannelUrl string, data map[string]interface{}) error // //SaveUp(query string, data map[string]interface{}) error // } // //MongoStorage // MongoStorage struct { // c *mongo.Collection // } // //调试模式下的数据存储 // ConsoleStorage struct { // } // //FileStorage // FileStorage struct { // f io.WriteCloser // } // //企业微信机器人消息 // WeChatMsgStorage struct { // RobotURL string // } // //输出到Excel文件中 // ExcelStorage struct { // FieldMapping map[string]string // ExcelFile *excelize.File // RowIndex int // SheetName string // FilePath string // } // //输出到EPUB文档 // EpubStorage struct { // EpubFilePath string // EpubFile *epub.Epub // } // ) // // NewMongoStorage // func NewMongoStorage(mongoURI, db, collection string) *MongoStorage { // client, err := mongo.Connect(context.TODO(), options.Client().ApplyURI(mongoURI)) // if err != nil { // panic(err) // } // return &MongoStorage{ // c: client.Database(db).Collection(collection), // } // } // // NewDebugerStorage // func NewConsoleStorage() *ConsoleStorage { // return &ConsoleStorage{} // } // // NewFileStorage // func NewFileStorage(f io.WriteCloser) *FileStorage { // return &FileStorage{f: f} // } // // NewWeChatMsgStorage // func NewWeChatMsgStorage(robotUrl string) *WeChatMsgStorage { // return &WeChatMsgStorage{robotUrl} // } // // Save // func (ms *MongoStorage) Save(siteName, spiderCode, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // data["site"] = siteName // data["spidercode"] = spiderCode // data["channel"] = siteChannelName // data["channelurl"] = siteChannelUrl // data["comeintime"] = time.Now().Add(time.Duration(1) * time.Hour).Unix() // t, _ := time.Parse("2006-01-02 15:04:05", fmt.Sprint(data["publishtime"])) // if t.Unix() > 0 { // data["l_np_publishtime"] = t.Unix() // } // _, err := ms.c.InsertOne(context.TODO(), data) // return err // } // func (ms *MongoStorage) SaveUp(query map[string]interface{}, data map[string]interface{}) error { // _, err := ms.c.UpdateOne(context.TODO(), query, data) // return err // } // // Save // func (ds *ConsoleStorage) Save(spiderCode, siteName, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // data["spider_code"] = spiderCode // data["site_name"] = siteName // data["site_channel_name"] = siteChannelName // data["site_channel_url"] = siteChannelUrl // fmt.Println(" data=") // for k, v := range data { // fmt.Printf("\t%s = %v \n", k, v) // } // return nil // } // // Save // func (fs *FileStorage) Save(spiderCode, siteName, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // data["spider_code"] = spiderCode // data["site_name"] = siteName // data["site_channel_name"] = siteChannelName // data["site_channel_url"] = siteChannelUrl // buf := new(bytes.Buffer) // json.NewEncoder(buf).Encode(data) // fs.f.Write([]byte(",")) // fs.f.Write(buf.Bytes()) // fs.f.Write([]byte("\n")) // return nil // } // // Save // func (ws *WeChatMsgStorage) Save(spiderCode, siteName, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // msg := fmt.Sprintf("%v", data) // body := struct { // MsgType string `json:"msgtype"` // Text struct { // Content string `json:"content"` // } `json:"text"` // }{ // MsgType: "text", // Text: struct { // Content string `json:"content"` // }{msg}, // } // bs, err := json.Marshal(body) // if err != nil { // return err // } // client := new(http.Client) // req, err := http.NewRequest("POST", ws.RobotURL, bytes.NewReader(bs)) // if err != nil { // return err // } // req.Header.Set("Content-Type", "application/json") // resp, err := client.Do(req) // if err != nil { // return err // } // return resp.Body.Close() // } // // NewExcelStorage // func NewExcelStorage(filepath string) *ExcelStorage { // f, err := excelize.OpenFile(filepath) // if err != nil { // panic(err) // } // sheetName := f.GetSheetName(0) // //第一行是表头 // rows, err := f.GetRows(sheetName) // if err != nil { // panic(err) // } // colnames := "ABCDEFGHIJKLMNOPQRSTUVWXYZ" // mapping := map[string]string{} // for i, title := range rows[0] { // mapping[title] = string(colnames[i]) // } // return &ExcelStorage{ // FieldMapping: mapping, // ExcelFile: f, // RowIndex: 1, // SheetName: sheetName, // FilePath: filepath, // } // } // // 保存数据 // func (es *ExcelStorage) Save(spiderCode, siteName, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // es.RowIndex += 1 // for k, v := range data { // if col, ok := es.FieldMapping[k]; ok { // cellName := fmt.Sprintf("%s%d", col, es.RowIndex) // value, ok := v.(string) // //fmt.Println(" write::", cellName, value) // if ok { // es.ExcelFile.SetCellStr(es.SheetName, cellName, value) // } // } // } // es.ExcelFile.Save() // return nil // } // // 电子书,存储 // func NewEpubStorage(filepath string) *EpubStorage { // output := epub.NewEpub("") // return &EpubStorage{ // EpubFilePath: filepath, // EpubFile: output, // } // } // // 保存数据 // func (es *EpubStorage) Save(spiderCode, siteName, siteChannelName, siteChannelUrl string, data map[string]interface{}) error { // title, _ := data["title"].(string) // body, _ := data["body"].(string) // index, _ := data["index"].(int64) // filename := fmt.Sprintf("%08d.xhtml", index) // if title == "init" { // bookName, _ := data["book_name"].(string) // author, _ := data["book_author"].(string) // log.Println("初始化电子书", bookName, author) // es.EpubFile.SetTitle(bookName) // es.EpubFile.SetAuthor(author) // } else if title == "save2epubfile" { // log.Println("开始写入epubfile", es.EpubFilePath) // es.EpubFile.Write(es.EpubFilePath) // } else { // es.EpubFile.AddSection(body, title, filename, "") // } // return nil // }