Skip to content

Commit 41fea0e

Browse files
rafbgarciarafbgarcia
committed
feat: manually trigger anon rules
1 parent dca48a1 commit 41fea0e

File tree

6 files changed

+311
-70
lines changed

6 files changed

+311
-70
lines changed

‎internal/anonymize/anonymize.go‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
package anonymize
22

33
import (
4+
"context"
45
"fmt"
6+
"os/exec"
57
"strings"
68

79
"github.com/branchd-dev/branchd/internal/models"
10+
"github.com/rs/zerolog"
11+
"gorm.io/gorm"
812
)
913

1014
// GenerateSQL generates anonymization SQL from rules
@@ -96,3 +100,78 @@ func renderTemplate(template string) string {
96100
func quoteIdentifier(name string) string {
97101
return fmt.Sprintf("\"%s\"", strings.ReplaceAll(name, "\"", "\"\""))
98102
}
103+
104+
// ApplyParams contains parameters needed to apply anonymization rules
105+
type ApplyParams struct {
106+
DatabaseName string
107+
PostgresVersion string
108+
PostgresPort int
109+
}
110+
111+
// Apply loads and applies anonymization rules to a database
112+
// Returns the number of rules applied and any error
113+
func Apply(ctx context.Context, db *gorm.DB, params ApplyParams, logger zerolog.Logger) (int, error) {
114+
// Load all anonymization rules
115+
var rules []models.AnonRule
116+
if err := db.Find(&rules).Error; err != nil {
117+
return 0, fmt.Errorf("failed to load anon rules: %w", err)
118+
}
119+
120+
if len(rules) == 0 {
121+
logger.Info().
122+
Str("database_name", params.DatabaseName).
123+
Msg("No anonymization rules configured, skipping")
124+
return 0, nil
125+
}
126+
127+
logger.Info().
128+
Str("database_name", params.DatabaseName).
129+
Int("rule_count", len(rules)).
130+
Msg("Applying anonymization rules")
131+
132+
// Generate SQL from rules
133+
sql := GenerateSQL(rules)
134+
if sql == "" {
135+
logger.Warn().Msg("Generated empty SQL from rules")
136+
return 0, nil
137+
}
138+
139+
// Execute anonymization SQL on the database
140+
script := fmt.Sprintf(`#!/bin/bash
141+
set -euo pipefail
142+
143+
DATABASE_NAME="%s"
144+
PG_VERSION="%s"
145+
PG_PORT="%d"
146+
PG_BIN="/usr/lib/postgresql/${PG_VERSION}/bin"
147+
148+
echo "Applying anonymization rules to database ${DATABASE_NAME}"
149+
150+
# Execute anonymization SQL with correct port
151+
sudo -u postgres ${PG_BIN}/psql -p ${PG_PORT} -d "${DATABASE_NAME}" <<'ANONYMIZE_SQL'
152+
%s
153+
ANONYMIZE_SQL
154+
155+
echo "Anonymization completed successfully"
156+
`, params.DatabaseName, params.PostgresVersion, params.PostgresPort, sql)
157+
158+
cmd := exec.CommandContext(ctx, "bash", "-c", script)
159+
outputBytes, err := cmd.CombinedOutput()
160+
output := string(outputBytes)
161+
if err != nil {
162+
logger.Error().
163+
Err(err).
164+
Str("output", output).
165+
Str("database_name", params.DatabaseName).
166+
Msg("Failed to execute anonymization script")
167+
return 0, fmt.Errorf("anonymization script execution failed: %w", err)
168+
}
169+
170+
logger.Info().
171+
Str("database_name", params.DatabaseName).
172+
Int("rule_count", len(rules)).
173+
Str("output", output).
174+
Msg("Anonymization rules applied successfully")
175+
176+
return len(rules), nil
177+
}

‎internal/server/restore_handlers.go‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/hibiken/asynq"
1313
"gorm.io/gorm"
1414

15+
"github.com/branchd-dev/branchd/internal/anonymize"
1516
"github.com/branchd-dev/branchd/internal/models"
1617
"github.com/branchd-dev/branchd/internal/tasks"
1718
)
@@ -269,3 +270,82 @@ func (s *Server) getRestoreLogs(c *gin.Context) {
269270
"exists": true,
270271
})
271272
}
273+
274+
// postgresVersionToPort maps PostgreSQL major version to its port
275+
func postgresVersionToPort(version string) int {
276+
switch version {
277+
case "14":
278+
return 5414
279+
case "15":
280+
return 5415
281+
case "16":
282+
return 5416
283+
case "17":
284+
return 5417
285+
default:
286+
// Default to 5432 for unknown versions
287+
return 5432
288+
}
289+
}
290+
291+
// @Summary Apply anonymization rules to restore
292+
// @Description Manually trigger anonymization rules on a specific restore
293+
// @Tags restores
294+
// @Produce json
295+
// @Security BearerAuth
296+
// @Param id path string true "Restore ID"
297+
// @Success 200 {object} map[string]interface{}
298+
// @Failure 404 {object} map[string]interface{}
299+
// @Failure 500 {object} map[string]interface{}
300+
// @Router /api/restores/{id}/anonymize [post]
301+
func (s *Server) applyAnonymization(c *gin.Context) {
302+
restoreID := c.Param("id")
303+
304+
// Find restore
305+
var restore models.Restore
306+
if err := s.db.Where("id = ?", restoreID).First(&restore).Error; err != nil {
307+
if err == gorm.ErrRecordNotFound {
308+
c.JSON(http.StatusNotFound, gin.H{"error": "Restore not found"})
309+
return
310+
}
311+
s.logger.Error().Err(err).Str("restore_id", restoreID).Msg("Failed to find restore")
312+
c.JSON(http.StatusInternalServerError, gin.H{"error": "Internal server error"})
313+
return
314+
}
315+
316+
// Load config to get PG version
317+
var config models.Config
318+
if err := s.db.First(&config).Error; err != nil {
319+
s.logger.Error().Err(err).Msg("Failed to load config")
320+
c.JSON(http.StatusInternalServerError, gin.H{"error": "Internal server error"})
321+
return
322+
}
323+
324+
s.logger.Info().
325+
Str("restore_id", restoreID).
326+
Str("restore_name", restore.Name).
327+
Msg("Manually triggering anonymization")
328+
329+
// Apply anonymization rules
330+
pgPort := postgresVersionToPort(config.PostgresVersion)
331+
rulesApplied, err := anonymize.Apply(c.Request.Context(), s.db, anonymize.ApplyParams{
332+
DatabaseName: restore.Name,
333+
PostgresVersion: config.PostgresVersion,
334+
PostgresPort: pgPort,
335+
}, s.logger)
336+
if err != nil {
337+
s.logger.Error().Err(err).Str("restore_id", restoreID).Msg("Failed to apply anonymization")
338+
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Failed to apply anonymization: %v", err)})
339+
return
340+
}
341+
342+
s.logger.Info().
343+
Str("restore_id", restoreID).
344+
Int("rules_applied", rulesApplied).
345+
Msg("Anonymization completed successfully")
346+
347+
c.JSON(http.StatusOK, gin.H{
348+
"message": "Anonymization completed successfully",
349+
"rules_applied": rulesApplied,
350+
})
351+
}

‎internal/server/server.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,7 @@ func (s *Server) setupRouter() {
260260
api.GET("/restores/:id/logs", s.getRestoreLogs)
261261
api.DELETE("/restores/:id", s.deleteRestore)
262262
api.POST("/restores/trigger-restore", s.triggerRestore)
263+
api.POST("/restores/:id/anonymize", s.applyAnonymization)
263264

264265
// Anonymization rules (global)
265266
api.GET("/anon-rules", s.listAnonRules)

‎internal/workers/pg_dump_restore.go‎

Lines changed: 7 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,13 @@ func HandlePgDumpRestoreWaitComplete(ctx context.Context, t *asynq.Task, client
304304
Msg("Restore completed successfully")
305305

306306
// Apply anonymization rules before marking as ready
307-
if err := applyAnonymizationRules(ctx, db, &config, &restoreModel, logger); err != nil {
307+
pgPort := postgresVersionToPort(config.PostgresVersion)
308+
_, err := anonymize.Apply(ctx, db, anonymize.ApplyParams{
309+
DatabaseName: restoreModel.Name,
310+
PostgresVersion: config.PostgresVersion,
311+
PostgresPort: pgPort,
312+
}, logger)
313+
if err != nil {
308314
logger.Error().Err(err).Msg("Failed to apply anonymization rules")
309315
return fmt.Errorf("failed to apply anonymization rules: %w", err)
310316
}
@@ -397,70 +403,3 @@ func calculateNextRefresh(cronExpr string, from time.Time) *time.Time {
397403
return &next
398404
}
399405

400-
// applyAnonymizationRules applies configured anonymization rules to a restored database
401-
func applyAnonymizationRules(ctx context.Context, db *gorm.DB, config *models.Config, database *models.Restore, logger zerolog.Logger) error {
402-
// Load global anonymization rules
403-
var rules []models.AnonRule
404-
if err := db.Find(&rules).Error; err != nil {
405-
return fmt.Errorf("failed to load anon rules: %w", err)
406-
}
407-
408-
if len(rules) == 0 {
409-
logger.Info().
410-
Str("restore_id", database.ID).
411-
Msg("No anonymization rules configured, skipping")
412-
return nil
413-
}
414-
415-
logger.Info().
416-
Str("restore_id", database.ID).
417-
Int("rule_count", len(rules)).
418-
Msg("Applying anonymization rules")
419-
420-
// Generate SQL from rules
421-
sql := anonymize.GenerateSQL(rules)
422-
if sql == "" {
423-
logger.Warn().Msg("Generated empty SQL from rules")
424-
return nil
425-
}
426-
427-
// Execute anonymization SQL on the database
428-
pgPort := postgresVersionToPort(config.PostgresVersion)
429-
script := fmt.Sprintf(`#!/bin/bash
430-
set -euo pipefail
431-
432-
DATABASE_NAME="%s"
433-
PG_VERSION="%s"
434-
PG_PORT="%d"
435-
PG_BIN="/usr/lib/postgresql/${PG_VERSION}/bin"
436-
437-
echo "Applying anonymization rules to database ${DATABASE_NAME}"
438-
439-
# Execute anonymization SQL with correct port
440-
sudo -u postgres ${PG_BIN}/psql -p ${PG_PORT} -d "${DATABASE_NAME}" <<'ANONYMIZE_SQL'
441-
%s
442-
ANONYMIZE_SQL
443-
444-
echo "Anonymization completed successfully"
445-
`, database.Name, config.PostgresVersion, pgPort, sql)
446-
447-
cmd := exec.CommandContext(ctx, "bash", "-c", script)
448-
outputBytes, err := cmd.CombinedOutput()
449-
output := string(outputBytes)
450-
if err != nil {
451-
logger.Error().
452-
Err(err).
453-
Str("output", output).
454-
Str("database_name", database.Name).
455-
Msg("Failed to execute anonymization script")
456-
return fmt.Errorf("anonymization script execution failed: %w", err)
457-
}
458-
459-
logger.Info().
460-
Str("database_name", database.Name).
461-
Int("rule_count", len(rules)).
462-
Str("output", output).
463-
Msg("Anonymization rules applied successfully")
464-
465-
return nil
466-
}

‎web/src/lib/openapi.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -566,6 +566,15 @@ export class Api<
566566
...params,
567567
}),
568568

569+
restoresAnonymizeCreate: (id: string, params: RequestParams = {}) =>
570+
this.request<Record<string, any>, Record<string, any>>({
571+
path: `/api/restores/${id}/anonymize`,
572+
method: "POST",
573+
secure: true,
574+
format: "json",
575+
...params,
576+
}),
577+
569578
restoresLogsList: (
570579
id: string,
571580
query?: {

0 commit comments

Comments
 (0)