[优化] 优化脚本

openflare_server/model/main_test.go

@@ -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") {
+	for _, column := range []string{
-		t.Fatal("expected migration to keep legacy nodes.agent_token column")
+		"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() {

openflare_server/model/migrate/migration.go

@@ -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
 }
 

openflare_server/model/migrate/v16.go

@@ -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")
 	}

openflare_server/model/migrations.go

@@ -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 {