Skip to content
This repository was archived by the owner on Nov 5, 2024. It is now read-only.

Commit 259b183

Browse files
iwindiwind
committed
使用MySQL、PostgreSQL时可以自动创建数据库
1 parent c4477f9 commit 259b183

7 files changed

Lines changed: 116 additions & 34 deletions

File tree

‎teadb/driver_mysql.go‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"github.com/TeaWeb/code/teadb/shared"
99
_ "github.com/go-sql-driver/mysql"
1010
"github.com/iwind/TeaGo/logs"
11+
"strings"
1112
)
1213

1314
type MySQLDriver struct {
@@ -108,7 +109,7 @@ func (this *MySQLDriver) CreateTable(table string, definitionSQL string) error {
108109
}
109110

110111
// 测试DSN
111-
func (this *MySQLDriver) TestDSN(dsn string) (message string, ok bool) {
112+
func (this *MySQLDriver) TestDSN(dsn string, autoCreateDB bool) (message string, ok bool) {
112113
dbInstance, err := sql.Open("mysql", dsn)
113114
if err != nil {
114115
message = "DSN解析错误：" + err.Error()
@@ -118,8 +119,39 @@ func (this *MySQLDriver) TestDSN(dsn string) (message string, ok bool) {
118119
_ = dbInstance.Close()
119120
}()
120121

122+
// 检查数据库
123+
if autoCreateDB {
124+
index := strings.Index(dsn, "/")
125+
if index == -1 {
126+
message = "invalid dsn"
127+
return
128+
}
129+
database := dsn[index+1:]
130+
index = strings.Index(database, "?")
131+
if index > -1 {
132+
database = database[:index]
133+
}
134+
if len(database) == 0 {
135+
message = "no database defined"
136+
return
137+
}
138+
newDSN := strings.Replace(dsn, "/"+database, "/", -1)
139+
newDBInstance, err := sql.Open("mysql", newDSN)
140+
if err != nil {
141+
message = err.Error()
142+
return
143+
}
144+
_, err = newDBInstance.ExecContext(context.Background(), "CREATE DATABASE IF NOT EXISTS `"+database+"`")
145+
if err != nil {
146+
message = err.Error()
147+
_ = newDBInstance.Close()
148+
return
149+
}
150+
_ = newDBInstance.Close()
151+
}
152+
121153
// 测试创建数据表
122-
_, err = dbInstance.ExecContext(context.Background(), "CREATE TABLE `teaweb.test` ( "+
154+
_, err = dbInstance.ExecContext(context.Background(), "CREATE TABLE `teaweb_test` ( "+
123155
"`id` bigint(20) unsigned NOT NULL AUTO_INCREMENT,"+
124156
"`_id` varchar(24) DEFAULT NULL,"+
125157
"PRIMARY KEY (`id`),"+
@@ -131,14 +163,14 @@ func (this *MySQLDriver) TestDSN(dsn string) (message string, ok bool) {
131163
}
132164

133165
// 测试写入数据表
134-
_, err = dbInstance.ExecContext(context.Background(), "INSERT INTO `teaweb.test` (`_id`) VALUES (\""+shared.NewObjectId().Hex()+"\")")
166+
_, err = dbInstance.ExecContext(context.Background(), "INSERT INTO `teaweb_test` (`_id`) VALUES (\""+shared.NewObjectId().Hex()+"\")")
135167
if err != nil {
136168
message = "尝试写入数据表失败：" + err.Error()
137169
return
138170
}
139171

140172
// 测试删除数据表
141-
_, err = dbInstance.ExecContext(context.Background(), "DROP TABLE `teaweb.test`")
173+
_, err = dbInstance.ExecContext(context.Background(), "DROP TABLE `teaweb_test`")
142174
if err != nil {
143175
message = "尝试删除数据表失败：" + err.Error()
144176
return

‎teadb/driver_mysql_test.go‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -182,19 +182,26 @@ func TestMySQLDriver_asSQL_Update(t *testing.T) {
182182
func TestMySQLDriver_TestDSN(t *testing.T) {
183183
driver := new(MySQLDriver)
184184
{
185-
message, ok := driver.TestDSN("root:abcdef@tcp(127.0.0.1:3306)/teaweb123?charset=utf8mb4&timeout=30s")
185+
message, ok := driver.TestDSN("root:abcdef@tcp(127.0.0.1:3306)/teaweb123?charset=utf8mb4&timeout=30s", false)
186186
t.Log(message, ok)
187187
}
188188
{
189-
message, ok := driver.TestDSN("root:123456@tcp(127.0.0.1:3306)/teaweb123?charset=utf8mb4&timeout=30s")
189+
message, ok := driver.TestDSN("root:123456@tcp(127.0.0.1:3306)/teaweb123?charset=utf8mb4&timeout=30s", false)
190190
t.Log(message, ok)
191191
}
192+
192193
{
193-
message, ok := driver.TestDSN("root:123456@tcp(127.0.0.1:3306)/teaweb?charset=utf8mb4&timeout=30s")
194+
message, ok := driver.TestDSN("root:123456@tcp(127.0.0.1:3306)/teaweb?charset=utf8mb4&timeout=30s", false)
194195
t.Log(message, ok)
195196
}
196197
}
197198

199+
func TestMySQLDriver_TestDSN_Create(t *testing.T) {
200+
driver := new(MySQLDriver)
201+
message, ok := driver.TestDSN("root:123456@tcp(127.0.0.1:3306)/teaweb?charset=utf8mb4&timeout=30s", true)
202+
t.Log(message, ok)
203+
}
204+
198205
func TestMySQLDriver_Ping(t *testing.T) {
199206
driver := new(MySQLDriver)
200207
driver.driver = "mysql"

‎teadb/driver_postgres.go‎

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"github.com/TeaWeb/code/teadb/shared"
99
"github.com/iwind/TeaGo/logs"
1010
_ "github.com/lib/pq"
11+
"net/url"
12+
"strings"
1113
)
1214

1315
type PostgresDriver struct {
@@ -109,7 +111,7 @@ func (this *PostgresDriver) CreateTable(table string, definitionSQL string) erro
109111
}
110112

111113
// 测试DSN
112-
func (this *PostgresDriver) TestDSN(dsn string) (message string, ok bool) {
114+
func (this *PostgresDriver) TestDSN(dsn string, autoCreateDB bool) (message string, ok bool) {
113115
dbInstance, err := sql.Open("postgres", dsn)
114116
if err != nil {
115117
message = "DSN解析错误：" + err.Error()
@@ -119,27 +121,57 @@ func (this *PostgresDriver) TestDSN(dsn string) (message string, ok bool) {
119121
_ = dbInstance.Close()
120122
}()
121123

124+
if autoCreateDB {
125+
u, err := url.Parse(dsn)
126+
if err != nil {
127+
message = err.Error()
128+
return
129+
}
130+
if len(u.Path) <= 1 {
131+
message = "database name should not be empty"
132+
return
133+
}
134+
database := u.Path[1:]
135+
u.Path = "/"
136+
137+
newDBInstance, err := sql.Open("postgres", u.String())
138+
if err != nil {
139+
message = err.Error()
140+
return
141+
}
142+
_, err = newDBInstance.ExecContext(context.Background(), `CREATE DATABASE "`+database+`"`)
143+
if err != nil {
144+
if !strings.Contains(err.Error(), "exists") {
145+
message = err.Error()
146+
_ = newDBInstance.Close()
147+
return
148+
}
149+
}
150+
151+
_ = newDBInstance.Close()
152+
}
153+
122154
// 测试创建数据表
123-
_, err = dbInstance.ExecContext(context.Background(), `CREATE TABLE "public"."teaweb.test" (
155+
_, err = dbInstance.ExecContext(context.Background(), `CREATE TABLE "public"."teaweb_test" (
124156
"id" serial8 primary key,
125157
"_id" varchar(24)
126158
);
127159
128-
CREATE UNIQUE INDEX "teaweb.test_id" ON "public"."teaweb.test" ("_id");`)
160+
CREATE UNIQUE INDEX "teaweb_test_id" ON "public"."teaweb_test" ("_id");`)
129161
if err != nil {
130162
message = "尝试创建数据表失败：" + err.Error()
131163
return
132164
}
133165

134166
// 测试写入数据表
135-
_, err = dbInstance.ExecContext(context.Background(), "INSERT INTO \"teaweb.test\" (\"_id\") VALUES ('"+shared.NewObjectId().Hex()+"')")
167+
_, err = dbInstance.ExecContext(context.Background(), "INSERT INTO \"teaweb_test\" (\"_id\") VALUES ('"+shared.NewObjectId().Hex()+"')")
136168
if err != nil {
137169
message = "尝试写入数据表失败：" + err.Error()
138170
return
139171
}
140172

141173
// 测试删除数据表
142-
_, err = dbInstance.ExecContext(context.Background(), "DROP TABLE \"teaweb.test\"")
174+
_, err = dbInstance.ExecContext(context.Background(), "DROP TABLE \"teaweb_test\"")
143175
if err != nil {
144176
message = "尝试删除数据表失败：" + err.Error()
145177
return

‎teadb/driver_postgres_test.go‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,15 +81,23 @@ func TestPostgresDriver_CreateTable(t *testing.T) {
8181
func TestPostgresDriver_TestDSN(t *testing.T) {
8282
driver := new(PostgresDriver)
8383
{
84-
message, ok := driver.TestDSN("postgres://postgres:@127.0.0.1:5432/teaweb?sslmode=disable")
84+
message, ok := driver.TestDSN("postgres://postgres:@127.0.0.1:5432/teaweb?sslmode=disable", false)
8585
t.Log(message, ok)
8686
}
8787
{
88-
message, ok := driver.TestDSN("postgres://postgres:123456@127.0.0.1:5432/teaweb123?sslmode=disable")
88+
message, ok := driver.TestDSN("postgres://postgres:123456@127.0.0.1:5432/teaweb123?sslmode=disable", false)
8989
t.Log(message, ok)
9090
}
9191
{
92-
message, ok := driver.TestDSN("postgres://postgres:123456@127.0.0.1:5432/teaweb?sslmode=disable")
92+
message, ok := driver.TestDSN("postgres://postgres:123456@127.0.0.1:5432/teaweb?sslmode=disable", false)
93+
t.Log(message, ok)
94+
}
95+
}
96+
97+
func TestPostgresDriver_TestDSN_Create(t *testing.T) {
98+
driver := new(PostgresDriver)
99+
{
100+
message, ok := driver.TestDSN("postgres://postgres:123456@127.0.0.1:5432/teaweb?sslmode=disable", true)
93101
t.Log(message, ok)
94102
}
95103
}

‎teaweb/actions/default/install/test.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,12 @@ type TestAction actions.Action
1111

1212
// 测试数据库
1313
func (this *TestAction) RunPost(params struct {
14-
DBType string `alias:"dbType"`
15-
Addr string
16-
Username string
17-
Password string
18-
DBName string `alias:"dbName"`
14+
DBType string `alias:"dbType"`
15+
Addr string
16+
Username string
17+
Password string
18+
DBName string `alias:"dbName"`
19+
AutoCreate bool
1920

2021
// mongodb专用
2122
AuthEnabled bool
@@ -57,7 +58,7 @@ func (this *TestAction) RunPost(params struct {
5758
config.DBName = params.DBName
5859
dsn := config.ComposeDSN()
5960
driver := new(teadb.MySQLDriver)
60-
message, ok := driver.TestDSN(dsn)
61+
message, ok := driver.TestDSN(dsn, params.AutoCreate)
6162
if ok {
6263
this.Success()
6364
}
@@ -71,7 +72,7 @@ func (this *TestAction) RunPost(params struct {
7172
config.DBName = params.DBName
7273
dsn := config.ComposeDSN()
7374
driver := new(teadb.PostgresDriver)
74-
message, ok := driver.TestDSN(dsn)
75+
message, ok := driver.TestDSN(dsn, params.AutoCreate)
7576
if ok {
7677
this.Success()
7778
}

‎teaweb/actions/default/settings/mysql/test.go‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,12 @@ type TestAction actions.Action
1111

1212
// 测试数据库
1313
func (this *TestAction) RunPost(params struct {
14-
DBType string `alias:"dbType"`
15-
Addr string
16-
Username string
17-
Password string
18-
DBName string `alias:"dbName"`
14+
DBType string `alias:"dbType"`
15+
Addr string
16+
Username string
17+
Password string
18+
DBName string `alias:"dbName"`
19+
AutoCreate bool
1920
}) {
2021
params.Addr = teautils.FormatAddress(params.Addr)
2122

@@ -34,7 +35,7 @@ func (this *TestAction) RunPost(params struct {
3435
config.DBName = params.DBName
3536
dsn := config.ComposeDSN()
3637
driver := new(teadb.MySQLDriver)
37-
message, ok := driver.TestDSN(dsn)
38+
message, ok := driver.TestDSN(dsn, params.AutoCreate)
3839
if ok {
3940
this.Success()
4041
}

‎teaweb/actions/default/settings/postgres/test.go‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,12 @@ type TestAction actions.Action
1111

1212
// 测试数据库
1313
func (this *TestAction) RunPost(params struct {
14-
DBType string `alias:"dbType"`
15-
Addr string
16-
Username string
17-
Password string
18-
DBName string `alias:"dbName"`
14+
DBType string `alias:"dbType"`
15+
Addr string
16+
Username string
17+
Password string
18+
DBName string `alias:"dbName"`
19+
AutoCreate bool
1920
}) {
2021
params.Addr = teautils.FormatAddress(params.Addr)
2122

@@ -34,7 +35,7 @@ func (this *TestAction) RunPost(params struct {
3435
config.DBName = params.DBName
3536
dsn := config.ComposeDSN()
3637
driver := new(teadb.PostgresDriver)
37-
message, ok := driver.TestDSN(dsn)
38+
message, ok := driver.TestDSN(dsn, params.AutoCreate)
3839
if ok {
3940
this.Success()
4041
}

0 commit comments

Comments
 (0)