fix(postgresql): check active sessions before dropping db (#12127)

agent/utils/postgresql/client/local.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/exec"
 	"path"
+	"strconv"
 	"strings"
 	"time"
 
@@ -99,6 +100,13 @@ func (r *Local) CreateUser(info CreateInfo, withDeleteDB bool) error {
 
 func (r *Local) Delete(info DeleteInfo) error {
 	if len(info.Name) != 0 {
+		inUse, err := r.isDatabaseInUse(info.Name, info.Timeout)
+		if err != nil && !info.ForceDelete {
+			return fmt.Errorf("check database connections failed, err: %v", err)
+		}
+		if inUse && !info.ForceDelete {
+			return buserr.WithDetail("ErrInUsed", info.Name, nil)
+		}
 		dropSql := fmt.Sprintf("DROP DATABASE \"%s\"", info.Name)
 		if err := r.ExecSQL(dropSql, info.Timeout); err != nil && !info.ForceDelete {
 			return fmt.Errorf("drop database failed, err: %v", err)
@@ -203,6 +211,30 @@ func (r *Local) SyncDB() ([]SyncDBInfo, error) {
 
 func (r *Local) Close() {}
 
+func (r *Local) isDatabaseInUse(name string, timeout uint) (bool, error) {
+	escapedName := strings.ReplaceAll(name, "'", "''")
+	checkSQL := fmt.Sprintf(
+		"SELECT COUNT(*) FROM pg_stat_activity WHERE datname='%s' AND pid <> pg_backend_pid()",
+		escapedName,
+	)
+	lines, err := r.ExecSQLForRows(checkSQL, timeout)
+	if err != nil {
+		return false, err
+	}
+	for _, line := range lines {
+		countStr := strings.TrimSpace(line)
+		if len(countStr) == 0 {
+			continue
+		}
+		count, parseErr := strconv.Atoi(countStr)
+		if parseErr != nil {
+			return false, parseErr
+		}
+		return count > 0, nil
+	}
+	return false, nil
+}
+
 func (r *Local) ExecSQL(command string, timeout uint) error {
 	itemCommand := r.PrefixCommand[:]
 	itemCommand = append(itemCommand, command)

agent/utils/postgresql/client/remote.go

@@ -95,6 +95,13 @@ func (r *Remote) CreateUser(info CreateInfo, withDeleteDB bool) error {
 
 func (r *Remote) Delete(info DeleteInfo) error {
 	if len(info.Name) != 0 {
+		inUse, err := r.isDatabaseInUse(info.Name, info.Timeout)
+		if err != nil && !info.ForceDelete {
+			return fmt.Errorf("check database connections failed, err: %v", err)
+		}
+		if inUse && !info.ForceDelete {
+			return buserr.WithDetail("ErrInUsed", info.Name, nil)
+		}
 		dropSql := fmt.Sprintf("DROP DATABASE \"%s\"", info.Name)
 		if err := r.ExecSQL(dropSql, info.Timeout); err != nil && !info.ForceDelete {
 			return fmt.Errorf("drop database failed, err: %v", err)
@@ -107,6 +114,24 @@ func (r *Remote) Delete(info DeleteInfo) error {
 	return nil
 }
 
+func (r *Remote) isDatabaseInUse(name string, timeout uint) (bool, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
+	defer cancel()
+
+	var count int
+	if err := r.Client.QueryRowContext(
+		ctx,
+		"SELECT COUNT(*) FROM pg_stat_activity WHERE datname = $1 AND pid <> pg_backend_pid()",
+		name,
+	).Scan(&count); err != nil {
+		return false, err
+	}
+	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+		return false, buserr.New("ErrExecTimeOut")
+	}
+	return count > 0, nil
+}
+
 func (r *Remote) ChangePrivileges(info Privileges) error {
 	super := "SUPERUSER"
 	if !info.SuperUser {