From 5f534eeba17987544919f9b105818c0495a66a35 Mon Sep 17 00:00:00 2001
From: zwwhdls <zww@hdls.me>
Date: Sun, 22 Dec 2024 22:53:38 +0800
Subject: [PATCH 1/3] feat: search with token

Signed-off-by: zwwhdls <zww@hdls.me>
---
 go.mod                             | 11 ++--
 go.sum                             | 22 ++++---
 pkg/dispatch/plugin/doc_process.go | 95 ++++++++++++++++++++++++++++++
 pkg/dispatch/plugin/header.go      |  2 +-
 pkg/dispatch/plugin/subcontent.go  | 17 ------
 pkg/dispatch/plugin/util.go        | 38 ++++++++++++
 pkg/models/doc/document.go         |  3 +
 pkg/store/postgres/document.go     | 12 +++-
 pkg/store/postgres/migrate.go      | 18 +++++-
 pkg/store/postgres/model.go        | 38 ++++++++++++
 10 files changed, 223 insertions(+), 33 deletions(-)
 create mode 100644 pkg/dispatch/plugin/doc_process.go
 create mode 100644 pkg/dispatch/plugin/util.go

diff --git a/go.mod b/go.mod
index 998fe6d..8165d8c 100644
--- a/go.mod
+++ b/go.mod
@@ -7,12 +7,13 @@ toolchain go1.23.3
 require (
 	github.com/PuerkitoBio/goquery v1.10.0
 	github.com/basenana/go-flow v0.0.0-20230801131009-d05f1f41b706
-	github.com/blevesearch/bleve/v2 v2.3.10
+	github.com/blevesearch/bleve/v2 v2.4.0
 	github.com/blevesearch/upsidedown_store_api v1.0.2
 	github.com/cdipaolo/goml v0.0.0-20220715001353-00e0c845ae1c
 	github.com/gin-gonic/gin v1.10.0
 	github.com/go-gormigrate/gormigrate/v2 v2.1.1
 	github.com/google/uuid v1.3.0
+	github.com/hyponet/jiebago v0.0.0-20240525141904-e34990856482
 	github.com/meilisearch/meilisearch-go v0.29.0
 	github.com/onsi/ginkgo v1.16.5
 	github.com/onsi/ginkgo/v2 v2.13.0
@@ -30,12 +31,13 @@ require (
 	github.com/andybalholm/brotli v1.1.0 // indirect
 	github.com/andybalholm/cascadia v1.3.2 // indirect
 	github.com/bits-and-blooms/bitset v1.2.0 // indirect
-	github.com/blevesearch/bleve_index_api v1.0.6 // indirect
-	github.com/blevesearch/geo v0.1.18 // indirect
+	github.com/blevesearch/bleve_index_api v1.1.6 // indirect
+	github.com/blevesearch/geo v0.1.20 // indirect
+	github.com/blevesearch/go-faiss v1.0.13 // indirect
 	github.com/blevesearch/go-porterstemmer v1.0.3 // indirect
 	github.com/blevesearch/gtreap v0.1.1 // indirect
 	github.com/blevesearch/mmap-go v1.0.4 // indirect
-	github.com/blevesearch/scorch_segment_api/v2 v2.1.6 // indirect
+	github.com/blevesearch/scorch_segment_api/v2 v2.2.9 // indirect
 	github.com/blevesearch/segment v0.9.1 // indirect
 	github.com/blevesearch/snowballstem v0.9.0 // indirect
 	github.com/blevesearch/vellum v1.0.10 // indirect
@@ -44,6 +46,7 @@ require (
 	github.com/blevesearch/zapx/v13 v13.3.10 // indirect
 	github.com/blevesearch/zapx/v14 v14.3.10 // indirect
 	github.com/blevesearch/zapx/v15 v15.3.13 // indirect
+	github.com/blevesearch/zapx/v16 v16.0.12 // indirect
 	github.com/bytedance/sonic v1.12.5 // indirect
 	github.com/bytedance/sonic/loader v0.2.1 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
diff --git a/go.sum b/go.sum
index 8fd3ad8..ab550f7 100644
--- a/go.sum
+++ b/go.sum
@@ -10,20 +10,22 @@ github.com/basenana/go-flow v0.0.0-20230801131009-d05f1f41b706 h1:FxXoMwMZsufBjS
 github.com/basenana/go-flow v0.0.0-20230801131009-d05f1f41b706/go.mod h1:Rs13PWsg/ITdXRiVJcI+yS0iqCfNHxCbIFEt5DCt/RQ=
 github.com/bits-and-blooms/bitset v1.2.0 h1:Kn4yilvwNtMACtf1eYDlG8H77R07mZSPbMjLyS07ChA=
 github.com/bits-and-blooms/bitset v1.2.0/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA=
-github.com/blevesearch/bleve/v2 v2.3.10 h1:z8V0wwGoL4rp7nG/O3qVVLYxUqCbEwskMt4iRJsPLgg=
-github.com/blevesearch/bleve/v2 v2.3.10/go.mod h1:RJzeoeHC+vNHsoLR54+crS1HmOWpnH87fL70HAUCzIA=
-github.com/blevesearch/bleve_index_api v1.0.6 h1:gyUUxdsrvmW3jVhhYdCVL6h9dCjNT/geNU7PxGn37p8=
-github.com/blevesearch/bleve_index_api v1.0.6/go.mod h1:YXMDwaXFFXwncRS8UobWs7nvo0DmusriM1nztTlj1ms=
-github.com/blevesearch/geo v0.1.18 h1:Np8jycHTZ5scFe7VEPLrDoHnnb9C4j636ue/CGrhtDw=
-github.com/blevesearch/geo v0.1.18/go.mod h1:uRMGWG0HJYfWfFJpK3zTdnnr1K+ksZTuWKhXeSokfnM=
+github.com/blevesearch/bleve/v2 v2.4.0 h1:2xyg+Wv60CFHYccXc+moGxbL+8QKT/dZK09AewHgKsg=
+github.com/blevesearch/bleve/v2 v2.4.0/go.mod h1:IhQHoFAbHgWKYavb9rQgQEJJVMuY99cKdQ0wPpst2aY=
+github.com/blevesearch/bleve_index_api v1.1.6 h1:orkqDFCBuNU2oHW9hN2YEJmet+TE9orml3FCGbl1cKk=
+github.com/blevesearch/bleve_index_api v1.1.6/go.mod h1:PbcwjIcRmjhGbkS/lJCpfgVSMROV6TRubGGAODaK1W8=
+github.com/blevesearch/geo v0.1.20 h1:paaSpu2Ewh/tn5DKn/FB5SzvH0EWupxHEIwbCk/QPqM=
+github.com/blevesearch/geo v0.1.20/go.mod h1:DVG2QjwHNMFmjo+ZgzrIq2sfCh6rIHzy9d9d0B59I6w=
+github.com/blevesearch/go-faiss v1.0.13 h1:zfFs7ZYD0NqXVSY37j0JZjZT1BhE9AE4peJfcx/NB4A=
+github.com/blevesearch/go-faiss v1.0.13/go.mod h1:jrxHrbl42X/RnDPI+wBoZU8joxxuRwedrxqswQ3xfU8=
 github.com/blevesearch/go-porterstemmer v1.0.3 h1:GtmsqID0aZdCSNiY8SkuPJ12pD4jI+DdXTAn4YRcHCo=
 github.com/blevesearch/go-porterstemmer v1.0.3/go.mod h1:angGc5Ht+k2xhJdZi511LtmxuEf0OVpvUUNrwmM1P7M=
 github.com/blevesearch/gtreap v0.1.1 h1:2JWigFrzDMR+42WGIN/V2p0cUvn4UP3C4Q5nmaZGW8Y=
 github.com/blevesearch/gtreap v0.1.1/go.mod h1:QaQyDRAT51sotthUWAH4Sj08awFSSWzgYICSZ3w0tYk=
 github.com/blevesearch/mmap-go v1.0.4 h1:OVhDhT5B/M1HNPpYPBKIEJaD0F3Si+CrEKULGCDPWmc=
 github.com/blevesearch/mmap-go v1.0.4/go.mod h1:EWmEAOmdAS9z/pi/+Toxu99DnsbhG1TIxUoRmJw/pSs=
-github.com/blevesearch/scorch_segment_api/v2 v2.1.6 h1:CdekX/Ob6YCYmeHzD72cKpwzBjvkOGegHOqhAkXp6yA=
-github.com/blevesearch/scorch_segment_api/v2 v2.1.6/go.mod h1:nQQYlp51XvoSVxcciBjtvuHPIVjlWrN1hX4qwK2cqdc=
+github.com/blevesearch/scorch_segment_api/v2 v2.2.9 h1:3nBaSBRFokjE4FtPW3eUDgcAu3KphBg1GP07zy/6Uyk=
+github.com/blevesearch/scorch_segment_api/v2 v2.2.9/go.mod h1:ckbeb7knyOOvAdZinn/ASbB7EA3HoagnJkmEV3J7+sg=
 github.com/blevesearch/segment v0.9.1 h1:+dThDy+Lvgj5JMxhmOVlgFfkUtZV2kw49xax4+jTfSU=
 github.com/blevesearch/segment v0.9.1/go.mod h1:zN21iLm7+GnBHWTao9I+Au/7MBiL8pPFtJBJTsk6kQw=
 github.com/blevesearch/snowballstem v0.9.0 h1:lMQ189YspGP6sXvZQ4WZ+MLawfV8wOmPoD/iWeNXm8s=
@@ -42,6 +44,8 @@ github.com/blevesearch/zapx/v14 v14.3.10 h1:SG6xlsL+W6YjhX5N3aEiL/2tcWh3DO75Bnz7
 github.com/blevesearch/zapx/v14 v14.3.10/go.mod h1:qqyuR0u230jN1yMmE4FIAuCxmahRQEOehF78m6oTgns=
 github.com/blevesearch/zapx/v15 v15.3.13 h1:6EkfaZiPlAxqXz0neniq35my6S48QI94W/wyhnpDHHQ=
 github.com/blevesearch/zapx/v15 v15.3.13/go.mod h1:Turk/TNRKj9es7ZpKK95PS7f6D44Y7fAFy8F4LXQtGg=
+github.com/blevesearch/zapx/v16 v16.0.12 h1:Uccxvjmn+hQ6ywQP+wIiTpdq9LnAviGoryJOmGwAo/I=
+github.com/blevesearch/zapx/v16 v16.0.12/go.mod h1:MYnOshRfSm4C4drxx1LGRI+MVFByykJ2anDY1fxdk9Q=
 github.com/bytedance/sonic v1.12.5 h1:hoZxY8uW+mT+OpkcUWw4k0fDINtOcVavEsGfzwzFU/w=
 github.com/bytedance/sonic v1.12.5/go.mod h1:B8Gt/XvtZ3Fqj+iSKMypzymZxw/FVwgIGKzMzT9r/rk=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
@@ -111,6 +115,8 @@ github.com/google/pprof v0.0.0-20230207041349-798e818bf904/go.mod h1:uglQLonpP8q
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
+github.com/hyponet/jiebago v0.0.0-20240525141904-e34990856482 h1:6phC4CbKBzBQ7PHsf3I21WLF9Fxy8feaLhXmCgh/lcs=
+github.com/hyponet/jiebago v0.0.0-20240525141904-e34990856482/go.mod h1:/QDKR3hYNk0E+DX4XWJHuBsiC1UTGdGynWjps4Kj+T0=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
diff --git a/pkg/dispatch/plugin/doc_process.go b/pkg/dispatch/plugin/doc_process.go
new file mode 100644
index 0000000..d8bcb42
--- /dev/null
+++ b/pkg/dispatch/plugin/doc_process.go
@@ -0,0 +1,95 @@
+/*
+ Copyright 2024 Friday Author.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+
+package plugin
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/PuerkitoBio/goquery"
+	"github.com/hyponet/jiebago"
+
+	"github.com/basenana/friday/pkg/models/doc"
+)
+
+type DocProcessPlugin struct {
+	seg jiebago.Segmenter
+}
+
+func NewDocProcessPlugin() *DocProcessPlugin {
+	seg := jiebago.Segmenter{}
+	seg.LoadDictionary("dict.txt")
+	return &DocProcessPlugin{
+		seg: seg,
+	}
+}
+
+var _ ChainPlugin = &DocProcessPlugin{}
+
+func (s *DocProcessPlugin) Name() string {
+	return "docProcess"
+}
+
+func (s *DocProcessPlugin) Run(ctx context.Context, doc *doc.Document) error {
+	var err error
+	// html analysis
+	doc.PureContent, err = trimContent(doc.Content)
+	if err != nil {
+		return fmt.Errorf("process doc with id %d error: %s", doc.EntryId, err)
+	}
+
+	// split title
+	doc.TitleTokens = s.splitTokens(doc.Name)
+
+	// split content
+	doc.ContentTokens = s.splitTokens(doc.PureContent)
+	return nil
+}
+
+func trimContent(content string) (string, error) {
+	query, err := goquery.NewDocumentFromReader(bytes.NewReader([]byte(content)))
+	if err != nil {
+		return "", err
+	}
+
+	query.Find("body").EachWithBreak(func(i int, selection *goquery.Selection) bool {
+		t := strings.TrimSpace(selection.Text())
+		if t != "" {
+			content = t
+		}
+		return true
+	})
+
+	content = ContentTrim("html", content)
+	content = strings.ReplaceAll(content, "'", "")
+	return content, nil
+}
+
+func (s *DocProcessPlugin) splitTokens(content string) []string {
+	contentCh := s.seg.CutForSearch(content, true)
+	tokens := make([]string, 0, len(contentCh))
+	for token := range contentCh {
+		tokens = append(tokens, token)
+	}
+	return tokens
+}
+
+func init() {
+	RegisterPlugin(NewDocProcessPlugin())
+}
diff --git a/pkg/dispatch/plugin/header.go b/pkg/dispatch/plugin/header.go
index 76bf408..a0ccd70 100644
--- a/pkg/dispatch/plugin/header.go
+++ b/pkg/dispatch/plugin/header.go
@@ -39,7 +39,7 @@ func (h *HeaderImgPlugin) Run(ctx context.Context, doc *doc.Document) error {
 	var headerImgUrl string
 	query, err := goquery.NewDocumentFromReader(bytes.NewReader([]byte(doc.Content)))
 	if err != nil {
-		return fmt.Errorf("build doc query with id %d error: %s", doc.EntryId, err)
+		return fmt.Errorf("get header image of doc with id %d error: %s", doc.EntryId, err)
 	}
 
 	query.Find("img").EachWithBreak(func(i int, selection *goquery.Selection) bool {
diff --git a/pkg/dispatch/plugin/subcontent.go b/pkg/dispatch/plugin/subcontent.go
index 8ae620b..9050173 100644
--- a/pkg/dispatch/plugin/subcontent.go
+++ b/pkg/dispatch/plugin/subcontent.go
@@ -19,7 +19,6 @@ package plugin
 import (
 	"bytes"
 	"context"
-	"regexp"
 	"strings"
 
 	"github.com/PuerkitoBio/goquery"
@@ -41,22 +40,6 @@ func (s *SubContentPlugin) Run(ctx context.Context, doc *doc.Document) error {
 
 var _ ChainPlugin = &SubContentPlugin{}
 
-var repeatSpace = regexp.MustCompile(`\s+`)
-var htmlCharFilterRegexp = regexp.MustCompile(`</?[!\w:]+((\s+[\w-]+(\s*=\s*(?:\\*".*?"|'.*?'|[^'">\s]+))?)+\s*|\s*)/?>`)
-
-func ContentTrim(contentType, content string) string {
-	switch contentType {
-	case "html", "htm", "webarchive", ".webarchive":
-		content = strings.ReplaceAll(content, "</p>", "</p>\n")
-		content = strings.ReplaceAll(content, "</P>", "</P>\n")
-		content = strings.ReplaceAll(content, "</div>", "</div>\n")
-		content = strings.ReplaceAll(content, "</DIV>", "</DIV>\n")
-		content = htmlCharFilterRegexp.ReplaceAllString(content, "")
-	}
-	content = repeatSpace.ReplaceAllString(content, " ")
-	return content
-}
-
 func GenerateContentSubContent(content string) string {
 	if subContent, err := slowPathContentSubContent([]byte(content)); err == nil {
 		return subContent
diff --git a/pkg/dispatch/plugin/util.go b/pkg/dispatch/plugin/util.go
new file mode 100644
index 0000000..20c4b1b
--- /dev/null
+++ b/pkg/dispatch/plugin/util.go
@@ -0,0 +1,38 @@
+/*
+ Copyright 2024 Friday Author.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+
+package plugin
+
+import (
+	"regexp"
+	"strings"
+)
+
+var repeatSpace = regexp.MustCompile(`\s+`)
+var htmlCharFilterRegexp = regexp.MustCompile(`</?[!\w:]+((\s+[\w-]+(\s*=\s*(?:\\*".*?"|'.*?'|[^'">\s]+))?)+\s*|\s*)/?>`)
+
+func ContentTrim(contentType, content string) string {
+	switch contentType {
+	case "html", "htm", "webarchive", ".webarchive":
+		content = strings.ReplaceAll(content, "</p>", "</p>\n")
+		content = strings.ReplaceAll(content, "</P>", "</P>\n")
+		content = strings.ReplaceAll(content, "</div>", "</div>\n")
+		content = strings.ReplaceAll(content, "</DIV>", "</DIV>\n")
+		content = htmlCharFilterRegexp.ReplaceAllString(content, "")
+	}
+	content = repeatSpace.ReplaceAllString(content, " ")
+	return content
+}
diff --git a/pkg/models/doc/document.go b/pkg/models/doc/document.go
index ccdb588..ab07a19 100644
--- a/pkg/models/doc/document.go
+++ b/pkg/models/doc/document.go
@@ -35,6 +35,9 @@ type Document struct {
 	WebUrl        string    `json:"web_url,omitempty"`
 	HeaderImage   string    `json:"header_image,omitempty"`
 	SubContent    string    `json:"sub_content,omitempty"`
+	PureContent   string    `json:"pure_content,omitempty"`
+	TitleTokens   []string  `json:"title_tokens,omitempty"`
+	ContentTokens []string  `json:"content_tokens,omitempty"`
 	Marked        *bool     `json:"marked,omitempty"`
 	Unread        *bool     `json:"unread,omitempty"`
 	CreatedAt     time.Time `json:"created_at"`
diff --git a/pkg/store/postgres/document.go b/pkg/store/postgres/document.go
index 92d1b76..387098e 100644
--- a/pkg/store/postgres/document.go
+++ b/pkg/store/postgres/document.go
@@ -40,6 +40,14 @@ func (p *PostgresClient) CreateDocument(ctx context.Context, doc *doc.Document)
 			if errors.Is(res.Error, gorm.ErrRecordNotFound) {
 				docMod = docMod.From(doc)
 				res = tx.Create(docMod)
+				if res.Error != nil {
+					return res.Error
+				}
+
+				// update tokens
+				docMod = docMod.Tokens(doc)
+				p.log.Info("update token", docMod.Token)
+				res = tx.Model(&Document{}).Where("id = ?", docMod.ID).Update("token", gorm.Expr(string(docMod.Token)))
 				return res.Error
 			}
 			return res.Error
@@ -144,7 +152,9 @@ func docQueryFilter(tx *gorm.DB, filter *doc.DocumentFilter) *gorm.DB {
 		tx = tx.Where("document.name LIKE ?", "%"+filter.FuzzyName+"%")
 	}
 	if filter.Search != "" {
-		tx = tx.Where("document.content LIKE ?", "%"+filter.Search+"%")
+		tx = tx.Where("document.token @@ to_tsquery('simple', ?)", filter.Search).
+			Select("*, ts_rank(document.token, to_tsquery('simple', ?)) as rank", filter.Search).
+			Order("rank DESC")
 	}
 	return tx
 }
diff --git a/pkg/store/postgres/migrate.go b/pkg/store/postgres/migrate.go
index 98bbe18..bbe48c9 100644
--- a/pkg/store/postgres/migrate.go
+++ b/pkg/store/postgres/migrate.go
@@ -59,8 +59,22 @@ func buildMigrations() []*gormigrate.Migration {
 				if err != nil {
 					return err
 				}
-				//_ = db.Exec(`CREATE INDEX name ON table USING gist(content);`)
-				//_ = db.Exec("CREATE INDEX idx_doc_content ON document USING GIN (content gin_trgm_ops);")
+				return nil
+			},
+			Rollback: func(db *gorm.DB) error {
+				return nil
+			},
+		},
+		{
+			ID: "2024122201",
+			Migrate: func(db *gorm.DB) error {
+				err := db.AutoMigrate(
+					&Document{},
+				)
+				if err != nil {
+					return err
+				}
+				_ = db.Exec("CREATE INDEX idx_document_token ON public.document USING gin (token);")
 				return nil
 			},
 			Rollback: func(db *gorm.DB) error {
diff --git a/pkg/store/postgres/model.go b/pkg/store/postgres/model.go
index fa48e00..5f401a3 100644
--- a/pkg/store/postgres/model.go
+++ b/pkg/store/postgres/model.go
@@ -17,7 +17,10 @@
 package postgres
 
 import (
+	"database/sql/driver"
 	"encoding/json"
+	"fmt"
+	"strings"
 	"time"
 
 	"github.com/basenana/friday/pkg/models/doc"
@@ -114,6 +117,28 @@ func (v *BleveKV) TableName() string {
 	return "friday_blevekv"
 }
 
+type TsVector string
+
+func (t *TsVector) Scan(value interface{}) error {
+	if value == nil {
+		*t = ""
+		return nil
+	}
+	switch v := value.(type) {
+	case []byte:
+		*t = TsVector(v)
+	case string:
+		*t = TsVector(v)
+	default:
+		return fmt.Errorf("cannot scan type %T into TsVector", value)
+	}
+	return nil
+}
+
+func (t *TsVector) Value() (driver.Value, error) {
+	return string(*t), nil
+}
+
 type Document struct {
 	ID            int64     `gorm:"column:id;primaryKey"`
 	OID           int64     `gorm:"column:oid;index:doc_oid"`
@@ -126,8 +151,10 @@ type Document struct {
 	Summary       string    `gorm:"column:summary"`
 	HeaderImage   string    `gorm:"column:header_image"`
 	SubContent    string    `gorm:"column:sub_content"`
+	PureContent   string    `gorm:"column:pure_content"`
 	Marked        bool      `gorm:"column:marked;index:doc_is_marked"`
 	Unread        bool      `gorm:"column:unread;index:doc_is_unread"`
+	Token         TsVector  `gorm:"column:token;type:tsvector"`
 	CreatedAt     time.Time `gorm:"column:created_at"`
 	ChangedAt     time.Time `gorm:"column:changed_at"`
 }
@@ -154,10 +181,21 @@ func (d *Document) From(document *doc.Document) *Document {
 		d.Unread = *document.Unread
 	}
 	d.HeaderImage = document.HeaderImage
+	d.PureContent = document.PureContent
 	d.SubContent = document.SubContent
 	return d
 }
 
+func (d *Document) Tokens(document *doc.Document) *Document {
+	d.Token = TsVector(
+		fmt.Sprintf("setweight(to_tsvector('simple', '%s'), 'A') || setweight(to_tsvector('simple', '%s'), 'B')",
+			strings.Join(document.TitleTokens, " "),
+			strings.Join(document.ContentTokens, " "),
+		),
+	)
+	return d
+}
+
 func (d *Document) UpdateFrom(document *doc.Document) *Document {
 	if document.Unread != nil {
 		d.Unread = *document.Unread

From 09756143ee34159f513ff47be5b1b4cdc0507ed4 Mon Sep 17 00:00:00 2001
From: zwwhdls <zww@hdls.me>
Date: Tue, 24 Dec 2024 21:07:50 +0800
Subject: [PATCH 2/3] add unittest

Signed-off-by: zwwhdls <zww@hdls.me>
---
 go.mod                                | 10 ++-
 go.sum                                | 19 ++++++
 pkg/service/chain.go                  |  9 ++-
 pkg/service/chain_test.go             |  8 ++-
 pkg/store/{postgres => db}/migrate.go |  4 +-
 pkg/store/{postgres => db}/model.go   |  4 +-
 pkg/store/interface.go                |  1 +
 pkg/store/meili/meili.go              |  4 ++
 pkg/store/meili/mock.go               |  4 ++
 pkg/store/memory/memory.go            | 98 +++++++++++++++++++++++++++
 pkg/store/postgres/bleve.go           | 26 +++----
 pkg/store/postgres/document.go        | 61 ++++++++++-------
 pkg/store/postgres/postgres.go        | 10 +--
 pkg/store/postgres/vector.go          | 19 +++---
 14 files changed, 219 insertions(+), 58 deletions(-)
 rename pkg/store/{postgres => db}/migrate.go (97%)
 rename pkg/store/{postgres => db}/model.go (99%)
 create mode 100644 pkg/store/memory/memory.go

diff --git a/go.mod b/go.mod
index 8165d8c..a85b3d6 100644
--- a/go.mod
+++ b/go.mod
@@ -23,7 +23,7 @@ require (
 	go.uber.org/zap v1.27.0
 	golang.org/x/time v0.4.0
 	gorm.io/driver/postgres v1.5.2
-	gorm.io/gorm v1.25.4
+	gorm.io/gorm v1.25.7
 )
 
 require (
@@ -51,9 +51,12 @@ require (
 	github.com/bytedance/sonic/loader v0.2.1 // indirect
 	github.com/cloudwego/base64x v0.1.4 // indirect
 	github.com/cloudwego/iasm v0.2.0 // indirect
+	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.7 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
+	github.com/glebarez/go-sqlite v1.21.2 // indirect
+	github.com/glebarez/sqlite v1.11.0 // indirect
 	github.com/go-logr/logr v1.2.4 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
@@ -84,6 +87,7 @@ require (
 	github.com/mschoch/smat v0.2.0 // indirect
 	github.com/nxadm/tail v1.4.8 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.3 // indirect
+	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
 	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
@@ -100,4 +104,8 @@ require (
 	google.golang.org/protobuf v1.35.2 // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	modernc.org/libc v1.22.5 // indirect
+	modernc.org/mathutil v1.5.0 // indirect
+	modernc.org/memory v1.5.0 // indirect
+	modernc.org/sqlite v1.23.1 // indirect
 )
diff --git a/go.sum b/go.sum
index ab550f7..5623244 100644
--- a/go.sum
+++ b/go.sum
@@ -61,6 +61,8 @@ github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
 github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
@@ -70,6 +72,10 @@ github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE
 github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI=
 github.com/gin-gonic/gin v1.10.0 h1:nTuyha1TYqgedzytsKYqna+DfLos46nTv2ygFy86HFU=
 github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y=
+github.com/glebarez/go-sqlite v1.21.2 h1:3a6LFC4sKahUunAmynQKLZceZCOzUthkRkEAl9gAXWo=
+github.com/glebarez/go-sqlite v1.21.2/go.mod h1:sfxdZyhQjTM2Wry3gVYWaW072Ri1WMdWJi0k6+3382k=
+github.com/glebarez/sqlite v1.11.0 h1:wSG0irqzP6VurnMEpFGer5Li19RpIRi2qvQz++w0GMw=
+github.com/glebarez/sqlite v1.11.0/go.mod h1:h8/o8j5wiAsqSPoWELDUdJXhjAhsVliSn7bWZjOhrgQ=
 github.com/go-gormigrate/gormigrate/v2 v2.1.1 h1:eGS0WTFRV30r103lU8JNXY27KbviRnqqIDobW3EV3iY=
 github.com/go-gormigrate/gormigrate/v2 v2.1.1/go.mod h1:L7nJ620PFDKei9QOhJzqA8kRCk+E3UbV2f5gv+1ndLc=
 github.com/go-logr/logr v1.2.4 h1:g01GSCwiDw2xSZfjJ2/T9M+S6pFdcNtFYsp+Y43HYDQ=
@@ -177,6 +183,9 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/redis/rueidis v1.0.0 h1:LrUhkD46Es7neMvpTgqyYGRpvlGG4F6dLIRq+nUw/ho=
 github.com/redis/rueidis v1.0.0/go.mod h1:yxbpgX+VYNxCvdE0KEQXDeUFcF2hB2Oz/TJiaqFxoEU=
+github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
+github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
 github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
@@ -308,4 +317,14 @@ gorm.io/driver/postgres v1.5.2 h1:ytTDxxEv+MplXOfFe3Lzm7SjG09fcdb3Z/c056DTBx0=
 gorm.io/driver/postgres v1.5.2/go.mod h1:fmpX0m2I1PKuR7mKZiEluwrP3hbs+ps7JIGMUBpCgl8=
 gorm.io/gorm v1.25.4 h1:iyNd8fNAe8W9dvtlgeRI5zSVZPsq3OpcTu37cYcpCmw=
 gorm.io/gorm v1.25.4/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k=
+gorm.io/gorm v1.25.7 h1:VsD6acwRjz2zFxGO50gPO6AkNs7KKnvfzUjHQhZDz/A=
+gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8=
+modernc.org/libc v1.22.5 h1:91BNch/e5B0uPbJFgqbxXuOnxBQjlS//icfQEGmvyjE=
+modernc.org/libc v1.22.5/go.mod h1:jj+Z7dTNX8fBScMVNRAYZ/jF91K8fdT2hYMThc3YjBY=
+modernc.org/mathutil v1.5.0 h1:rV0Ko/6SfM+8G+yKiyI830l3Wuz1zRutdslNoQ0kfiQ=
+modernc.org/mathutil v1.5.0/go.mod h1:mZW8CKdRPY1v87qxC/wUdX5O1qDzXMP5TH3wjfpga6E=
+modernc.org/memory v1.5.0 h1:N+/8c5rE6EqugZwHii4IFsaJ7MUhoWX07J5tC/iI5Ds=
+modernc.org/memory v1.5.0/go.mod h1:PkUhL0Mugw21sHPeskwZW4D6VscE/GQJOnIpCnW6pSU=
+modernc.org/sqlite v1.23.1 h1:nrSBg4aRQQwq59JpvGEQ15tNxoO5pX/kUjcRNwSAGQM=
+modernc.org/sqlite v1.23.1/go.mod h1:OrDj17Mggn6MhE+iPbBNf7RGKODDE9NFT0f3EwDzJqk=
 nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50=
diff --git a/pkg/service/chain.go b/pkg/service/chain.go
index 64d9d9d..217ba60 100644
--- a/pkg/service/chain.go
+++ b/pkg/service/chain.go
@@ -93,7 +93,11 @@ func (c *Chain) CreateDocument(ctx context.Context, document *doc.Document) erro
 			}
 		}
 		c.Log.Debugf("create document: %+v", document.Name)
-		return c.DocClient.CreateDocument(ctx, document)
+		if err := c.DocClient.CreateDocument(ctx, document); err != nil {
+			c.Log.Errorf("create document error: %s", err)
+			return err
+		}
+		return c.DocClient.UpdateTokens(ctx, document)
 	})
 }
 
@@ -110,6 +114,9 @@ func (c *Chain) GetDocument(ctx context.Context, namespace string, entryId int64
 	ctx = c.WithNamespace(ctx, namespace)
 	doc, err := c.DocClient.GetDocument(ctx, entryId)
 	if err != nil {
+		if err == models.ErrNotFound {
+			return nil, nil
+		}
 		c.Log.Errorf("get document error: %s", err)
 		return nil, err
 	}
diff --git a/pkg/service/chain_test.go b/pkg/service/chain_test.go
index cce87b5..677cdbf 100644
--- a/pkg/service/chain_test.go
+++ b/pkg/service/chain_test.go
@@ -28,7 +28,7 @@ import (
 	_ "github.com/basenana/friday/pkg/dispatch/plugin"
 	"github.com/basenana/friday/pkg/models/doc"
 	"github.com/basenana/friday/pkg/service"
-	"github.com/basenana/friday/pkg/store/meili"
+	"github.com/basenana/friday/pkg/store/memory"
 	"github.com/basenana/friday/pkg/utils/logger"
 )
 
@@ -49,8 +49,10 @@ var _ = Describe("Chain", func() {
 	BeforeEach(func() {
 		service.ChainPool = dispatch.NewPool(10)
 		logger.InitLog()
+		client, err := memory.NewMemoryMetaStore()
+		Expect(err).Should(BeNil())
 		Chain = &service.Chain{
-			DocClient: &meili.MockClient{},
+			DocClient: client,
 			Plugins:   []plugin.ChainPlugin{},
 			Log:       logger.NewLog("test"),
 		}
@@ -84,7 +86,7 @@ var _ = Describe("Chain", func() {
 			CreatedAt:     t,
 			ChangedAt:     t,
 		}
-		err := Chain.CreateDocument(context.TODO(), doc11)
+		err = Chain.CreateDocument(context.TODO(), doc11)
 		Expect(err).Should(BeNil())
 		err = Chain.CreateDocument(context.TODO(), doc12)
 		Expect(err).Should(BeNil())
diff --git a/pkg/store/postgres/migrate.go b/pkg/store/db/migrate.go
similarity index 97%
rename from pkg/store/postgres/migrate.go
rename to pkg/store/db/migrate.go
index bbe48c9..3bd1094 100644
--- a/pkg/store/postgres/migrate.go
+++ b/pkg/store/db/migrate.go
@@ -1,5 +1,5 @@
 /*
- Copyright 2023 Friday Author.
+ Copyright 2024 Friday Author.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  limitations under the License.
 */
 
-package postgres
+package db
 
 import (
 	"github.com/go-gormigrate/gormigrate/v2"
diff --git a/pkg/store/postgres/model.go b/pkg/store/db/model.go
similarity index 99%
rename from pkg/store/postgres/model.go
rename to pkg/store/db/model.go
index 5f401a3..ee97461 100644
--- a/pkg/store/postgres/model.go
+++ b/pkg/store/db/model.go
@@ -1,5 +1,5 @@
 /*
- Copyright 2023 Friday Author.
+ Copyright 2024 Friday Author.
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  limitations under the License.
 */
 
-package postgres
+package db
 
 import (
 	"database/sql/driver"
diff --git a/pkg/store/interface.go b/pkg/store/interface.go
index 138bf7b..c002a0a 100644
--- a/pkg/store/interface.go
+++ b/pkg/store/interface.go
@@ -31,6 +31,7 @@ type VectorStore interface {
 
 type DocStoreInterface interface {
 	CreateDocument(ctx context.Context, doc *doc.Document) error
+	UpdateTokens(ctx context.Context, doc *doc.Document) error
 	UpdateDocument(ctx context.Context, doc *doc.Document) error
 	GetDocument(ctx context.Context, entryId int64) (*doc.Document, error)
 	FilterDocuments(ctx context.Context, filter *doc.DocumentFilter) ([]*doc.Document, error)
diff --git a/pkg/store/meili/meili.go b/pkg/store/meili/meili.go
index 5f421e5..93320e4 100644
--- a/pkg/store/meili/meili.go
+++ b/pkg/store/meili/meili.go
@@ -155,6 +155,10 @@ func (c *Client) CreateDocument(ctx context.Context, doc *doc.Document) error {
 	return nil
 }
 
+func (c *Client) UpdateTokens(ctx context.Context, doc *doc.Document) error {
+	return nil
+}
+
 func (c *Client) UpdateDocument(ctx context.Context, doc *doc.Document) error {
 	// delete document attr
 	newAttrsQuery := (&DocumentAttrQuery{}).FromModel(doc)
diff --git a/pkg/store/meili/mock.go b/pkg/store/meili/mock.go
index b0f3b21..558574e 100644
--- a/pkg/store/meili/mock.go
+++ b/pkg/store/meili/mock.go
@@ -40,6 +40,10 @@ func (m *MockClient) CreateDocument(ctx context.Context, doc *doc.Document) erro
 	return nil
 }
 
+func (m *MockClient) UpdateTokens(ctx context.Context, doc *doc.Document) error {
+	return nil
+}
+
 func (m *MockClient) UpdateDocument(ctx context.Context, doc *doc.Document) error {
 	aq := (&DocumentAttrQuery{}).FromModel(doc)
 	result := []*DocumentAttr{}
diff --git a/pkg/store/memory/memory.go b/pkg/store/memory/memory.go
new file mode 100644
index 0000000..5cba58f
--- /dev/null
+++ b/pkg/store/memory/memory.go
@@ -0,0 +1,98 @@
+/*
+ Copyright 2024 Friday Author.
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+     http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+*/
+
+package memory
+
+import (
+	"context"
+	"time"
+
+	"github.com/glebarez/sqlite"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+
+	"github.com/basenana/friday/pkg/models/doc"
+	"github.com/basenana/friday/pkg/store"
+	"github.com/basenana/friday/pkg/store/db"
+	"github.com/basenana/friday/pkg/store/postgres"
+	"github.com/basenana/friday/pkg/store/utils"
+	"github.com/basenana/friday/pkg/utils/logger"
+)
+
+type MemoryClient struct {
+	log     *zap.SugaredLogger
+	dbStore postgres.PostgresClient
+}
+
+var _ store.DocStoreInterface = &MemoryClient{}
+
+func NewMemoryMetaStore() (*MemoryClient, error) {
+	log := logger.NewLog("memory")
+	dbEntity, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: utils.NewDbLogger()})
+	if err != nil {
+		return nil, err
+	}
+
+	dbConn, err := dbEntity.DB()
+	if err != nil {
+		return nil, err
+	}
+
+	if err = dbConn.Ping(); err != nil {
+		return nil, err
+	}
+
+	dbConn.SetMaxIdleConns(1)
+	dbConn.SetMaxOpenConns(1)
+	dbConn.SetConnMaxLifetime(time.Hour)
+
+	dbEnt, err := db.NewDbEntity(dbEntity, db.Migrate)
+	if err != nil {
+		return nil, err
+	}
+
+	return &MemoryClient{
+		dbStore: postgres.PostgresClient{
+			Logger:  log,
+			DEntity: dbEnt,
+		},
+	}, nil
+}
+
+func (m *MemoryClient) CreateDocument(ctx context.Context, doc *doc.Document) error {
+	return m.dbStore.CreateDocument(ctx, doc)
+}
+
+func (m *MemoryClient) UpdateTokens(ctx context.Context, doc *doc.Document) error {
+	return nil
+}
+
+func (m *MemoryClient) UpdateDocument(ctx context.Context, doc *doc.Document) error {
+	return m.dbStore.UpdateDocument(ctx, doc)
+}
+
+func (m *MemoryClient) GetDocument(ctx context.Context, entryId int64) (*doc.Document, error) {
+	return m.dbStore.GetDocument(ctx, entryId)
+}
+
+func (m *MemoryClient) FilterDocuments(ctx context.Context, filter *doc.DocumentFilter) ([]*doc.Document, error) {
+	filter.Search = ""
+	return m.dbStore.FilterDocuments(ctx, filter)
+}
+
+func (m *MemoryClient) DeleteDocument(ctx context.Context, docId int64) error {
+	return m.dbStore.DeleteDocument(ctx, docId)
+}
diff --git a/pkg/store/postgres/bleve.go b/pkg/store/postgres/bleve.go
index 2e6b576..d75ca09 100644
--- a/pkg/store/postgres/bleve.go
+++ b/pkg/store/postgres/bleve.go
@@ -9,6 +9,8 @@ import (
 	"github.com/blevesearch/bleve/v2/registry"
 	"github.com/blevesearch/upsidedown_store_api"
 	"gorm.io/gorm"
+
+	"github.com/basenana/friday/pkg/store/db"
 )
 
 const (
@@ -75,9 +77,9 @@ func (p *PgKVWriter) ExecuteBatch(batch store.KVBatch) error {
 		return fmt.Errorf("wrong type of batch")
 	}
 
-	err := p.dEntity.DB.WithContext(context.Background()).Transaction(func(tx *gorm.DB) error {
+	err := p.DEntity.DB.WithContext(context.Background()).Transaction(func(tx *gorm.DB) error {
 		for k, mergeOps := range emulatedBatch.Merger.Merges {
-			mod := &BleveKV{ID: bytes2Str([]byte(k)), Key: []byte(k)}
+			mod := &db.BleveKV{ID: bytes2Str([]byte(k)), Key: []byte(k)}
 
 			res := tx.First(mod)
 			if res.Error != nil && !errors.Is(res.Error, gorm.ErrRecordNotFound) {
@@ -99,7 +101,7 @@ func (p *PgKVWriter) ExecuteBatch(batch store.KVBatch) error {
 		}
 
 		for _, op := range emulatedBatch.Ops {
-			mod := &BleveKV{ID: bytes2Str(op.K), Key: op.K, Value: op.V}
+			mod := &db.BleveKV{ID: bytes2Str(op.K), Key: op.K, Value: op.V}
 			var res *gorm.DB
 			if op.V != nil {
 				res = tx.Save(mod)
@@ -127,8 +129,8 @@ type PgKVReader struct {
 var _ store.KVReader = &PgKVReader{}
 
 func (p *PgKVReader) Get(key []byte) ([]byte, error) {
-	mod := &BleveKV{ID: bytes2Str(key)}
-	res := p.dEntity.DB.WithContext(context.Background()).First(mod)
+	mod := &db.BleveKV{ID: bytes2Str(key)}
+	res := p.DEntity.DB.WithContext(context.Background()).First(mod)
 	if res.Error != nil {
 		if errors.Is(res.Error, gorm.ErrRecordNotFound) {
 			return nil, nil
@@ -144,8 +146,8 @@ func (p *PgKVReader) MultiGet(keys [][]byte) ([][]byte, error) {
 		idList[i] = bytes2Str(k)
 	}
 
-	var modList []BleveKV
-	res := p.dEntity.DB.WithContext(context.Background()).Where("id IN ?", idList).Find(modList)
+	var modList []db.BleveKV
+	res := p.DEntity.DB.WithContext(context.Background()).Where("id IN ?", idList).Find(modList)
 	if res.Error != nil {
 		return nil, res.Error
 	}
@@ -157,8 +159,8 @@ func (p *PgKVReader) MultiGet(keys [][]byte) ([][]byte, error) {
 }
 
 func (p *PgKVReader) PrefixIterator(prefix []byte) store.KVIterator {
-	var modList []BleveKV
-	res := p.PostgresClient.dEntity.WithContext(context.Background()).
+	var modList []db.BleveKV
+	res := p.PostgresClient.DEntity.WithContext(context.Background()).
 		Where("id LIKE ?", bytes2Str(prefix)+"%").
 		Order("key DESC").Find(&modList)
 	if res.Error != nil {
@@ -175,8 +177,8 @@ func (p *PgKVReader) PrefixIterator(prefix []byte) store.KVIterator {
 }
 
 func (p *PgKVReader) RangeIterator(start, end []byte) store.KVIterator {
-	var modList []BleveKV
-	res := p.PostgresClient.dEntity.WithContext(context.Background()).
+	var modList []db.BleveKV
+	res := p.PostgresClient.DEntity.WithContext(context.Background()).
 		Where("id >= ? AND id < ?", bytes2Str(start), bytes2Str(end)).
 		Order("key DESC").Find(&modList)
 	if res.Error != nil {
@@ -201,7 +203,7 @@ func NewPgKVReader(cli *PostgresClient) *PgKVReader {
 }
 
 type KVIterator struct {
-	listKV     []BleveKV
+	listKV     []db.BleveKV
 	start, end []byte
 	prefix     []byte
 	crt        int
diff --git a/pkg/store/postgres/document.go b/pkg/store/postgres/document.go
index 387098e..187d086 100644
--- a/pkg/store/postgres/document.go
+++ b/pkg/store/postgres/document.go
@@ -26,15 +26,16 @@ import (
 	"github.com/basenana/friday/pkg/models"
 	"github.com/basenana/friday/pkg/models/doc"
 	"github.com/basenana/friday/pkg/store"
+	"github.com/basenana/friday/pkg/store/db"
 	"github.com/basenana/friday/pkg/store/utils"
 )
 
 var _ store.DocStoreInterface = &PostgresClient{}
 
 func (p *PostgresClient) CreateDocument(ctx context.Context, doc *doc.Document) error {
-	defer trace.StartRegion(ctx, "metastore.sql.SaveDocument").End()
-	err := p.dEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		docMod := &Document{}
+	defer trace.StartRegion(ctx, "store.doc.CreateDocument").End()
+	err := p.DEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		docMod := &db.Document{}
 		res := tx.Where("oid = ?", doc.EntryId).First(docMod)
 		if res.Error != nil {
 			if errors.Is(res.Error, gorm.ErrRecordNotFound) {
@@ -44,11 +45,6 @@ func (p *PostgresClient) CreateDocument(ctx context.Context, doc *doc.Document)
 					return res.Error
 				}
 
-				// update tokens
-				docMod = docMod.Tokens(doc)
-				p.log.Info("update token", docMod.Token)
-				res = tx.Model(&Document{}).Where("id = ?", docMod.ID).Update("token", gorm.Expr(string(docMod.Token)))
-				return res.Error
 			}
 			return res.Error
 		}
@@ -62,10 +58,29 @@ func (p *PostgresClient) CreateDocument(ctx context.Context, doc *doc.Document)
 	return nil
 }
 
+func (p *PostgresClient) UpdateTokens(ctx context.Context, doc *doc.Document) error {
+	defer trace.StartRegion(ctx, "store.doc.UpdateTokens").End()
+	err := p.DEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		docMod := &db.Document{}
+		res := tx.Where("oid = ?", doc.EntryId).First(docMod)
+		if res.Error != nil {
+			return res.Error
+		}
+		docMod = docMod.Tokens(doc)
+		p.Logger.Info("update token", docMod.Token)
+		res = tx.Model(&db.Document{}).Where("id = ?", docMod.ID).Update("token", gorm.Expr(string(docMod.Token)))
+		return res.Error
+	})
+	if err != nil {
+		return utils.SqlError2Error(err)
+	}
+	return nil
+}
+
 func (p *PostgresClient) UpdateDocument(ctx context.Context, doc *doc.Document) error {
-	defer trace.StartRegion(ctx, "metastore.sql.SaveDocument").End()
-	err := p.dEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		docMod := &Document{}
+	defer trace.StartRegion(ctx, "store.doc.UpdateDocument").End()
+	err := p.DEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		docMod := &db.Document{}
 		res := tx.Where("oid = ?", doc.EntryId).First(docMod)
 		if res.Error != nil {
 			if errors.Is(res.Error, gorm.ErrRecordNotFound) {
@@ -86,9 +101,9 @@ func (p *PostgresClient) UpdateDocument(ctx context.Context, doc *doc.Document)
 }
 
 func (p *PostgresClient) GetDocument(ctx context.Context, entryId int64) (*doc.Document, error) {
-	defer trace.StartRegion(ctx, "metastore.sql.GetDocument").End()
-	doc := &Document{}
-	res := p.dEntity.WithNamespace(ctx).Where("oid = ?", entryId).First(doc)
+	defer trace.StartRegion(ctx, "store.doc.GetDocument").End()
+	doc := &db.Document{}
+	res := p.DEntity.WithNamespace(ctx).Where("oid = ?", entryId).First(doc)
 	if res.Error != nil {
 		return nil, utils.SqlError2Error(res.Error)
 	}
@@ -96,13 +111,13 @@ func (p *PostgresClient) GetDocument(ctx context.Context, entryId int64) (*doc.D
 }
 
 func (p *PostgresClient) FilterDocuments(ctx context.Context, filter *doc.DocumentFilter) ([]*doc.Document, error) {
-	defer trace.StartRegion(ctx, "metastore.sql.ListDocument").End()
-	docList := make([]Document, 0)
+	defer trace.StartRegion(ctx, "store.doc.FilterDocuments").End()
+	docList := make([]db.Document, 0)
 	q := p.WithNamespace(ctx)
 	if page := models.GetPagination(ctx); page != nil {
 		q = q.Offset(page.Offset()).Limit(page.Limit())
 	}
-	res := docOrder(docQueryFilter(q, filter), &filter.Order).Find(&docList)
+	res := docOrder(p.docQueryFilter(q, filter), &filter.Order).Find(&docList)
 	if res.Error != nil {
 		return nil, utils.SqlError2Error(res.Error)
 	}
@@ -115,15 +130,15 @@ func (p *PostgresClient) FilterDocuments(ctx context.Context, filter *doc.Docume
 }
 
 func (p *PostgresClient) DeleteDocument(ctx context.Context, entryId int64) error {
-	defer trace.StartRegion(ctx, "metastore.sql.DeleteDocument").End()
-	err := p.dEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		res := namespaceQuery(ctx, tx).Where("oid = ?", entryId).Delete(&Document{})
+	defer trace.StartRegion(ctx, "store.doc.DeleteDocument").End()
+	err := p.DEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		res := namespaceQuery(ctx, tx).Where("oid = ?", entryId).Delete(&db.Document{})
 		return res.Error
 	})
 	return utils.SqlError2Error(err)
 }
 
-func docQueryFilter(tx *gorm.DB, filter *doc.DocumentFilter) *gorm.DB {
+func (p *PostgresClient) docQueryFilter(tx *gorm.DB, filter *doc.DocumentFilter) *gorm.DB {
 	if filter.ParentID != nil {
 		tx = tx.Where("document.parent_entry_id = ?", filter.ParentID)
 	}
@@ -183,7 +198,7 @@ func namespaceQuery(ctx context.Context, tx *gorm.DB) *gorm.DB {
 func (p *PostgresClient) WithNamespace(ctx context.Context) *gorm.DB {
 	ns := models.GetNamespace(ctx)
 	if ns.String() == models.DefaultNamespaceValue {
-		return p.dEntity.WithContext(ctx)
+		return p.DEntity.WithContext(ctx)
 	}
-	return p.dEntity.WithContext(ctx).Where("namespace = ?", ns.String())
+	return p.DEntity.WithContext(ctx).Where("namespace = ?", ns.String())
 }
diff --git a/pkg/store/postgres/postgres.go b/pkg/store/postgres/postgres.go
index 661ca0c..ae5ef08 100644
--- a/pkg/store/postgres/postgres.go
+++ b/pkg/store/postgres/postgres.go
@@ -31,8 +31,8 @@ import (
 const defaultNamespace = "global"
 
 type PostgresClient struct {
-	log     *zap.SugaredLogger
-	dEntity *db.Entity
+	Logger  *zap.SugaredLogger
+	DEntity *db.Entity
 }
 
 func NewPostgresClient(postgresUrl string) (*PostgresClient, error) {
@@ -55,13 +55,13 @@ func NewPostgresClient(postgresUrl string) (*PostgresClient, error) {
 		return nil, err
 	}
 
-	dbEnt, err := db.NewDbEntity(dbObj, Migrate)
+	dbEnt, err := db.NewDbEntity(dbObj, db.Migrate)
 	if err != nil {
 		return nil, err
 	}
 
 	return &PostgresClient{
-		log:     log,
-		dEntity: dbEnt,
+		Logger:  log,
+		DEntity: dbEnt,
 	}, nil
 }
diff --git a/pkg/store/postgres/vector.go b/pkg/store/postgres/vector.go
index 9d5f367..1aa607c 100644
--- a/pkg/store/postgres/vector.go
+++ b/pkg/store/postgres/vector.go
@@ -28,6 +28,7 @@ import (
 
 	"github.com/basenana/friday/pkg/models/vector"
 	"github.com/basenana/friday/pkg/store"
+	"github.com/basenana/friday/pkg/store/db"
 	"github.com/basenana/friday/pkg/utils"
 )
 
@@ -36,7 +37,7 @@ func (p *PostgresClient) Store(ctx context.Context, element *vector.Element, ext
 	if namespace == nil {
 		namespace = defaultNamespace
 	}
-	return p.dEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+	return p.DEntity.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
 		if extra == nil {
 			extra = make(map[string]interface{})
 		}
@@ -48,7 +49,7 @@ func (p *PostgresClient) Store(ctx context.Context, element *vector.Element, ext
 			return err
 		}
 
-		var v *Index
+		var v *db.Index
 		v, err = v.From(element)
 		if err != nil {
 			return err
@@ -57,7 +58,7 @@ func (p *PostgresClient) Store(ctx context.Context, element *vector.Element, ext
 		v.Extra = string(b)
 		v.Namespace = namespace.(string)
 
-		vModel := &Index{}
+		vModel := &db.Index{}
 		res := tx.Where("namespace = ? AND name = ? AND idx_group = ?", namespace, element.Name, element.Group).First(vModel)
 		if res.Error != nil && res.Error != gorm.ErrRecordNotFound {
 			return res.Error
@@ -91,9 +92,9 @@ func (p *PostgresClient) Search(ctx context.Context, query vector.VectorDocQuery
 		vectors64 = append(vectors64, float64(v))
 	}
 	// query from db
-	existIndexes := make([]Index, 0)
+	existIndexes := make([]db.Index, 0)
 
-	res := p.dEntity.WithNamespace(ctx)
+	res := p.DEntity.WithNamespace(ctx)
 	if query.ParentId != 0 {
 		res = res.Where("parent_entry_id = ?", query.ParentId)
 	}
@@ -128,7 +129,7 @@ func (p *PostgresClient) Search(ctx context.Context, query vector.VectorDocQuery
 	}
 	results := make([]*vector.Doc, 0)
 	for _, index := range minKIndexes {
-		i := index.Object.(Index)
+		i := index.Object.(db.Index)
 		results = append(results, i.ToDoc())
 	}
 
@@ -136,8 +137,8 @@ func (p *PostgresClient) Search(ctx context.Context, query vector.VectorDocQuery
 }
 
 func (p *PostgresClient) Get(ctx context.Context, oid int64, name string, group int) (*vector.Element, error) {
-	vModel := &Index{}
-	tx := p.dEntity.WithNamespace(ctx).Where("name = ? AND idx_group = ?", name, group)
+	vModel := &db.Index{}
+	tx := p.DEntity.WithNamespace(ctx).Where("name = ? AND idx_group = ?", name, group)
 	if oid != 0 {
 		tx = tx.Where("oid = ?", oid)
 	}
@@ -155,7 +156,7 @@ var _ store.VectorStore = &PostgresClient{}
 
 func (p *PostgresClient) Inited(ctx context.Context) (bool, error) {
 	var count int64
-	res := p.dEntity.WithContext(ctx).Model(&BleveKV{}).Count(&count)
+	res := p.DEntity.WithContext(ctx).Model(&db.BleveKV{}).Count(&count)
 	if res.Error != nil {
 		return false, res.Error
 	}

From d0281370d4c3808c2c11cd7ef1bc1635179c79a4 Mon Sep 17 00:00:00 2001
From: zwwhdls <zww@hdls.me>
Date: Tue, 24 Dec 2024 21:12:12 +0800
Subject: [PATCH 3/3] update migrate

Signed-off-by: zwwhdls <zww@hdls.me>
---
 pkg/store/db/migrate.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pkg/store/db/migrate.go b/pkg/store/db/migrate.go
index 3bd1094..956f530 100644
--- a/pkg/store/db/migrate.go
+++ b/pkg/store/db/migrate.go
@@ -74,6 +74,7 @@ func buildMigrations() []*gormigrate.Migration {
 				if err != nil {
 					return err
 				}
+				_ = db.Exec("CREATE EXTENSION pg_trgm;")
 				_ = db.Exec("CREATE INDEX idx_document_token ON public.document USING gin (token);")
 				return nil
 			},