From e0b4f96b5bce6e5b9526d0a51a5509507276306a Mon Sep 17 00:00:00 2001 From: JustSong Date: Sat, 12 Aug 2023 19:20:12 +0800 Subject: [PATCH] feat: support PostgreSQL now --- README.md | 6 ++++-- go.mod | 6 +++++- go.sum | 10 ++++++++++ model/main.go | 41 +++++++++++++++++++++++++---------------- 4 files changed, 44 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 1f3a4a5e..a53c8b9d 100644 --- a/README.md +++ b/README.md @@ -276,8 +276,10 @@ graph LR + 例子：`REDIS_CONN_STRING=redis://default:redispw@localhost:49153` 2. `SESSION_SECRET`：设置之后将使用固定的会话密钥，这样系统重新启动后已登录用户的 cookie 将依旧有效。 + 例子：`SESSION_SECRET=random_string` -3. `SQL_DSN`：设置之后将使用指定数据库而非 SQLite，请使用 MySQL 8.0 版本。 - + 例子：`SQL_DSN=root:123456@tcp(localhost:3306)/oneapi` +3. `SQL_DSN`：设置之后将使用指定数据库而非 SQLite，请使用 MySQL 或 PostgreSQL。 + + 例子： + + MySQL：`SQL_DSN=root:123456@tcp(localhost:3306)/oneapi` + + PostgreSQL：`SQL_DSN=postgres://postgres:123456@localhost:5432/oneapi` + 注意需要提前建立数据库 `oneapi`，无需手动建表，程序将自动建表。 + 如果使用本地数据库：部署命令可添加 `--network="host"` 以使得容器内的程序可以访问到宿主机上的 MySQL。 + 如果使用云数据库：如果云服务器需要验证身份，需要在连接参数中添加 `?tls=skip-verify`。 diff --git a/go.mod b/go.mod index 1d08a7d3..79b01f93 100644 --- a/go.mod +++ b/go.mod @@ -18,7 +18,7 @@ require ( golang.org/x/crypto v0.9.0 gorm.io/driver/mysql v1.4.3 gorm.io/driver/sqlite v1.4.3 - gorm.io/gorm v1.24.0 + gorm.io/gorm v1.25.0 ) require ( @@ -36,6 +36,9 @@ require ( github.com/gorilla/context v1.1.1 // indirect github.com/gorilla/securecookie v1.1.1 // indirect github.com/gorilla/sessions v1.2.1 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect + github.com/jackc/pgx/v5 v5.3.1 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect @@ -54,4 +57,5 @@ require ( golang.org/x/text v0.9.0 // indirect google.golang.org/protobuf v1.30.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect + gorm.io/driver/postgres v1.5.2 // indirect ) diff --git a/go.sum b/go.sum index b4281cb6..810e7819 100644 --- a/go.sum +++ b/go.sum @@ -69,6 +69,12 @@ github.com/gorilla/sessions v1.2.1 h1:DHd3rPN5lE3Ts3D8rKkQ8x/0kqfeNmBAaiSi+o7Fsg github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM= github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= +github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.3.1 h1:Fcr8QJ1ZeLi5zsPZqQeUZhNhxfkkKBOgJuYkJHoBOtU= +github.com/jackc/pgx/v5 v5.3.1/go.mod h1:t3JDKnCBlYIc0ewLF0Q7B8MXmoIaBOZj/ic7iHozM/8= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= github.com/jinzhu/now v1.1.4/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= @@ -187,9 +193,13 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gorm.io/driver/mysql v1.4.3 h1:/JhWJhO2v17d8hjApTltKNADm7K7YI2ogkR7avJUL3k= gorm.io/driver/mysql v1.4.3/go.mod h1:sSIebwZAVPiT+27jK9HIwvsqOGKx3YMPmrA3mBJR10c= +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/driver/sqlite v1.4.3 h1:HBBcZSDnWi5BW3B3rwvVTc510KGkBkexlOg0QrmLUuU= gorm.io/driver/sqlite v1.4.3/go.mod h1:0Aq3iPO+v9ZKbcdiz8gLWRw5VOPcBOPUQJFLq5e2ecI= gorm.io/gorm v1.23.8/go.mod h1:l2lP/RyAtc1ynaTjFksBde/O8v9oOGIApu2/xRitmZk= gorm.io/gorm v1.24.0 h1:j/CoiSm6xpRpmzbFJsQHYj+I8bGYWLXVHeYEyyKlF74= gorm.io/gorm v1.24.0/go.mod h1:DVrVomtaYTbqs7gB/x2uVvqnXzv0nqjB396B8cG4dBA= +gorm.io/gorm v1.25.0 h1:+KtYtb2roDz14EQe4bla8CbQlmb9dN3VejSai3lprfU= +gorm.io/gorm v1.25.0/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= rsc.io/pdf v0.1.1/go.mod h1:n8OzWcQ6Sp37PL01nO98y4iUCRdTGarVfzxY20ICaU4= diff --git a/model/main.go b/model/main.go index fcc134d8..213db58c 100644 --- a/model/main.go +++ b/model/main.go @@ -2,10 +2,12 @@ package model import ( "gorm.io/driver/mysql" + "gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" "one-api/common" "os" + "strings" "time" ) @@ -34,28 +36,35 @@ func createRootAccountIfNeed() error { return nil } -func CountTable(tableName string) (num int64) { - DB.Table(tableName).Count(&num) - return -} - -func InitDB() (err error) { - var db *gorm.DB +func chooseDB() (*gorm.DB, error) { if os.Getenv("SQL_DSN") != "" { + dsn := os.Getenv("SQL_DSN") + if strings.HasPrefix(dsn, "postgres://") { + // Use PostgreSQL + common.SysLog("using PostgreSQL as database") + return gorm.Open(postgres.New(postgres.Config{ + DSN: dsn, + PreferSimpleProtocol: true, // disables implicit prepared statement usage + }), &gorm.Config{ + PrepareStmt: true, // precompile SQL + }) + } // Use MySQL common.SysLog("using MySQL as database") - db, err = gorm.Open(mysql.Open(os.Getenv("SQL_DSN")), &gorm.Config{ - PrepareStmt: true, // precompile SQL - }) - } else { - // Use SQLite - common.SysLog("SQL_DSN not set, using SQLite as database") - common.UsingSQLite = true - db, err = gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{ + return gorm.Open(mysql.Open(dsn), &gorm.Config{ PrepareStmt: true, // precompile SQL }) } - common.SysLog("database connected") + // Use SQLite + common.SysLog("SQL_DSN not set, using SQLite as database") + common.UsingSQLite = true + return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{ + PrepareStmt: true, // precompile SQL + }) +} + +func InitDB() (err error) { + db, err := chooseDB() if err == nil { if common.DebugEnabled { db = db.Debug()