dify-plugin-daemon/internal/db/pg/pg.go

package pg

import (
	"fmt"
	"strings"
	"time"

	"gorm.io/driver/postgres"
	"gorm.io/gorm"
)

type PGConfig struct {
	Host               string
	Port               int
	DBName             string
	DefaultDBName      string
	User               string
	Pass               string
	SSLMode            string
	MaxIdleConns       int
	MaxOpenConns       int
	ConnMaxLifetime    int
	Charset            string
	Extras             string
	PreparedStatements bool
}

func InitPluginDB(config *PGConfig) (*gorm.DB, error) {
	// first try to connect to target database
	dsn := buildDSN(config, false)
	gormConfig := &gorm.Config{
		PrepareStmt: config.PreparedStatements,
	}
	db, err := gorm.Open(postgres.Open(dsn), gormConfig)
	if err != nil {
		// if connection fails, try to create database
		dsn = buildDSN(config, true)
		db, err = gorm.Open(postgres.Open(dsn), gormConfig)
		if err != nil {
			return nil, err
		}

		pgsqlDB, err := db.DB()
		if err != nil {
			return nil, err
		}
		defer pgsqlDB.Close()

		// check if the db exists
		rows, err := pgsqlDB.Query(fmt.Sprintf("SELECT 1 FROM pg_database WHERE datname = '%s'", config.DBName))
		if err != nil {
			return nil, err
		}

		if !rows.Next() {
			// create database
			_, err = pgsqlDB.Exec(fmt.Sprintf("CREATE DATABASE %s", config.DBName))
			if err != nil {
				return nil, err
			}
		}

		// connect to the new db
		dsn = fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=%s", config.Host, config.Port, config.User, config.Pass, config.DBName, config.SSLMode)
		db, err = gorm.Open(postgres.Open(dsn), gormConfig)
		if err != nil {
			return nil, err
		}
	}

	pgsqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}

	// configure connection pool
	pgsqlDB.SetMaxIdleConns(config.MaxIdleConns)
	pgsqlDB.SetMaxOpenConns(config.MaxOpenConns)
	pgsqlDB.SetConnMaxLifetime(time.Duration(config.ConnMaxLifetime) * time.Second)

	return db, nil
}

func buildDSN(config *PGConfig, useDefaultDB bool) string {
	dsn := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=%s",
		config.Host,
		config.Port,
		config.User,
		config.Pass,
		func() string {
			if useDefaultDB {
				return config.DefaultDBName
			}
			return config.DBName
		}(),
		config.SSLMode,
	)
	if config.Charset != "" {
		dsn = fmt.Sprintf("%s client_encoding=%s", dsn, config.Charset)
	}
	if config.Extras != "" {
		extra := strings.ReplaceAll(config.Extras, "options=", "")
		dsn = fmt.Sprintf("%s options='%s'", dsn, extra)
	}
	return dsn
}