Skip to content

Commit 067660c

Browse files
rafbgarciarafbgarcia
committed
feat: post restore sql
1 parent 8d73363 commit 067660c

File tree

12 files changed

+253
-94
lines changed

12 files changed

+253
-94
lines changed

‎internal/cli/client/client.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,3 +323,50 @@ func (c *Client) UpdateAnonRules(serverIP string, rules []AnonRule) error {
323323

324324
return nil
325325
}
326+
327+
// UpdateConfigRequest represents the config update request
328+
type UpdateConfigRequest struct {
329+
PostRestoreSQL *string `json:"postRestoreSQL,omitempty"`
330+
}
331+
332+
// UpdateConfig updates server configuration (e.g., post-restore SQL)
333+
func (c *Client) UpdateConfig(serverIP string, postRestoreSQL *string) error {
334+
token, err := auth.LoadToken(serverIP)
335+
if err != nil {
336+
return err
337+
}
338+
339+
reqBody := UpdateConfigRequest{
340+
PostRestoreSQL: postRestoreSQL,
341+
}
342+
343+
jsonData, err := json.Marshal(reqBody)
344+
if err != nil {
345+
return fmt.Errorf("failed to marshal request: %w", err)
346+
}
347+
348+
req, err := http.NewRequest(
349+
"PATCH",
350+
fmt.Sprintf("%s/api/config", c.baseURL),
351+
bytes.NewBuffer(jsonData),
352+
)
353+
if err != nil {
354+
return fmt.Errorf("failed to create request: %w", err)
355+
}
356+
357+
req.Header.Set("Content-Type", "application/json")
358+
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
359+
360+
resp, err := c.httpClient.Do(req)
361+
if err != nil {
362+
return fmt.Errorf("failed to send request: %w", err)
363+
}
364+
defer resp.Body.Close()
365+
366+
if resp.StatusCode != http.StatusOK {
367+
body, _ := io.ReadAll(resp.Body)
368+
return fmt.Errorf("failed to update config (status %d): %s", resp.StatusCode, string(body))
369+
}
370+
371+
return nil
372+
}

‎internal/cli/commands/update_anon_rules.go‎

Lines changed: 0 additions & 72 deletions
This file was deleted.

‎internal/cli/commands/update_config.go‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package commands
2+
3+
import (
4+
"fmt"
5+
6+
"github.com/branchd-dev/branchd/internal/cli/client"
7+
"github.com/branchd-dev/branchd/internal/cli/config"
8+
"github.com/spf13/cobra"
9+
)
10+
11+
// NewUpdateConfigCmd creates the update-config command
12+
func NewUpdateConfigCmd() *cobra.Command {
13+
cmd := &cobra.Command{
14+
Use: "update-config",
15+
Short: "Update configuration (anon rules, post-restore SQL) on all servers",
16+
RunE: func(cmd *cobra.Command, args []string) error {
17+
return runUpdateConfig()
18+
},
19+
}
20+
21+
return cmd
22+
}
23+
24+
func runUpdateConfig() error {
25+
// Load config from current directory
26+
cfg, err := config.LoadFromCurrentDir()
27+
if err != nil {
28+
return fmt.Errorf("failed to load config: %w\nRun 'branchd init' to create a configuration file", err)
29+
}
30+
31+
if len(cfg.Servers) == 0 {
32+
return fmt.Errorf("no servers configured. Run 'branchd init' to add a server")
33+
}
34+
35+
// Check if there's anything to update
36+
hasAnonRules := cfg.AnonRules != nil && len(cfg.AnonRules) > 0
37+
hasPostRestoreSQL := cfg.PostRestoreSQL != ""
38+
39+
if !hasAnonRules && !hasPostRestoreSQL {
40+
return fmt.Errorf("no anonRules or postRestoreSQL defined in branchd.json")
41+
}
42+
43+
// Convert config rules to client rules
44+
var rules []client.AnonRule
45+
if hasAnonRules {
46+
for _, rule := range cfg.AnonRules {
47+
rules = append(rules, client.AnonRule{
48+
Table: rule.Table,
49+
Column: rule.Column,
50+
Template: rule.Template,
51+
Type: rule.Type,
52+
})
53+
}
54+
}
55+
56+
// Update all servers
57+
for _, server := range cfg.Servers {
58+
if server.IP == "" {
59+
continue
60+
}
61+
62+
fmt.Printf("Updating configuration on server '%s' (%s)... ", server.Alias, server.IP)
63+
64+
// Create API client
65+
apiClient := client.New(server.IP)
66+
67+
// Update anon rules if defined
68+
if hasAnonRules {
69+
if err := apiClient.UpdateAnonRules(server.IP, rules); err != nil {
70+
fmt.Printf("Failed: %v\n", err)
71+
continue
72+
}
73+
}
74+
75+
// Update post-restore SQL if defined
76+
if hasPostRestoreSQL {
77+
postRestoreSQL := cfg.PostRestoreSQL
78+
if err := apiClient.UpdateConfig(server.IP, &postRestoreSQL); err != nil {
79+
fmt.Printf("Failed: %v\n", err)
80+
continue
81+
}
82+
}
83+
84+
fmt.Printf("Done\n")
85+
}
86+
87+
return nil
88+
}

‎internal/cli/commands/update_server.go‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,22 +35,20 @@ func runUpdateServer() error {
3535
// Update all servers
3636
for _, server := range cfg.Servers {
3737
if server.IP == "" {
38-
fmt.Printf("⚠ Skipping server '%s' (no IP configured)\n", server.Alias)
38+
fmt.Printf("Skipping server '%s' (no IP configured)\n", server.Alias)
3939
continue
4040
}
4141

42-
fmt.Printf("Updating server '%s' (%s)...\n", server.Alias, server.IP)
43-
4442
// Create API client
4543
apiClient := client.New(server.IP)
4644

4745
// Trigger update
4846
if err := apiClient.UpdateServer(server.IP); err != nil {
49-
fmt.Printf("✗ Failed to update server '%s': %v\n", server.Alias, err)
47+
fmt.Printf("Failed to update server '%s': %v\n", server.Alias, err)
5048
continue
5149
}
5250

53-
fmt.Printf("✓ Update triggered for server '%s' (takes ~1 minute)\n", server.Alias)
51+
fmt.Printf("Update triggered on server '%s'\n", server.Alias)
5452
}
5553

5654
return nil

‎internal/cli/config/config.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,8 +141,9 @@ func (r *AnonRule) Parse() (ParsedAnonRule, error) {
141141

142142
// Config represents the CLI configuration file
143143
type Config struct {
144-
Servers []Server `json:"servers"`
145-
AnonRules []AnonRule `json:"anonRules,omitempty"`
144+
Servers []Server `json:"servers"`
145+
AnonRules []AnonRule `json:"anonRules,omitempty"`
146+
PostRestoreSQL string `json:"postRestoreSQL,omitempty"`
146147
}
147148

148149
// DefaultConfig returns a default configuration with example servers

‎internal/cli/root.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ func init() {
4747
rootCmd.AddCommand(commands.NewSelectServerCmd())
4848
rootCmd.AddCommand(commands.NewUpdateCmd(version))
4949
rootCmd.AddCommand(commands.NewUpdateServerCmd())
50-
rootCmd.AddCommand(commands.NewUpdateAnonRulesCmd())
50+
rootCmd.AddCommand(commands.NewUpdateConfigCmd())
5151
}
5252

5353
// Execute runs the root command

‎internal/models/models.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,9 @@ type Config struct {
5656
Domain string `json:"domain"` // Custom domain (e.g. "db.company.com"), empty = use self-signed cert
5757
LetsEncryptEmail string `json:"lets_encrypt_email"` // Email for Let's Encrypt ACME, required if Domain is set
5858

59+
// Post-restore SQL (executed after restore, before anonymization)
60+
PostRestoreSQL string `json:"post_restore_sql" gorm:"type:text"` // SQL statements to run after restore (e.g., TRUNCATE, ANALYZE)
61+
5962
// Computed fields (populated at runtime, not persisted)
6063
DatabaseName string `json:"database_name" gorm:"-"` // Extracted from ConnectionString
6164
}

‎internal/server/config_handlers.go‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -41,20 +41,22 @@ type ConfigResponse struct {
4141
CrunchyBridgeAPIKey string `json:"crunchy_bridge_api_key"`
4242
CrunchyBridgeClusterName string `json:"crunchy_bridge_cluster_name"`
4343
CrunchyBridgeDatabaseName string `json:"crunchy_bridge_database_name"`
44+
PostRestoreSQL string `json:"post_restore_sql"`
4445
}
4546

4647
// UpdateConfigRequest represents the request to update configuration
4748
type UpdateConfigRequest struct {
48-
ConnectionString string `json:"connectionString"`
49-
PostgresVersion string `json:"postgresVersion"`
50-
SchemaOnly *bool `json:"schemaOnly"`
51-
RefreshSchedule string `json:"refreshSchedule"`
52-
Domain string `json:"domain"`
53-
LetsEncryptEmail string `json:"letsEncryptEmail"`
54-
MaxRestores *int `json:"maxRestores"`
55-
CrunchyBridgeAPIKey string `json:"crunchyBridgeApiKey"`
56-
CrunchyBridgeClusterName string `json:"crunchyBridgeClusterName"`
57-
CrunchyBridgeDatabaseName string `json:"crunchyBridgeDatabaseName"`
49+
ConnectionString string `json:"connectionString"`
50+
PostgresVersion string `json:"postgresVersion"`
51+
SchemaOnly *bool `json:"schemaOnly"`
52+
RefreshSchedule string `json:"refreshSchedule"`
53+
Domain string `json:"domain"`
54+
LetsEncryptEmail string `json:"letsEncryptEmail"`
55+
MaxRestores *int `json:"maxRestores"`
56+
CrunchyBridgeAPIKey string `json:"crunchyBridgeApiKey"`
57+
CrunchyBridgeClusterName string `json:"crunchyBridgeClusterName"`
58+
CrunchyBridgeDatabaseName string `json:"crunchyBridgeDatabaseName"`
59+
PostRestoreSQL *string `json:"postRestoreSQL"`
5860
}
5961

6062
// @Summary Get configuration
@@ -94,6 +96,7 @@ func (s *Server) getConfig(c *gin.Context) {
9496
CrunchyBridgeAPIKey: redactSecret(config.CrunchyBridgeAPIKey),
9597
CrunchyBridgeClusterName: config.CrunchyBridgeClusterName,
9698
CrunchyBridgeDatabaseName: config.CrunchyBridgeDatabaseName,
99+
PostRestoreSQL: config.PostRestoreSQL,
97100
})
98101
}
99102

@@ -261,6 +264,11 @@ func (s *Server) updateConfig(c *gin.Context) {
261264
config.Domain = req.Domain
262265
config.LetsEncryptEmail = req.LetsEncryptEmail
263266

267+
// Update post-restore SQL if provided (allow empty string to clear)
268+
if req.PostRestoreSQL != nil {
269+
config.PostRestoreSQL = *req.PostRestoreSQL
270+
}
271+
264272
// If domain is set, configure Caddy with Let's Encrypt
265273
if req.Domain != "" {
266274
if err := s.configureCaddy(req.Domain, req.LetsEncryptEmail); err != nil {
@@ -299,6 +307,7 @@ func (s *Server) updateConfig(c *gin.Context) {
299307
CrunchyBridgeAPIKey: redactSecret(config.CrunchyBridgeAPIKey),
300308
CrunchyBridgeClusterName: config.CrunchyBridgeClusterName,
301309
CrunchyBridgeDatabaseName: config.CrunchyBridgeDatabaseName,
310+
PostRestoreSQL: config.PostRestoreSQL,
302311
})
303312
}
304313

‎internal/workers/crunchy_bridge_restore.sh‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -293,10 +293,7 @@ done
293293
log "Cleaning up schemas and extensions..."
294294

295295
sudo -u postgres ${PG_BIN}/psql -p ${PG_PORT} -h 127.0.0.1 -d ${DATABASE_NAME} << 'EOSQL'
296-
DROP SCHEMA IF EXISTS perfsnap CASCADE;
297-
DROP SCHEMA IF EXISTS hint_plan CASCADE;
298296
DROP EXTENSION IF EXISTS crunchy_pooler CASCADE;
299-
DROP EXTENSION IF EXISTS pg_stat_statements CASCADE;
300297
EOSQL
301298

302299
log "Cleaning up pgBackRest config..."

0 commit comments

Comments
 (0)