feat: Add PostgreSQL support and refactor MySQL DSN handling (#284)

cmd/apimain.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"github.com/go-redis/redis/v8"
 	"github.com/lejianwen/rustdesk-api/v2/config"
 	"github.com/lejianwen/rustdesk-api/v2/global"
@@ -140,18 +141,40 @@ func InitGlobal() {
 	}
 	//gorm
 	if global.Config.Gorm.Type == config.TypeMysql {
-		dns := global.Config.Mysql.Username + ":" + global.Config.Mysql.Password + "@(" + global.Config.Mysql.Addr + ")/" + global.Config.Mysql.Dbname + "?charset=utf8mb4&parseTime=True&loc=Local"
+
+		dsn := fmt.Sprintf("%s:%s@(%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
+			global.Config.Mysql.Username,
+			global.Config.Mysql.Password,
+			global.Config.Mysql.Addr,
+			global.Config.Mysql.Dbname,
+		)
+
 		global.DB = orm.NewMysql(&orm.MysqlConfig{
-			Dns:          dns,
+			Dsn:          dsn,
 			MaxIdleConns: global.Config.Gorm.MaxIdleConns,
 			MaxOpenConns: global.Config.Gorm.MaxOpenConns,
-		})
+		}, global.Logger)
+	} else if global.Config.Gorm.Type == config.TypePostgresql {
+		dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s TimeZone=%s",
+			global.Config.Postgresql.Host,
+			global.Config.Postgresql.Port,
+			global.Config.Postgresql.User,
+			global.Config.Postgresql.Password,
+			global.Config.Postgresql.Dbname,
+			global.Config.Postgresql.Sslmode,
+			global.Config.Postgresql.TimeZone,
+		)
+		global.DB = orm.NewPostgresql(&orm.PostgresqlConfig{
+			Dsn:          dsn,
+			MaxIdleConns: global.Config.Gorm.MaxIdleConns,
+			MaxOpenConns: global.Config.Gorm.MaxOpenConns,
+		}, global.Logger)
 	} else {
 		//sqlite
 		global.DB = orm.NewSqlite(&orm.SqliteConfig{
 			MaxIdleConns: global.Config.Gorm.MaxIdleConns,
 			MaxOpenConns: global.Config.Gorm.MaxOpenConns,
-		})
+		}, global.Logger)
 	}
 
 	//validator
@@ -197,11 +220,17 @@ func DatabaseAutoUpdate() {
 		if dbName == "" {
 			dbName = global.Config.Mysql.Dbname
 			// 移除 DSN 中的数据库名称，以便初始连接时不指定数据库
-			dsnWithoutDB := global.Config.Mysql.Username + ":" + global.Config.Mysql.Password + "@(" + global.Config.Mysql.Addr + ")/?charset=utf8mb4&parseTime=True&loc=Local"
+			dsnWithoutDB := fmt.Sprintf("%s:%s@(%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
+				global.Config.Mysql.Username,
+				global.Config.Mysql.Password,
+				global.Config.Mysql.Addr,
+				"",
+			)
+
 			//新链接
 			dbWithoutDB := orm.NewMysql(&orm.MysqlConfig{
-				Dns: dsnWithoutDB,
+				Dsn: dsnWithoutDB,
-			})
+			}, global.Logger)
 			// 获取底层的 *sql.DB 对象，并确保在程序退出时关闭连接
 			sqlDBWithoutDB, err := dbWithoutDB.DB()
 			if err != nil {

conf/config.yaml

@@ -11,7 +11,7 @@ app:
   disable-pwd-login: false #禁用密码登录
 
 admin:
-  title: "RustDesk Api Admin"
+  title: "RustDesk API Admin"
   hello-file: "./conf/admin/hello.html"  #优先使用file
   hello: ""
   # ID Server and Relay Server ports https://github.com/lejianwen/rustdesk-api/issues/257
@@ -31,6 +31,16 @@ mysql:
   password: ""
   addr: ""
   dbname: ""
+
+postgresql:
+  host: "127.0.0.1"
+  port: "5432"
+  user: ""
+  password: ""
+  dbname: "postgres"
+  sslmode: "disable" # disable, require, verify-ca, verify-full
+  time-zone: "Asia/Shanghai" # Time zone for PostgreSQL connection
+
 rustdesk:
   id-server: "192.168.1.66:21116"
   relay-server: "192.168.1.66:21117"

config/config.go

@@ -32,20 +32,21 @@ type Admin struct {
 	RelayServerPort int    `mapstructure:"relay-server-port"`
 }
 type Config struct {
-	Lang     string `mapstructure:"lang"`
+	Lang       string `mapstructure:"lang"`
-	App      App
+	App        App
-	Admin    Admin
+	Admin      Admin
-	Gorm     Gorm
+	Gorm       Gorm
-	Mysql    Mysql
+	Mysql      Mysql
-	Gin      Gin
+	Postgresql Postgresql
-	Logger   Logger
+	Gin        Gin
-	Redis    Redis
+	Logger     Logger
-	Cache    Cache
+	Redis      Redis
-	Oss      Oss
+	Cache      Cache
-	Jwt      Jwt
+	Oss        Oss
-	Rustdesk Rustdesk
+	Jwt        Jwt
-	Proxy    Proxy
+	Rustdesk   Rustdesk
-	Ldap     Ldap
+	Proxy      Proxy
+	Ldap       Ldap
 }
 
 func (a *Admin) Init() {

config/gorm.go

@@ -1,8 +1,9 @@
 package config
 
 const (
-	TypeSqlite = "sqlite"
+	TypeSqlite     = "sqlite"
-	TypeMysql  = "mysql"
+	TypeMysql      = "mysql"
+	TypePostgresql = "postgresql"
 )
 
 type Gorm struct {
@@ -17,3 +18,13 @@ type Mysql struct {
 	Password string `mapstructure:"password"`
 	Dbname   string `mapstructure:"dbname"`
 }
+
+type Postgresql struct {
+	Host     string `mapstructure:"host"`
+	Port     string `mapstructure:"port"`
+	User     string `mapstructure:"user"`
+	Password string `mapstructure:"password"`
+	Dbname   string `mapstructure:"dbname"`
+	Sslmode  string `mapstructure:"sslmode"`   // "disable", "require", "verify-ca", "verify-full"
+	TimeZone string `mapstructure:"time-zone"` // e.g., "Asia/Shanghai"
+}

go.mod

@@ -1,19 +1,23 @@
 module github.com/lejianwen/rustdesk-api/v2
 
-go 1.22
+go 1.23
+
+toolchain go1.23.10
 
 require (
 	github.com/BurntSushi/toml v1.3.2
 	github.com/antonfisher/nested-logrus-formatter v1.3.1
-	github.com/fsnotify/fsnotify v1.5.1
+	github.com/coreos/go-oidc/v3 v3.12.0
 	github.com/fvbock/endless v0.0.0-20170109170031-447134032cb6
 	github.com/gin-gonic/gin v1.9.0
+	github.com/go-ldap/ldap/v3 v3.4.10
 	github.com/go-playground/locales v0.14.1
 	github.com/go-playground/universal-translator v0.18.1
 	github.com/go-playground/validator/v10 v10.26.0
 	github.com/go-redis/redis/v8 v8.11.4
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/google/uuid v1.6.0
+	github.com/mojocn/base64Captcha v1.3.6
 	github.com/nicksnyder/go-i18n/v2 v2.4.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/spf13/cobra v1.8.1
@@ -24,8 +28,9 @@ require (
 	golang.org/x/oauth2 v0.23.0
 	golang.org/x/text v0.22.0
 	gorm.io/driver/mysql v1.5.7
+	gorm.io/driver/postgres v1.6.0
 	gorm.io/driver/sqlite v1.5.6
-	gorm.io/gorm v1.25.7
+	gorm.io/gorm v1.25.10
 )
 
 require (
@@ -36,13 +41,12 @@ require (
 	github.com/bytedance/sonic v1.8.0 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
-	github.com/coreos/go-oidc/v3 v3.12.0 // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
+	github.com/fsnotify/fsnotify v1.5.1 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.8 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/go-asn1-ber/asn1-ber v1.5.7 // indirect
 	github.com/go-jose/go-jose/v4 v4.0.2 // indirect
-	github.com/go-ldap/ldap/v3 v3.4.10 // indirect
 	github.com/go-openapi/jsonpointer v0.19.5 // indirect
 	github.com/go-openapi/jsonreference v0.19.6 // indirect
 	github.com/go-openapi/spec v0.20.4 // indirect
@@ -52,6 +56,10 @@ require (
 	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/pgx/v5 v5.6.0 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
@@ -65,9 +73,9 @@ require (
 	github.com/mitchellh/mapstructure v1.4.2 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/mojocn/base64Captcha v1.3.6 // indirect
 	github.com/pelletier/go-toml v1.9.4 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.6 // indirect
+	github.com/rogpeppe/go-internal v1.14.1 // indirect
 	github.com/spf13/afero v1.6.0 // indirect
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
@@ -79,8 +87,9 @@ require (
 	golang.org/x/crypto v0.33.0 // indirect
 	golang.org/x/image v0.13.0 // indirect
 	golang.org/x/net v0.34.0 // indirect
+	golang.org/x/sync v0.11.0 // indirect
 	golang.org/x/sys v0.30.0 // indirect
-	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
+	golang.org/x/tools v0.26.0 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect
 	gopkg.in/ini.v1 v1.63.2 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect

lib/orm/mysql.go

@@ -2,7 +2,6 @@ package orm
 
 import (
 	"fmt"
-	"github.com/lejianwen/rustdesk-api/v2/global"
 	"gorm.io/driver/mysql"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
@@ -10,14 +9,14 @@ import (
 )
 
 type MysqlConfig struct {
-	Dns          string
+	Dsn          string
 	MaxIdleConns int
 	MaxOpenConns int
 }
 
-func NewMysql(mysqlConf *MysqlConfig) *gorm.DB {
+func NewMysql(mysqlConf *MysqlConfig, logwriter logger.Writer) *gorm.DB {
 	db, err := gorm.Open(mysql.New(mysql.Config{
-		DSN:               mysqlConf.Dns, // DSN data source name
+		DSN:               mysqlConf.Dsn, // DSN data source name
 		DefaultStringSize: 256,           // string 类型字段的默认长度
 		//DisableDatetimePrecision:  true,                    // 禁用 datetime 精度，MySQL 5.6 之前的数据库不支持
 		//DontSupportRenameIndex:    true,                    // 重命名索引时采用删除并新建的方式，MySQL 5.7 之前的数据库和 MariaDB 不支持重命名索引
@@ -26,7 +25,7 @@ func NewMysql(mysqlConf *MysqlConfig) *gorm.DB {
 	}), &gorm.Config{
 		DisableForeignKeyConstraintWhenMigrating: true,
 		Logger: logger.New(
-			global.Logger, // io writer
+			logwriter, // io writer
 			logger.Config{
 				SlowThreshold:             time.Second, // Slow SQL threshold
 				LogLevel:                  logger.Warn, // Log level

lib/orm/postgresql.go

@@ -0,0 +1,45 @@
+package orm
+
+import (
+	"fmt"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+	"time"
+)
+
+type PostgresqlConfig struct {
+	Dsn          string
+	MaxIdleConns int
+	MaxOpenConns int
+}
+
+func NewPostgresql(conf *PostgresqlConfig, logwriter logger.Writer) *gorm.DB {
+	db, err := gorm.Open(postgres.Open(conf.Dsn), &gorm.Config{
+		DisableForeignKeyConstraintWhenMigrating: true,
+		Logger: logger.New(
+			logwriter, // io writer
+			logger.Config{
+				SlowThreshold: time.Second, // Slow SQL threshold
+				LogLevel:      logger.Warn, // Log level
+				//IgnoreRecordNotFoundError: true,        // Ignore ErrRecordNotFound error for logger
+				ParameterizedQueries: true, // Don't include params in the SQL log
+				Colorful:             true,
+			},
+		),
+	})
+	if err != nil {
+		fmt.Println(err)
+	}
+	sqlDB, err2 := db.DB()
+	if err2 != nil {
+		fmt.Println(err2)
+	}
+	// SetMaxIdleConns 设置空闲连接池中连接的最大数量
+	sqlDB.SetMaxIdleConns(conf.MaxIdleConns)
+
+	// SetMaxOpenConns 设置打开数据库连接的最大数量。
+	sqlDB.SetMaxOpenConns(conf.MaxOpenConns)
+
+	return db
+}

lib/orm/sqlite.go

@@ -2,7 +2,6 @@ package orm
 
 import (
 	"fmt"
-	"github.com/lejianwen/rustdesk-api/v2/global"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
@@ -14,11 +13,11 @@ type SqliteConfig struct {
 	MaxOpenConns int
 }
 
-func NewSqlite(sqliteConf *SqliteConfig) *gorm.DB {
+func NewSqlite(sqliteConf *SqliteConfig, logwriter logger.Writer) *gorm.DB {
 	db, err := gorm.Open(sqlite.Open("./data/rustdeskapi.db"), &gorm.Config{
 		DisableForeignKeyConstraintWhenMigrating: true,
 		Logger: logger.New(
-			global.Logger, // io writer
+			logwriter, // io writer
 			logger.Config{
 				SlowThreshold:             time.Second, // Slow SQL threshold
 				LogLevel:                  logger.Warn, // Log level