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]) + } +}