🐛 fix(postgres): 修复删除数据库误判当前连接占用

- PostgreSQL 类数据库 DROP DATABASE 自动切换到维护库执行 - 避免前端传入目标库名时被误判为当前连接正在使用 - 同步修复 ALTER DATABASE RENAME 的同类误判 - 补充 PostgreSQL 删除和重命名数据库回归测试 Close #567
2026-06-20 21:43:56 +08:00 · 2026-06-16 09:07:19 +08:00
parent 0816702084
commit 093b3cae1f
2 changed files with 174 additions and 8 deletions
--- a/internal/app/methods_db.go
+++ b/internal/app/methods_db.go
@@ -213,6 +213,45 @@ func isPostgresSchemaDDLDBType(dbType string) bool {
 	}
 }

+func resolvePGLikeDatabaseDDLCandidates(dbType string, user string) []string {
+	switch resolveDDLDBType(connection.ConnectionConfig{Type: dbType}) {
+	case "kingbase":
+		return []string{"test", "template1", strings.TrimSpace(user)}
+	case "vastbase":
+		return []string{"vastbase", "postgres", "template1", strings.TrimSpace(user)}
+	default:
+		return []string{"postgres", "template1", strings.TrimSpace(user)}
+	}
+}
+
+func resolvePGLikeDatabaseDDLRunConfig(config connection.ConnectionConfig, dbType string, targetDatabase string) connection.ConnectionConfig {
+	runConfig := config
+	target := strings.TrimSpace(targetDatabase)
+	current := strings.TrimSpace(runConfig.Database)
+	if current != "" && !strings.EqualFold(current, target) {
+		return runConfig
+	}
+
+	candidates := resolvePGLikeDatabaseDDLCandidates(dbType, runConfig.User)
+	seen := make(map[string]struct{}, len(candidates))
+	for _, candidate := range candidates {
+		name := strings.TrimSpace(candidate)
+		if name == "" || strings.EqualFold(name, target) {
+			continue
+		}
+		normalized := strings.ToLower(name)
+		if _, ok := seen[normalized]; ok {
+			continue
+		}
+		seen[normalized] = struct{}{}
+		runConfig.Database = name
+		return runConfig
+	}
+
+	runConfig.Database = ""
+	return runConfig
+}
+
 func buildCreateSchemaSQL(dbType string, schemaName string) (string, error) {
 	schemaName = strings.TrimSpace(schemaName)
 	if schemaName == "" {
@@ -571,10 +610,7 @@ func (a *App) RenameDatabase(config connection.ConnectionConfig, oldName string,
 	case "mysql", "mariadb", "oceanbase", "starrocks", "sphinx":
 		return connection.QueryResult{Success: false, Message: "MySQL/MariaDB/OceanBase/StarRocks/Sphinx 不支持直接重命名数据库，请新建库后迁移数据"}
 	case "postgres", "kingbase", "highgo", "vastbase", "opengauss", "gaussdb":
-		if strings.EqualFold(strings.TrimSpace(config.Database), oldName) {
-			return connection.QueryResult{Success: false, Message: "当前连接正在使用目标数据库，请先连接到其他数据库后再重命名"}
-		}
-		runConfig := config
+		runConfig := resolvePGLikeDatabaseDDLRunConfig(config, dbType, oldName)
 		dbInst, err := a.getDatabase(runConfig)
 		if err != nil {
 			return connection.QueryResult{Success: false, Message: err.Error()}
@@ -606,10 +642,7 @@ func (a *App) DropDatabase(config connection.ConnectionConfig, dbName string) co
 		runConfig.Database = ""
 		sql = fmt.Sprintf("DROP DATABASE %s", quoteIdentByType(dbType, dbName))
 	case "postgres", "kingbase", "highgo", "vastbase", "opengauss", "gaussdb":
-		if strings.EqualFold(strings.TrimSpace(config.Database), dbName) {
-			return connection.QueryResult{Success: false, Message: "当前连接正在使用目标数据库，请先连接到其他数据库后再删除"}
-		}
-		runConfig = config
+		runConfig = resolvePGLikeDatabaseDDLRunConfig(config, dbType, dbName)
 		sql = fmt.Sprintf("DROP DATABASE %s", quoteIdentByType(dbType, dbName))
 	default:
 		return connection.QueryResult{Success: false, Message: fmt.Sprintf("当前数据源(%s)暂不支持删除数据库", dbType)}
--- a/internal/app/methods_db_rename_test.go
+++ b/internal/app/methods_db_rename_test.go
@@ -91,3 +91,136 @@ func TestRenameDatabase_DorisUsesNativeRenameSQL(t *testing.T) {
 		t.Fatalf("unexpected Doris rename SQL, want %q got %q", want, fakeDB.execQueries[0])
 	}
 }
+
+func TestResolvePGLikeDatabaseDDLRunConfig_UsesCompatibleMaintenanceDatabase(t *testing.T) {
+	kingbase := resolvePGLikeDatabaseDDLRunConfig(connection.ConnectionConfig{
+		Type:     "kingbase",
+		User:     "system",
+		Database: "test",
+	}, "kingbase", "test")
+	if kingbase.Database != "template1" {
+		t.Fatalf("expected Kingbase target test to use template1 maintenance database, got %q", kingbase.Database)
+	}
+
+	vastbase := resolvePGLikeDatabaseDDLRunConfig(connection.ConnectionConfig{
+		Type:     "vastbase",
+		User:     "vastbase",
+		Database: "tenant_db",
+	}, "vastbase", "tenant_db")
+	if vastbase.Database != "vastbase" {
+		t.Fatalf("expected Vastbase target tenant_db to use vastbase maintenance database, got %q", vastbase.Database)
+	}
+}
+
+func TestDropDatabase_PostgresUsesMaintenanceDatabaseWhenConfigTargetsDroppedDB(t *testing.T) {
+	originalNewDatabaseFunc := newDatabaseFunc
+	originalResolveDialConfigWithProxyFunc := resolveDialConfigWithProxyFunc
+	t.Cleanup(func() {
+		newDatabaseFunc = originalNewDatabaseFunc
+		resolveDialConfigWithProxyFunc = originalResolveDialConfigWithProxyFunc
+	})
+
+	fakeDB := &fakeRenameDatabaseDB{}
+	newDatabaseFunc = func(dbType string) (db.Database, error) {
+		return fakeDB, nil
+	}
+	resolveDialConfigWithProxyFunc = func(raw connection.ConnectionConfig) (connection.ConnectionConfig, error) {
+		return raw, nil
+	}
+
+	app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test"))
+	result := app.DropDatabase(connection.ConnectionConfig{
+		Type:     "postgres",
+		Host:     "127.0.0.1",
+		Port:     5432,
+		User:     "postgres",
+		Database: "test",
+	}, "test")
+
+	if !result.Success {
+		t.Fatalf("expected PostgreSQL drop database success, got failure: %s", result.Message)
+	}
+	if fakeDB.connectConfig.Database != "postgres" {
+		t.Fatalf("expected PostgreSQL drop database to connect to maintenance database postgres, got %q", fakeDB.connectConfig.Database)
+	}
+	if len(fakeDB.execQueries) != 1 {
+		t.Fatalf("expected one drop statement, got %d: %#v", len(fakeDB.execQueries), fakeDB.execQueries)
+	}
+	const want = `DROP DATABASE "test"`
+	if fakeDB.execQueries[0] != want {
+		t.Fatalf("unexpected PostgreSQL drop SQL, want %q got %q", want, fakeDB.execQueries[0])
+	}
+}
+
+func TestDropDatabase_PostgresUsesTemplateWhenDroppingPostgresDatabase(t *testing.T) {
+	originalNewDatabaseFunc := newDatabaseFunc
+	originalResolveDialConfigWithProxyFunc := resolveDialConfigWithProxyFunc
+	t.Cleanup(func() {
+		newDatabaseFunc = originalNewDatabaseFunc
+		resolveDialConfigWithProxyFunc = originalResolveDialConfigWithProxyFunc
+	})
+
+	fakeDB := &fakeRenameDatabaseDB{}
+	newDatabaseFunc = func(dbType string) (db.Database, error) {
+		return fakeDB, nil
+	}
+	resolveDialConfigWithProxyFunc = func(raw connection.ConnectionConfig) (connection.ConnectionConfig, error) {
+		return raw, nil
+	}
+
+	app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test"))
+	result := app.DropDatabase(connection.ConnectionConfig{
+		Type:     "postgres",
+		Host:     "127.0.0.1",
+		Port:     5432,
+		User:     "postgres",
+		Database: "postgres",
+	}, "postgres")
+
+	if !result.Success {
+		t.Fatalf("expected PostgreSQL drop database success, got failure: %s", result.Message)
+	}
+	if fakeDB.connectConfig.Database != "template1" {
+		t.Fatalf("expected PostgreSQL drop postgres database to connect to template1, got %q", fakeDB.connectConfig.Database)
+	}
+}
+
+func TestRenameDatabase_PostgresUsesMaintenanceDatabaseWhenConfigTargetsRenamedDB(t *testing.T) {
+	originalNewDatabaseFunc := newDatabaseFunc
+	originalResolveDialConfigWithProxyFunc := resolveDialConfigWithProxyFunc
+	t.Cleanup(func() {
+		newDatabaseFunc = originalNewDatabaseFunc
+		resolveDialConfigWithProxyFunc = originalResolveDialConfigWithProxyFunc
+	})
+
+	fakeDB := &fakeRenameDatabaseDB{}
+	newDatabaseFunc = func(dbType string) (db.Database, error) {
+		return fakeDB, nil
+	}
+	resolveDialConfigWithProxyFunc = func(raw connection.ConnectionConfig) (connection.ConnectionConfig, error) {
+		return raw, nil
+	}
+
+	app := NewAppWithSecretStore(secretstore.NewUnavailableStore("test"))
+	result := app.RenameDatabase(connection.ConnectionConfig{
+		Type:     "postgres",
+		Host:     "127.0.0.1",
+		Port:     5432,
+		User:     "postgres",
+		Database: "test",
+	}, "test", "test_new")
+
+	if !result.Success {
+		t.Fatalf("expected PostgreSQL rename database success, got failure: %s", result.Message)
+	}
+	if fakeDB.connectConfig.Database != "postgres" {
+		t.Fatalf("expected PostgreSQL rename database to connect to maintenance database postgres, got %q", fakeDB.connectConfig.Database)
+	}
+	if len(fakeDB.execQueries) != 1 {
+		t.Fatalf("expected one rename statement, got %d: %#v", len(fakeDB.execQueries), fakeDB.execQueries)
+	}
+	const want = `ALTER DATABASE "test" RENAME TO "test_new"`
+	if fakeDB.execQueries[0] != want {
+		t.Fatalf("unexpected PostgreSQL rename SQL, want %q got %q", want, fakeDB.execQueries[0])
+	}
+}