diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..485dee6 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.idea diff --git a/go.mod b/go.mod index 70f6c3c..08697a0 100644 --- a/go.mod +++ b/go.mod @@ -6,4 +6,7 @@ toolchain go1.22.1 require gopkg.in/telebot.v3 v3.3.8 -require github.com/stretchr/testify v1.10.0 // indirect +require ( + github.com/lib/pq v1.10.9 // indirect + github.com/stretchr/testify v1.10.0 // indirect +) diff --git a/go.sum b/go.sum index 59e21c5..f34785d 100644 --- a/go.sum +++ b/go.sum @@ -266,6 +266,8 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII= +github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= +github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= github.com/magiconair/properties v1.8.6/go.mod h1:y3VJvCyxH9uVvJTWEGAELF3aiYNyPKd5NZ3oSwXrF60= github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE= diff --git a/handlers.go b/handlers.go index e267ea4..6b21a85 100644 --- a/handlers.go +++ b/handlers.go @@ -1,6 +1,7 @@ package main import ( + "database/sql" "encoding/json" "fmt" "github.com/purofle/remake_bot/quotely" @@ -8,6 +9,7 @@ import ( "math/rand" crand "math/rand" "os" + "strconv" "strings" "sync" "time" @@ -177,7 +179,6 @@ func CommandEat(c tele.Context) error { } func CommandOnText(c tele.Context) error { - if c.Chat().ID != -1001965344356 { return nil } @@ -190,3 +191,81 @@ func CommandOnText(c tele.Context) error { } return nil } + +func getQuote(text string) (error, []string, []string) { + var rows *sql.Rows + var err error + if text == "" { + query := "select text, \"from\" from result_new where from_id not like 'channel%' order by random() limit 50" + rows, err = database.Query(query) + } else { + query := "select text, \"from\" from result_new where from_id not like 'channel%' AND text like '%' || $1 || '%' order by random() limit 50" + rows, err = database.Query(query, text) + } + if err != nil { + return err, nil, nil + } + + defer func(rows *sql.Rows) { + err := rows.Close() + if err != nil { + return + } + }(rows) + + var resultText []string + var from []string + for rows.Next() { + var t, f string + if err := rows.Scan(&t, &f); err != nil { + return err, nil, nil + } + resultText = append(resultText, t) + from = append(from, f) + } + + return nil, resultText, from +} + +func InlineQuery(c tele.Context) error { + member, err := c.Bot().ChatMemberOf( + &tele.Chat{ID: -1001965344356}, + c.Sender(), + ) + if err != nil { + fmt.Println(fmt.Sprintf("sender: %s 不在群", c.Sender().FirstName)) + return err + } + if member.Role != "creator" { + fmt.Println(fmt.Sprintf("sender: %s 不在群", c.Sender().FirstName)) + return nil + } + + var resultText []string + var from []string + + if c.Query().Text == "" { + err, resultText, from = getQuote("") + } else { + err, resultText, from = getQuote(c.Query().Text) + } + results := make(tele.Results, len(resultText)) + + if err != nil { + return err + } + + for i, text := range resultText { + results[i] = &tele.ArticleResult{ + Title: text, + Text: fmt.Sprintf("%s: %s", from[i], text), + Description: fmt.Sprintf("来自 %s", from[i]), + } + results[i].SetResultID(strconv.Itoa(i)) + } + + return c.Answer(&tele.QueryResponse{ + Results: results, + CacheTime: 0, + }) +} diff --git a/main.go b/main.go index 50e18a3..0c4ca27 100644 --- a/main.go +++ b/main.go @@ -1,11 +1,14 @@ package main import ( + "database/sql" "log" "os" "time" tele "gopkg.in/telebot.v3" + + _ "github.com/lib/pq" ) type RemakeData struct { @@ -15,6 +18,7 @@ type RemakeData struct { } var remakeCount map[int64]*RemakeData +var database *sql.DB func main() { pref := tele.Settings{ @@ -22,7 +26,14 @@ func main() { Poller: &tele.LongPoller{Timeout: 10 * time.Second}, } - err := initList() + connStr := "postgresql://postgres:114514@localhost:5432/postgres?sslmode=disable" + db, err := sql.Open("postgres", connStr) + if err != nil { + log.Fatal(err) + } + database = db + + err = initList() if err != nil { log.Fatal(err) } @@ -35,6 +46,7 @@ func main() { return } + b.Handle(tele.OnQuery, InlineQuery) b.Handle("/remake", CommandRemake) b.Handle("/remake_data", CommandRemakeData) b.Handle("/eat", CommandEat)