From 093b3cae1f99e65a23dae7780e0090429d3a9561 Mon Sep 17 00:00:00 2001 From: Syngnat Date: Tue, 16 Jun 2026 09:07:19 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(postgres):=20=E4=BF=AE?= =?UTF-8?q?=E5=A4=8D=E5=88=A0=E9=99=A4=E6=95=B0=E6=8D=AE=E5=BA=93=E8=AF=AF?= =?UTF-8?q?=E5=88=A4=E5=BD=93=E5=89=8D=E8=BF=9E=E6=8E=A5=E5=8D=A0=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - PostgreSQL 类数据库 DROP DATABASE 自动切换到维护库执行 - 避免前端传入目标库名时被误判为当前连接正在使用 - 同步修复 ALTER DATABASE RENAME 的同类误判 - 补充 PostgreSQL 删除和重命名数据库回归测试 Close #567 --- internal/app/methods_db.go | 49 +++++++-- internal/app/methods_db_rename_test.go | 133 +++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 8 deletions(-) diff --git a/internal/app/methods_db.go b/internal/app/methods_db.go index 67c1001..e2035ca 100644 --- 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)} diff --git a/internal/app/methods_db_rename_test.go b/internal/app/methods_db_rename_test.go index 952ebf8..91db6b9 100644 --- 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]) + } +}