fix: #658 (#659)

* fix:  Fix the issue of the program getting stuck when starting. Add a MySQL connection timeout configuration

* refactoring the DSN query string construction to use net/url.Values~

* fix : Add unit support for mysql database connection timeout configuration items

Author: shootercheng

.env.example

@@ -139,6 +139,12 @@ DB_CHARSET=
 # "info" level displays SQL statements
 # Refer to gorm.io/gorm/logger#LogLevel for detailed information.
 DB_GORM_LOG_LEVEL=
+# database connection timeout, default 5s
+DB_CONNECT_TIMEOUT=5s
+# database read timeout, default 30s
+DB_READ_TIMEOUT=30s
+# database write timeout, default 30s
+DB_WRITE_TIMEOUT=30s
 
 DIFY_INVOCATION_CONNECTION_IDLE_TIMEOUT=120
 

internal/db/init.go

@@ -98,6 +98,9 @@ func Init(config *app.Config) {
 			Charset:         config.DBCharset,
 			Extras:          config.DBExtras,
 			LogLevel:        config.DBGormLogLevel,
+			ConnectTimeout:  config.DBConnectTimeout,
+			ReadTimeout:     config.DBReadTimeout,
+			WriteTimeout:    config.DBWriteTimeout,
 		})
 	default:
 		log.Panic("unsupported database type", "type", config.DBType)

internal/db/mysql/mysql.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"time"
 
+	"net/url"
+
 	gormConfig "github.com/langgenius/dify-plugin-daemon/internal/db/config"
 	"gorm.io/driver/mysql"
 	"gorm.io/gorm"
@@ -24,17 +26,23 @@ type MySQLConfig struct {
 	Charset         string
 	Extras          string
 	LogLevel        string
+	ConnectTimeout  time.Duration
+	ReadTimeout     time.Duration
+	WriteTimeout    time.Duration
 }
 
 func InitPluginDB(config *MySQLConfig) (*gorm.DB, error) {
 	// TODO: MySQL dose not support DB_EXTRAS now
 	initializer := mysqlDbInitializer{
-		host:     config.Host,
-		port:     config.Port,
-		user:     config.User,
-		password: config.Pass,
-		sslMode:  config.SSLMode,
-		logLevel: config.LogLevel,
+		host:           config.Host,
+		port:           config.Port,
+		user:           config.User,
+		password:       config.Pass,
+		sslMode:        config.SSLMode,
+		logLevel:       config.LogLevel,
+		connectTimeout: config.ConnectTimeout,
+		readTimeout:    config.ReadTimeout,
+		writeTimeout:   config.WriteTimeout,
 	}
 
 	// first try to connect to target database
@@ -73,16 +81,27 @@ func InitPluginDB(config *MySQLConfig) (*gorm.DB, error) {
 
 // mysqlDbInitializer initializes database for MySQL.
 type mysqlDbInitializer struct {
-	host     string
-	port     int
-	user     string
-	password string
-	sslMode  string
-	logLevel string
+	host           string
+	port           int
+	user           string
+	password       string
+	sslMode        string
+	logLevel       string
+	connectTimeout time.Duration
+	readTimeout    time.Duration
+	writeTimeout   time.Duration
 }
 
 func (m *mysqlDbInitializer) connect(dbName string) (*gorm.DB, error) {
-	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=true&tls=%v", m.user, m.password, m.host, m.port, dbName, m.sslMode == "require")
+	query := url.Values{}
+	query.Set("charset", "utf8mb4")
+	query.Set("parseTime", "true")
+	query.Set("tls", fmt.Sprintf("%v", m.sslMode == "require"))
+	query.Set("timeout", fmt.Sprintf("%s", m.connectTimeout))
+	query.Set("readTimeout", fmt.Sprintf("%s", m.readTimeout))
+	query.Set("writeTimeout", fmt.Sprintf("%s", m.writeTimeout))
+	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?%s",
+		m.user, m.password, m.host, m.port, dbName, query.Encode())
 
 	config := &gorm.Config{}
 	if m.logLevel != "" {

internal/types/app/config.go

@@ -7,6 +7,8 @@ import (
 	"os"
 	"strings"
 
+	"time"
+
 	"github.com/go-playground/validator/v10"
 	gormConfig "github.com/langgenius/dify-plugin-daemon/internal/db/config"
 	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
@@ -147,12 +149,15 @@ type Config struct {
 	DBSslMode         string `envconfig:"DB_SSL_MODE" validate:"required,oneof=disable require"`
 
 	// database connection pool settings
-	DBMaxIdleConns    int    `envconfig:"DB_MAX_IDLE_CONNS" default:"10"`
-	DBMaxOpenConns    int    `envconfig:"DB_MAX_OPEN_CONNS" default:"30"`
-	DBConnMaxLifetime int    `envconfig:"DB_CONN_MAX_LIFETIME" default:"3600"`
-	DBExtras          string `envconfig:"DB_EXTRAS"`
-	DBCharset         string `envconfig:"DB_CHARSET"`
-	DBGormLogLevel    string `envconfig:"DB_GORM_LOG_LEVEL"`
+	DBMaxIdleConns    int           `envconfig:"DB_MAX_IDLE_CONNS" default:"10"`
+	DBMaxOpenConns    int           `envconfig:"DB_MAX_OPEN_CONNS" default:"30"`
+	DBConnMaxLifetime int           `envconfig:"DB_CONN_MAX_LIFETIME" default:"3600"`
+	DBExtras          string        `envconfig:"DB_EXTRAS"`
+	DBCharset         string        `envconfig:"DB_CHARSET"`
+	DBGormLogLevel    string        `envconfig:"DB_GORM_LOG_LEVEL"`
+	DBConnectTimeout  time.Duration `envconfig:"DB_CONNECT_TIMEOUT" default:"5s"`
+	DBReadTimeout     time.Duration `envconfig:"DB_READ_TIMEOUT" default:"30s"`
+	DBWriteTimeout    time.Duration `envconfig:"DB_WRITE_TIMEOUT" default:"30s"`
 
 	// persistence storage
 	PersistenceStoragePath    string `envconfig:"PERSISTENCE_STORAGE_PATH"`