feat(cmd): add source

app/source/source.go

@@ -0,0 +1,199 @@
+package source
+
+import (
+	"errors"
+	"fmt"
+	"github.com/bcicen/jstream"
+	"github.com/fatih/color"
+	"github.com/iyear/searchx/pkg/keygen"
+	"github.com/iyear/searchx/pkg/models"
+	"github.com/iyear/searchx/pkg/storage"
+	"github.com/iyear/searchx/pkg/storage/search"
+	"github.com/mitchellh/mapstructure"
+	"log"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const (
+	keyID       = "id"
+	keyType     = "type"
+	supergroup  = "private_supergroup"
+	channel     = "private_channel"
+	typeMessage = "message"
+)
+
+type message struct {
+	ID   int         `mapstructure:"id"`
+	Type string      `mapstructure:"type"`
+	Time string      `mapstructure:"date_unixtime"`
+	From string      `mapstructure:"from_id"`
+	Text interface{} `mapstructure:"-"`
+}
+
+func Start(src, searchDriver string, searchOptions map[string]string) error {
+	if searchDriver == "" {
+		return errors.New("search driver can not be empty")
+	}
+
+	options := make(map[string]interface{})
+	if err := mapstructure.Decode(searchOptions, &options); err != nil {
+		return err
+	}
+
+	_search, err := search.New(searchDriver, options)
+	if err != nil {
+		return err
+	}
+
+	chatType, chatID, err := getChatInfo(src)
+	if err != nil {
+		return err
+	}
+
+	color.Blue("Type: %s, ID: %d\n", chatType, chatID)
+
+	start := time.Now()
+	if err = index(src, chatID, _search); err != nil {
+		return err
+	}
+	color.Blue("Index Succ... Time: %v", time.Since(start))
+
+	return nil
+
+}
+
+func index(src string, chatID int64, search storage.Search) error {
+	f, err := os.Open(src)
+	if err != nil {
+		return err
+	}
+
+	defer func(f *os.File) {
+		if err = f.Close(); err != nil {
+			log.Fatalln(err)
+		}
+	}(f)
+
+	d := jstream.NewDecoder(f, 2)
+
+	msg := message{}
+	batchSize := 50
+	items := make([]*storage.SearchItem, 0, batchSize)
+
+	for mv := range d.Stream() {
+		if mv.ValueType != jstream.Object {
+			continue
+		}
+
+		if err = mapstructure.Decode(mv.Value, &msg); err != nil {
+			for k, v := range mv.Value.(map[string]interface{}) {
+				fmt.Println(k, v)
+			}
+			return err
+		}
+
+		if msg.ID < 0 || msg.Type != typeMessage {
+			continue
+		}
+
+		//fmt.Printf("%T, %v\n", msg.Text, msg.Text)
+
+		msg.Text = mv.Value.(map[string]interface{})["text"]
+		text := ""
+
+		switch r := msg.Text.(type) {
+		case string:
+			text = r
+		case []interface{}:
+			for _, tt := range r {
+				switch t := tt.(type) {
+				case string:
+					text += t
+				case map[string]interface{}:
+					text += " " + t["text"].(string) + " "
+				}
+			}
+		}
+
+		if !strings.HasPrefix(msg.From, "user") {
+			continue
+		}
+
+		if text != "" {
+			items = append(items, &storage.SearchItem{
+				ID: keygen.SearchMsgID(chatID, msg.ID),
+				Data: &models.SearchMsg{
+					ID:     strconv.Itoa(msg.ID),
+					Chat:   strconv.FormatInt(chatID, 10),
+					Text:   text,
+					Sender: strings.TrimPrefix(msg.From, "user"),
+					Date:   msg.Time,
+				},
+			})
+		}
+
+		if len(items) == batchSize {
+			if err = search.Index(items); err != nil {
+				return err
+			}
+			items = make([]*storage.SearchItem, 0, batchSize)
+		}
+	}
+
+	if len(items) > 0 {
+		if err = search.Index(items); err != nil {
+			return err
+		}
+	}
+
+	return nil
+
+}
+
+func getChatInfo(src string) (string, int64, error) {
+	f, err := os.Open(src)
+	if err != nil {
+		return "", 0, err
+	}
+	defer func(f *os.File) {
+		if err = f.Close(); err != nil {
+			log.Fatalln(err)
+		}
+	}(f)
+
+	d := jstream.NewDecoder(f, 1).EmitKV()
+
+	var chatType = ""
+	var chatID int64 = 0
+
+	for mv := range d.Stream() {
+		kv, ok := mv.Value.(jstream.KV)
+		if !ok {
+			continue
+		}
+
+		if kv.Key == keyType {
+			chatType = kv.Value.(string)
+			if chatType != supergroup && chatType != channel {
+				return "", 0, errors.New("chat type should be supergroup or channel")
+			}
+		}
+
+		if kv.Key == keyID {
+			chatID = -int64(kv.Value.(float64)) - 1e12
+		}
+
+		if chatType != "" && chatID != 0 {
+			break
+		}
+	}
+
+	if chatType == "" || chatID == 0 {
+		return "", 0, errors.New("can not get chat type or chat id")
+	}
+
+	return chatType, chatID, nil
+}

cmd/root.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"github.com/fatih/color"
 	"github.com/iyear/searchx/cmd/run"
+	"github.com/iyear/searchx/cmd/source"
 	"github.com/iyear/searchx/global"
 	"github.com/spf13/cobra"
 )
@@ -23,6 +24,7 @@ var cmd = &cobra.Command{
 
 func init() {
 	cmd.AddCommand(run.Cmd)
+	cmd.AddCommand(source.Cmd)
 
 	cmd.PersistentFlags().BoolVarP(&version, "version", "v", false, "check the version of pure-live")
 

cmd/source/source.go

@@ -0,0 +1,31 @@
+package source
+
+import (
+	"github.com/fatih/color"
+	"github.com/iyear/searchx/app/source"
+	"github.com/spf13/cobra"
+)
+
+var (
+	src           string
+	searchDriver  string
+	searchOptions map[string]string
+)
+
+var Cmd = &cobra.Command{
+	Use:   "source",
+	Short: "",
+	Long:  ``,
+	Run: func(cmd *cobra.Command, args []string) {
+		if err := source.Start(src, searchDriver, searchOptions); err != nil {
+			color.Red("error happens: %v", err)
+			return
+		}
+	},
+}
+
+func init() {
+	Cmd.PersistentFlags().StringVarP(&src, "file", "f", "result.json", "The path to the JSON file exported by Telegram")
+	Cmd.PersistentFlags().StringVarP(&searchDriver, "driver", "d", "", "Used search driver")
+	Cmd.PersistentFlags().StringToStringVarP(&searchOptions, "options", "o", make(map[string]string), "")
+}