Create & Init Project...

app/service/main/riot-search/dao/BUILD

@@ -0,0 +1,58 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_library",
+    "go_test",
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "dao.go",
+        "indexer.go",
+        "searcher.go",
+    ],
+    importpath = "go-common/app/service/main/riot-search/dao",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//app/service/main/riot-search/conf:go_default_library",
+        "//app/service/main/riot-search/model:go_default_library",
+        "//library/database/sql:go_default_library",
+        "//library/log:go_default_library",
+        "//vendor/github.com/go-ego/riot:go_default_library",
+        "//vendor/github.com/go-ego/riot/types:go_default_library",
+    ],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)
+
+go_test(
+    name = "go_default_test",
+    srcs = [
+        "dao_test.go",
+        "indexer_test.go",
+        "searcher_test.go",
+    ],
+    embed = [":go_default_library"],
+    rundir = ".",
+    tags = ["automanaged"],
+    deps = [
+        "//app/service/main/riot-search/conf:go_default_library",
+        "//app/service/main/riot-search/model:go_default_library",
+        "//vendor/github.com/smartystreets/goconvey/convey:go_default_library",
+    ],
+)

app/service/main/riot-search/dao/dao.go

@@ -0,0 +1,99 @@
+package dao
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"go-common/app/service/main/riot-search/conf"
+	"go-common/app/service/main/riot-search/model"
+	"go-common/library/database/sql"
+	"go-common/library/log"
+
+	"github.com/go-ego/riot"
+	"github.com/go-ego/riot/types"
+)
+
+// 过审的增量数据
+var _selIncrement = "SELECT id, title from archive where mtime>? and mtime<=?"
+
+// Dao dao
+type Dao struct {
+	c        *conf.Config
+	searcher *riot.Engine
+	db       *sql.DB
+}
+
+// New init mysql db
+func New(c *conf.Config) (dao *Dao) {
+	validateConfig(c.Riot)
+	if c.UT {
+		dao = &Dao{
+			c:        c,
+			searcher: &riot.Engine{},
+			db:       sql.NewMySQL(c.Mysql),
+		}
+		dao.searcher.Init(types.EngineOpts{})
+		return
+	}
+	dao = &Dao{
+		c:        c,
+		searcher: &riot.Engine{},
+		// db
+		db: sql.NewMySQL(c.Mysql),
+	}
+	dao.searcher.Init(types.EngineOpts{
+		GseDict:       c.Riot.Dict,
+		StopTokenFile: c.Riot.StopToken,
+		NumShards:     c.Riot.NumShards,
+		IndexerOpts: &types.IndexerOpts{
+			IndexType:    types.FrequenciesIndex,
+			DocCacheSize: 5000,
+		},
+	})
+	return
+}
+
+func validateConfig(conf *conf.RiotConfig) {
+	if conf.Dict == "" || conf.StopToken == "" {
+		panic("must provide a dict and stop_token file")
+	}
+	if conf.FlushTime <= 0 {
+		panic("flush time must larger than 0")
+	}
+}
+
+// Close close the resource.
+func (d *Dao) Close() {
+	d.db.Close()
+}
+
+// Ping dao ping
+func (d *Dao) Ping(c context.Context) error {
+	return d.db.Ping(c)
+}
+
+// IncrementBackup select mtime>now-24h data
+func (d *Dao) IncrementBackup(c context.Context, stime, etime time.Time) (docs []*model.Document, err error) {
+	var states []int
+	for k := range model.PubStates.LegalStates {
+		states = append(states, k)
+	}
+	query := _selIncrement + " and state in (" + strings.Trim(strings.Join(strings.Split(fmt.Sprint(states), " "), ","), "[]") + ")" + " order by id asc"
+	rows, err := d.db.Query(c, query, stime, etime)
+	log.Info("exec query(%s) args(stime:%v, etime:%v)", query, stime, etime)
+	if err != nil {
+		return
+	}
+	defer rows.Close()
+	for rows.Next() {
+		doc := &model.Document{}
+		if err = rows.Scan(&doc.ID, &doc.Content); err != nil {
+			return
+		}
+		docs = append(docs, doc)
+	}
+	err = rows.Err()
+	return
+}

app/service/main/riot-search/dao/dao_test.go

@@ -0,0 +1,67 @@
+package dao
+
+import (
+	"context"
+	"flag"
+	"os"
+	"testing"
+	"time"
+
+	"go-common/app/service/main/riot-search/conf"
+
+	"github.com/smartystreets/goconvey/convey"
+)
+
+var (
+	d *Dao
+)
+
+func TestMain(m *testing.M) {
+	if os.Getenv("DEPLOY_ENV") != "" {
+		flag.Set("app_id", "main.search.riot-search")
+		flag.Set("conf_token", "7cac78a7fdfe78c053879bf4dff0171b")
+		flag.Set("tree_id", "55087")
+		flag.Set("conf_version", "ut")
+		flag.Set("deploy_env", "uat")
+		flag.Set("conf_host", "config.bilibili.co")
+		flag.Set("conf_path", "/tmp")
+		flag.Set("region", "sh")
+		flag.Set("zone", "sh001")
+	} else {
+		flag.Set("conf", "../cmd/ut.toml")
+	}
+	flag.Parse()
+	if err := conf.Init(); err != nil {
+		panic(err)
+	}
+	d = New(conf.Conf)
+	os.Exit(m.Run())
+}
+
+// IncrementBackup ...
+func TestIncrementBackup(t *testing.T) {
+	convey.Convey("IncrementBackup", t, func(ctx convey.C) {
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			p1, err := d.IncrementBackup(context.Background(), time.Now(), time.Now())
+			ctx.Convey("Error should be nil", func(ctx convey.C) {
+				ctx.So(err, convey.ShouldBeNil)
+				ctx.So(p1, convey.ShouldBeNil)
+			})
+		})
+	})
+	convey.Convey("Ping", t, func(ctx convey.C) {
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			err := d.Ping(context.Background())
+			ctx.Convey("Error should be nil", func(ctx convey.C) {
+				ctx.So(err, convey.ShouldBeNil)
+			})
+		})
+	})
+	convey.Convey("Close", t, func(ctx convey.C) {
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			d.Close()
+			ctx.Convey("No return values", func(ctx convey.C) {
+			})
+		})
+	})
+}

app/service/main/riot-search/dao/indexer.go

@@ -0,0 +1,20 @@
+package dao
+
+import (
+	"github.com/go-ego/riot/types"
+)
+
+// Insert doc into index...
+func (d *Dao) Insert(id uint64, content string, forceUpdate bool) {
+	d.searcher.Index(id, types.DocData{Content: content}, forceUpdate)
+}
+
+// Flush force update data from cache to index
+func (d *Dao) Flush() {
+	d.searcher.Flush()
+}
+
+// Remove remove a doc from index
+func (d *Dao) Remove(id uint64, forceUpdate bool) {
+	d.searcher.RemoveDoc(id, forceUpdate)
+}

app/service/main/riot-search/dao/indexer_test.go

@@ -0,0 +1,46 @@
+package dao
+
+import (
+	"testing"
+
+	"github.com/smartystreets/goconvey/convey"
+)
+
+func TestDaoInsert(t *testing.T) {
+	convey.Convey("Insert", t, func(ctx convey.C) {
+		var (
+			id          = uint64(0)
+			content     = ""
+			forceUpdate bool
+		)
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			d.Insert(id, content, forceUpdate)
+			ctx.Convey("No return values", func(ctx convey.C) {
+			})
+		})
+	})
+}
+
+func TestDaoFlush(t *testing.T) {
+	convey.Convey("Flush", t, func(ctx convey.C) {
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			d.Flush()
+			ctx.Convey("No return values", func(ctx convey.C) {
+			})
+		})
+	})
+}
+
+func TestDaoRemove(t *testing.T) {
+	convey.Convey("Remove", t, func(ctx convey.C) {
+		var (
+			id          = uint64(0)
+			forceUpdate bool
+		)
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			d.Remove(id, forceUpdate)
+			ctx.Convey("No return values", func(ctx convey.C) {
+			})
+		})
+	})
+}

app/service/main/riot-search/dao/searcher.go

@@ -0,0 +1,108 @@
+package dao
+
+import (
+	"go-common/app/service/main/riot-search/model"
+
+	"github.com/go-ego/riot/types"
+)
+
+// SearchIDOnly only return aids
+func (d *Dao) SearchIDOnly(arg *model.RiotSearchReq) *model.IDsResp {
+	if arg.Keyword == "" {
+		return nil
+	}
+	var docIDs map[uint64]bool
+	if len(arg.IDs) != 0 {
+		docIDs = make(map[uint64]bool, len(arg.IDs))
+		for _, id := range arg.IDs {
+			docIDs[id] = true
+		}
+	}
+	output := d.searcher.Search(types.SearchReq{
+		Text:    arg.Keyword,
+		DocIds:  docIDs,
+		Timeout: d.c.Riot.Timeout,
+		RankOpts: &types.RankOpts{
+			// 从第几条结果开始输出
+			OutputOffset: (arg.Pn - 1) * arg.Ps,
+			// 最大输出的搜索结果数，为 0 时无限制
+			MaxOutputs: arg.Ps,
+		},
+	})
+	docLength := len(output.Docs.(types.ScoredDocs))
+	tokenLength := len(output.Tokens)
+	res := &model.IDsResp{
+		IDs:    make([]uint64, docLength),
+		Tokens: make([]string, tokenLength),
+		Page: &model.Page{
+			PageNum:  arg.Pn,
+			PageSize: arg.Ps,
+			Total:    docLength,
+		},
+	}
+	for i, doc := range output.Docs.(types.ScoredDocs) {
+		res.IDs[i] = doc.DocId
+	}
+	copy(res.Tokens, output.Tokens)
+	res.IDs, res.Page.Total = uniqueIDs(res.IDs)
+	return res
+}
+
+func uniqueIDs(IDs []uint64) (uIDs []uint64, length int) {
+	m := make(map[uint64]struct{})
+	for _, ID := range IDs {
+		if _, ok := m[ID]; !ok {
+			m[ID] = struct{}{}
+			uIDs = append(uIDs, ID)
+			length++
+		}
+	}
+	return
+}
+
+// Search return archives info
+func (d *Dao) Search(arg *model.RiotSearchReq) *model.DocumentsResp {
+	if arg.Keyword == "" {
+		return nil
+	}
+	var docIDs map[uint64]bool
+	if len(arg.IDs) != 0 {
+		docIDs = make(map[uint64]bool, len(arg.IDs))
+		for _, id := range arg.IDs {
+			docIDs[id] = true
+		}
+	}
+	output := d.searcher.Search(types.SearchReq{
+		Text:    arg.Keyword,
+		DocIds:  docIDs,
+		Timeout: d.c.Riot.Timeout,
+		RankOpts: &types.RankOpts{
+			// 从第几条结果开始输出
+			OutputOffset: (arg.Pn - 1) * arg.Ps,
+			// 最大输出的搜索结果数，为 0 时无限制
+			MaxOutputs: arg.Ps,
+		},
+	})
+	docLength := len(output.Docs.(types.ScoredDocs))
+	tokenLength := len(output.Tokens)
+	res := &model.DocumentsResp{
+		Documents: make([]model.Document, docLength),
+		Tokens:    make([]string, tokenLength),
+		Page: &model.Page{
+			PageNum:  arg.Pn,
+			PageSize: arg.Ps,
+			Total:    docLength,
+		},
+	}
+	for i, doc := range output.Docs.(types.ScoredDocs) {
+		res.Documents[i].ID = doc.DocId
+		res.Documents[i].Content = doc.Content
+	}
+	copy(res.Tokens, output.Tokens)
+	return res
+}
+
+// Has return DocId exists
+func (d *Dao) Has(id uint64) bool {
+	return d.searcher.HasDoc(id)
+}

app/service/main/riot-search/dao/searcher_test.go

@@ -0,0 +1,57 @@
+package dao
+
+import (
+	"go-common/app/service/main/riot-search/model"
+	"testing"
+
+	"github.com/smartystreets/goconvey/convey"
+)
+
+func TestDaoSearchIDOnly(t *testing.T) {
+	convey.Convey("SearchIDOnly", t, func(ctx convey.C) {
+		var (
+			arg1 = &model.RiotSearchReq{}
+			arg2 = &model.RiotSearchReq{Keyword: "test", IDs: []uint64{1}}
+		)
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			p1 := d.SearchIDOnly(arg1)
+			ctx.Convey("Then p1 should be nil.", func(ctx convey.C) {
+				ctx.So(p1, convey.ShouldBeNil)
+			})
+			p2 := d.SearchIDOnly(arg2)
+			ctx.Convey("Then p2 should not be nil.", func(ctx convey.C) {
+				ctx.So(p2, convey.ShouldNotBeNil)
+			})
+		})
+	})
+}
+
+func TestDaoSearch(t *testing.T) {
+	convey.Convey("Search", t, func(ctx convey.C) {
+		var (
+			arg1 = &model.RiotSearchReq{}
+			arg2 = &model.RiotSearchReq{Keyword: "test", IDs: []uint64{1}}
+		)
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			p1 := d.Search(arg1)
+			ctx.Convey("Then p1 should be nil.", func(ctx convey.C) {
+				ctx.So(p1, convey.ShouldBeNil)
+			})
+			p2 := d.Search(arg2)
+			ctx.Convey("Then p2 should not be nil.", func(ctx convey.C) {
+				ctx.So(p2, convey.ShouldNotBeNil)
+			})
+		})
+	})
+}
+
+func TestDaoHas(t *testing.T) {
+	convey.Convey("Search", t, func(ctx convey.C) {
+		ctx.Convey("When everything goes positive", func(ctx convey.C) {
+			p1 := d.Has(1)
+			ctx.Convey("Then p1 should be false.", func(ctx convey.C) {
+				ctx.So(p1, convey.ShouldBeFalse)
+			})
+		})
+	})
+}