115 lines
2.3 KiB
Go
115 lines
2.3 KiB
Go
package db
|
||
|
||
import (
|
||
"database/sql"
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
func driverKey(d string) string {
|
||
return strings.ToLower(strings.TrimSpace(d))
|
||
}
|
||
|
||
func isMySQLFamily(d string) bool {
|
||
switch driverKey(d) {
|
||
case "mysql", "mariadb", "tidb", "oceanbase", "polardb", "tdsql":
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func isPostgresFamily(d string) bool {
|
||
switch driverKey(d) {
|
||
case "postgresql", "postgres", "gaussdb", "opengauss", "kingbase":
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func isMSSQL(d string) bool {
|
||
switch driverKey(d) {
|
||
case "sqlserver", "mssql":
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
// Open 按驱动打开 *sql.DB
|
||
func Open(c ConnReq) (*sql.DB, error) {
|
||
var (
|
||
db *sql.DB
|
||
err error
|
||
)
|
||
switch {
|
||
case isMySQLFamily(c.Driver):
|
||
db, err = openMySQL(c)
|
||
case isPostgresFamily(c.Driver):
|
||
db, err = openPostgres(c)
|
||
case driverKey(c.Driver) == "sqlserver" || driverKey(c.Driver) == "mssql":
|
||
db, err = openMSSQL(c)
|
||
case driverKey(c.Driver) == "oracle":
|
||
db, err = openOracle(c)
|
||
case driverKey(c.Driver) == "sqlite":
|
||
db, err = openSQLite(c)
|
||
default:
|
||
return nil, fmt.Errorf("驱动「%s」尚未实现（当前支持: mysql/mariadb/tidb、postgresql、sqlserver、oracle、sqlite、redis）", c.Driver)
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
db.SetConnMaxLifetime(2 * time.Minute)
|
||
db.SetConnMaxIdleTime(5 * time.Minute) // 防空闲连接堆积
|
||
db.SetMaxOpenConns(5)
|
||
db.SetMaxIdleConns(2)
|
||
return db, nil
|
||
}
|
||
|
||
func TestConnection(c ConnReq) error {
|
||
if isAgentDriver(c.Driver) {
|
||
// JDBC agent 桥的 TCP 由 JVM 侧发起，SSH 隧道（本地端口转发）对 agent 不生效，明确拒绝
|
||
if c.SSHEnabled {
|
||
return fmt.Errorf("%s 走 JDBC agent，暂不支持 SSH 隧道", c.Driver)
|
||
}
|
||
if err := agentPing(c); err != nil {
|
||
return fmt.Errorf("数据库: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
dial, tun, err := ApplySSHTunnel(c)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if tun != nil {
|
||
defer tun.Close()
|
||
}
|
||
if err := testDial(dial); err != nil {
|
||
return fmt.Errorf("数据库: %w", err)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func testDial(c ConnReq) error {
|
||
if isRedis(c.Driver) {
|
||
return testRedis(c)
|
||
}
|
||
if isES(c.Driver) {
|
||
return esPing(c)
|
||
}
|
||
if isInflux(c.Driver) {
|
||
return ifxPing(c)
|
||
}
|
||
if isAgentDriver(c.Driver) {
|
||
return agentPing(c)
|
||
}
|
||
db, err := Open(c)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer db.Close()
|
||
return db.Ping()
|
||
}
|