apply postgresql(with bun orm)

main.go

@@ -6,15 +6,17 @@ import (
 	"slices"
 	"strings"
 
+	"sshchat/db"
 	"sshchat/utils"
 
 	"github.com/gliderlabs/ssh"
 	"github.com/oschwald/geoip2-golang"
+	"github.com/uptrace/bun"
 )
 
 var config = utils.GetConfig()
 
-func sessionHandler(s ssh.Session, geoip *geoip2.Reader) {
+func sessionHandler(s ssh.Session, geoip *geoip2.Reader, pgDb *bun.DB) {
 	ptyReq, _, isPty := s.Pty()
 	if !isPty {
 		_, _ = fmt.Fprintln(s, "Err: PTY requires. Reconnect with -t option.")
@@ -61,11 +63,17 @@ func sessionHandler(s ssh.Session, geoip *geoip2.Reader) {
 
 func main() {
 	geoip, err := utils.GetDB(config.Geoip)
-	port := config.Port
 	if err != nil {
 		log.Fatalf("Geoip db is error: %v", err)
 	}
 
+	pgDb, err := db.GetDB(config.PgDsn)
+	if err != nil {
+		log.Fatalf("DB Connection error: %v", err)
+	}
+
+	port := config.Port
+
 	keys, err := utils.CheckHostKey()
 	if err != nil {
 		log.Print("Failed to check SSH keys: generate one.\n", err)
@@ -83,13 +91,17 @@ func main() {
 	s := &ssh.Server{
 		Addr: ":" + port,
 		Handler: func(s ssh.Session) {
-			sessionHandler(s, geoip)
+			sessionHandler(s, geoip, pgDb)
 		},
 	}
 	for _, key := range keys {
 		s.AddHostKey(key)
 	}
 
+	defer func() {
+		_ = pgDb.Close()
+	}()
+
 	log.Print("Listening on :" + port)
 	log.Fatal(s.ListenAndServe())
 }