[优化] 优化脚本

This commit is contained in:
ryan
2026-06-01 20:07:50 +08:00
parent e00d67f2d9
commit 2635a47d29
4 changed files with 138 additions and 3 deletions
+52 -2
View File
@@ -592,8 +592,22 @@ func TestEnsureDatabaseSchemaUpToDateV16BackfillsNodeColumnsWhenNewColumnsAlread
if node.ExtVersion != "openresty/1.25.3" {
t.Fatalf("unexpected ext_version: got %q", node.ExtVersion)
}
if !db.Migrator().HasColumn(&Node{}, "agent_token") {
t.Fatal("expected migration to keep legacy nodes.agent_token column")
for _, column := range []string{
"agent_token",
"agent_version",
"nginx_version",
"relay_version",
"relay_frp_version",
"relay_frps_connections",
"relay_frps_proxy_count",
} {
exists, err := databaseColumnExists(db, "nodes", column)
if err != nil {
t.Fatalf("inspect legacy nodes.%s: %v", column, err)
}
if exists {
t.Fatalf("expected migration to drop legacy nodes.%s column", column)
}
}
version, exists, err := loadDatabaseSchemaVersion(db)
if err != nil {
@@ -607,6 +621,42 @@ func TestEnsureDatabaseSchemaUpToDateV16BackfillsNodeColumnsWhenNewColumnsAlread
}
}
func TestEnsureDatabaseSchemaUpToDateV16DropsLegacyNodeColumnsWhenAlreadyCurrent(t *testing.T) {
db := openBareTestSQLiteDB(t, "node-v16-current-legacy-columns.db")
if err := registerSharding(db, "sqlite"); err != nil {
t.Fatalf("register sharding: %v", err)
}
if err := applyCurrentSchema(db, "sqlite"); err != nil {
t.Fatalf("apply current schema: %v", err)
}
for _, stmt := range []string{
`ALTER TABLE nodes ADD COLUMN agent_token text`,
`ALTER TABLE nodes ADD COLUMN agent_version text`,
`ALTER TABLE nodes ADD COLUMN nginx_version text`,
} {
if err := db.Exec(stmt).Error; err != nil {
t.Fatalf("prepare legacy node column with %q: %v", stmt, err)
}
}
if err := saveDatabaseSchemaVersion(db, currentDatabaseSchemaVersion); err != nil {
t.Fatalf("save schema version: %v", err)
}
if err := ensureDatabaseSchemaUpToDate(db, "sqlite"); err != nil {
t.Fatalf("ensureDatabaseSchemaUpToDate: %v", err)
}
for _, column := range []string{"agent_token", "agent_version", "nginx_version"} {
exists, err := databaseColumnExists(db, "nodes", column)
if err != nil {
t.Fatalf("inspect legacy nodes.%s: %v", column, err)
}
if exists {
t.Fatalf("expected current-schema cleanup to drop legacy nodes.%s column", column)
}
}
}
func TestAllRegisteredMigrationsHaveValidationDefined(t *testing.T) {
ctx := databaseSchemaMigrationContext{}
for _, migration := range databaseSchemaMigrations() {
@@ -18,6 +18,7 @@ type Context interface {
BackfillProxyRouteDomainCertificateFields(db *gorm.DB) error
EnsureDefaultGitHubAuthSource(db *gorm.DB) error
EnsureDefaultWAFRuleGroup(db *gorm.DB) error
DropLegacyNodeColumns(db *gorm.DB, backend string) error
ValidateDatabaseSchemaVersion(db *gorm.DB, backend string, version int) error
}
+16
View File
@@ -77,6 +77,9 @@ func migrateV16(ctx Context, db *gorm.DB, backend string) error {
return fmt.Errorf("backfill nodes.ext_version from nginx_version: %w", err)
}
}
if err := ctx.DropLegacyNodeColumns(db, backend); err != nil {
return err
}
if err := db.Exec("UPDATE nodes SET node_type = 'edge_node' WHERE node_type = '' OR node_type IS NULL").Error; err != nil {
return fmt.Errorf("backfill nodes.node_type: %w", err)
@@ -145,6 +148,19 @@ func validateV16(ctx Context, db *gorm.DB, backend string) error {
if migrator.HasTable(&tunnelV16{}) {
return fmt.Errorf("table tunnels should not exist in v16")
}
for _, column := range []string{
"agent_token",
"agent_version",
"nginx_version",
"relay_version",
"relay_frp_version",
"relay_frps_connections",
"relay_frps_proxy_count",
} {
if migrator.HasColumn(&nodeV16{}, column) {
return fmt.Errorf("column nodes.%s should not exist in v16", column)
}
}
if !migrator.HasTable(&wafIPGroupV16{}) {
return fmt.Errorf("table waf_ip_groups is missing")
}
+69 -1
View File
@@ -60,6 +60,10 @@ func (databaseSchemaMigrationContext) EnsureDefaultWAFRuleGroup(db *gorm.DB) err
return ensureDefaultWAFRuleGroup(db)
}
func (databaseSchemaMigrationContext) DropLegacyNodeColumns(db *gorm.DB, backend string) error {
return dropLegacyNodeColumns(db, backend)
}
func (databaseSchemaMigrationContext) ValidateDatabaseSchemaVersion(db *gorm.DB, backend string, version int) error {
switch version {
case 7:
@@ -181,6 +185,50 @@ func applyCurrentSchema(db *gorm.DB, backend string) error {
return applyCurrentSchemaExcept(db, backend)
}
func databaseColumnExists(db *gorm.DB, tableName string, columnName string) (bool, error) {
columnTypes, err := db.Migrator().ColumnTypes(tableName)
if err != nil {
return false, err
}
for _, columnType := range columnTypes {
if strings.EqualFold(columnType.Name(), columnName) {
return true, nil
}
}
return false, nil
}
func dropLegacyNodeColumns(db *gorm.DB, backend string) error {
if db == nil || !db.Migrator().HasTable(&Node{}) {
return nil
}
legacyColumns := []struct {
column string
}{
{column: "agent_token"},
{column: "agent_version"},
{column: "nginx_version"},
{column: "relay_version"},
{column: "relay_frp_version"},
{column: "relay_frps_connections"},
{column: "relay_frps_proxy_count"},
}
for _, item := range legacyColumns {
exists, err := databaseColumnExists(db, "nodes", item.column)
if err != nil {
return fmt.Errorf("inspect legacy nodes.%s failed: %w", item.column, err)
}
if !exists {
continue
}
if err := db.Exec(fmt.Sprintf(`ALTER TABLE "nodes" DROP COLUMN "%s"`, item.column)).Error; err != nil {
return fmt.Errorf("drop legacy nodes.%s failed: %w", item.column, err)
}
}
_ = backend
return nil
}
func applyCurrentSchemaExcept(db *gorm.DB, backend string, excludedTables ...string) error {
excluded := make(map[string]bool, len(excludedTables))
for _, table := range excludedTables {
@@ -1221,6 +1269,23 @@ func validateDatabaseSchemaV16(db *gorm.DB, backend string) error {
if db.Migrator().HasColumn(&ProxyRoute{}, "tunnel_id") {
return fmt.Errorf("column proxy_routes.tunnel_id should not exist in v16")
}
for _, column := range []string{
"agent_token",
"agent_version",
"nginx_version",
"relay_version",
"relay_frp_version",
"relay_frps_connections",
"relay_frps_proxy_count",
} {
exists, err := databaseColumnExists(db, "nodes", column)
if err != nil {
return fmt.Errorf("inspect legacy nodes.%s failed: %w", column, err)
}
if exists {
return fmt.Errorf("column nodes.%s should not exist in v16", column)
}
}
if !db.Migrator().HasTable(&WAFIPGroup{}) {
return fmt.Errorf("table waf_ip_groups is missing")
}
@@ -1372,7 +1437,10 @@ func ensureDatabaseSchemaUpToDate(db *gorm.DB, backend string) error {
return err
}
if exists {
return upgradeDatabaseSchema(db, backend, version)
if err := upgradeDatabaseSchema(db, backend, version); err != nil {
return err
}
return dropLegacyNodeColumns(db, backend)
}
empty, err := isDatabaseEmpty(db)
if err != nil {