From 32b53c0b18eabb25304a7af16b34c72ec4ad7c5d Mon Sep 17 00:00:00 2001
From: Rogee <rogeecn@qq.com>
Date: Mon, 2 Sep 2024 18:47:12 +0800
Subject: [PATCH] store messages

---
 config/config.go       |  1 +
 internal/db_channel.go | 31 ++++++++++++++++++++-----------
 main.go                |  5 +++++
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/config/config.go b/config/config.go
index 75c9f31..664366b 100644
--- a/config/config.go
+++ b/config/config.go
@@ -14,6 +14,7 @@ type Config struct {
 	BotToken    string `mapstructure:"bot_token"`
 	SessionFile string `mapstructure:"session_file"`
 	LogFile     string `mapstructure:"log_file"`
+	DSN         string `mapstructure:"dsn"`
 }
 
 func Load(path string) error {
diff --git a/internal/db_channel.go b/internal/db_channel.go
index 495d462..8397385 100644
--- a/internal/db_channel.go
+++ b/internal/db_channel.go
@@ -2,7 +2,6 @@ package internal
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -12,6 +11,8 @@ import (
 	"exporter/database/telegram_resource/public/table"
 
 	. "github.com/go-jet/jet/v2/postgres"
+	"github.com/go-jet/jet/v2/qrm"
+	"github.com/lib/pq"
 	"github.com/pkg/errors"
 	"github.com/samber/lo"
 )
@@ -58,13 +59,13 @@ func (c *DBChannel) Asset(assetID int64, ext string) string {
 func (c *DBChannel) Get(ctx context.Context) error {
 	tbl := table.Channels
 
-	var m *model.Channels
+	var m model.Channels
 
-	err := tbl.SELECT(tbl.AllColumns).QueryContext(ctx, db, &m)
+	err := tbl.SELECT(tbl.AllColumns).WHERE(tbl.UUID.EQ(Int64(c.UUID))).QueryContext(ctx, db, &m)
 	if err != nil {
-		if errors.Is(err, sql.ErrNoRows) {
+		if errors.Is(err, qrm.ErrNoRows) {
 			// create new channel with default value
-			m = &model.Channels{
+			m = model.Channels{
 				UUID:      c.UUID,
 				Username:  c.Username,
 				Title:     c.Title,
@@ -74,7 +75,7 @@ func (c *DBChannel) Get(ctx context.Context) error {
 				UpdatedAt: lo.ToPtr(time.Now()),
 			}
 
-			if _, err := tbl.INSERT(tbl.AllColumns).MODEL(m).ExecContext(ctx, db); err != nil {
+			if _, err := tbl.INSERT(tbl.AllColumns.Except(tbl.ID)).MODEL(m).ExecContext(ctx, db); err != nil {
 				return errors.Wrap(err, "insert channel")
 			}
 		} else {
@@ -95,10 +96,13 @@ func (c *DBChannel) Update(ctx context.Context, offsetID int) error {
 	}
 
 	tbl := table.Channels
-	_, err := tbl.UPDATE().SET(
-		tbl.Offset.SET(Int(int64(c.Offset))),
-		tbl.MinID.SET(Int(int64(c.MinID))),
-	).ExecContext(ctx, db)
+	_, err := tbl.UPDATE().
+		SET(
+			tbl.Offset.SET(Int64(int64(c.Offset))),
+			tbl.MinID.SET(Int64(int64(c.MinID))),
+		).
+		WHERE(tbl.UUID.EQ(Int64(c.UUID))).
+		ExecContext(ctx, db)
 	if err != nil {
 		return errors.Wrap(err, "update channel")
 	}
@@ -118,8 +122,13 @@ func (c *DBChannel) SaveMessage(ctx context.Context, msg *ChannelMessage) error
 
 	tbl := table.ChannelMessages
 
-	_, err := tbl.INSERT(tbl.AllColumns).MODEL(message).ExecContext(ctx, db)
+	_, err := tbl.INSERT(tbl.AllColumns.Except(tbl.ID)).MODEL(message).ExecContext(ctx, db)
 	if err != nil {
+		if e, ok := err.(*pq.Error); ok {
+			if e.Code == "23505" {
+				return nil
+			}
+		}
 		return errors.Wrap(err, "insert message")
 	}
 	return nil
diff --git a/main.go b/main.go
index 51f0aee..4360c62 100644
--- a/main.go
+++ b/main.go
@@ -23,6 +23,11 @@ func main() {
 				return errors.Wrap(err, "load config")
 			}
 
+			fmt.Println("init db")
+			if err := internal.InitDB(config.C.DSN); err != nil {
+				return errors.Wrap(err, "init db")
+			}
+
 			fmt.Println("init client")
 			defer fmt.Println("init client done")
 			return internal.InitClient(config.C)