Skip to content

Commit bc5c318

Browse files
rafbgarciarafbgarcia
committed
feat: support other anon rule column types
1 parent 93c4c2a commit bc5c318

File tree

8 files changed

+702
-79
lines changed

8 files changed

+702
-79
lines changed

‎internal/anonymize/anonymize.go‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ func generateTableUpdateSQL(table string, rules []models.AnonRule, pkColumn stri
9494
var setClauses []string
9595
var whereConditions []string
9696
for _, rule := range rules {
97-
setValue := renderTemplate(rule.Template)
97+
setValue := renderTemplate(rule.Template, rule.ColumnType)
9898
columnQuoted := quoteIdentifier(rule.Column)
9999

100100
// Add SET clause
@@ -134,7 +134,43 @@ WHERE %s.ctid = numbered_rows.ctid
134134

135135
// renderTemplate converts template string to SQL expression
136136
// Replaces ${index} with row number reference
137-
func renderTemplate(template string) string {
137+
// Handles different column types: text, integer, boolean, null
138+
func renderTemplate(template string, columnType string) string {
139+
// Handle NULL type - ignore template and return SQL NULL
140+
if columnType == "null" {
141+
return "NULL"
142+
}
143+
144+
// Handle boolean type
145+
if columnType == "boolean" {
146+
// Return unquoted true/false
147+
return template
148+
}
149+
150+
// Handle integer type
151+
if columnType == "integer" {
152+
// Check if template contains ${index}
153+
if strings.Contains(template, "${index}") {
154+
// For integer columns with ${index}, we need to cast to text for concatenation,
155+
// then cast back to integer
156+
parts := strings.Split(template, "${index}")
157+
var sqlParts []string
158+
for i, part := range parts {
159+
if part != "" {
160+
sqlParts = append(sqlParts, "'"+part+"'")
161+
}
162+
if i < len(parts)-1 {
163+
sqlParts = append(sqlParts, "numbered_rows._row_num::text")
164+
}
165+
}
166+
// Concatenate and cast to integer
167+
return "(" + strings.Join(sqlParts, " || ") + ")::integer"
168+
}
169+
// No placeholder, return as unquoted integer
170+
return template
171+
}
172+
173+
// Handle text type (default)
138174
// Check if template contains ${index}
139175
if !strings.Contains(template, "${index}") {
140176
// No placeholder, return as quoted string

‎internal/anonymize/anonymize_test.go‎

Lines changed: 64 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -9,64 +9,43 @@ import (
99

1010
func TestGenerateSQL(t *testing.T) {
1111
tests := []struct {
12-
name string
13-
rules []models.AnonRule
14-
primaryKeys map[string]string
15-
want string
12+
name string
13+
rules []models.AnonRule
14+
want string
1615
}{
1716
{
18-
name: "empty rules",
19-
rules: []models.AnonRule{},
20-
primaryKeys: map[string]string{},
21-
want: "",
17+
name: "empty rules",
18+
rules: []models.AnonRule{},
19+
want: "",
2220
},
2321
{
24-
name: "single column anonymization without PK",
22+
name: "single column anonymization",
2523
rules: []models.AnonRule{
26-
{Table: "users", Column: "email", Template: "user_${index}@example.com"},
24+
{Table: "users", Column: "email", Template: "user_${index}@example.com", ColumnType: "text"},
2725
},
28-
primaryKeys: map[string]string{},
29-
want: "UPDATE",
30-
},
31-
{
32-
name: "single column anonymization with PK",
33-
rules: []models.AnonRule{
34-
{Table: "users", Column: "email", Template: "user_${index}@example.com"},
35-
},
36-
primaryKeys: map[string]string{"users": "id"},
37-
want: "ORDER BY \"id\"",
26+
want: "UPDATE",
3827
},
3928
{
4029
name: "multiple columns same table",
4130
rules: []models.AnonRule{
42-
{Table: "users", Column: "email", Template: "user_${index}@example.com"},
43-
{Table: "users", Column: "name", Template: "User ${index}"},
31+
{Table: "users", Column: "email", Template: "user_${index}@example.com", ColumnType: "text"},
32+
{Table: "users", Column: "name", Template: "User ${index}", ColumnType: "text"},
4433
},
45-
primaryKeys: map[string]string{},
46-
want: "numbered_rows._row_num",
34+
want: "numbered_rows._row_num",
4735
},
4836
{
49-
name: "multiple tables with mixed PKs",
37+
name: "multiple tables",
5038
rules: []models.AnonRule{
51-
{Table: "users", Column: "email", Template: "user_${index}@example.com"},
52-
{Table: "orders", Column: "reference", Template: "ORD-${index}"},
39+
{Table: "users", Column: "email", Template: "user_${index}@example.com", ColumnType: "text"},
40+
{Table: "orders", Column: "reference", Template: "ORD-${index}", ColumnType: "text"},
5341
},
54-
primaryKeys: map[string]string{"users": "id"}, // Only users has PK
55-
want: "-- Anonymize table:",
56-
},
57-
{
58-
name: "idempotency - IS DISTINCT FROM clause",
59-
rules: []models.AnonRule{
60-
{Table: "users", Column: "email", Template: "user_${index}@example.com"},
61-
},
62-
primaryKeys: map[string]string{},
63-
want: "IS DISTINCT FROM",
42+
want: "-- Anonymize table:",
6443
},
6544
}
6645

6746
for _, tt := range tests {
6847
t.Run(tt.name, func(t *testing.T) {
69-
got := GenerateSQL(tt.rules, tt.primaryKeys)
48+
got := GenerateSQL(tt.rules, make(map[string]string))
7049
if tt.want != "" && !strings.Contains(got, tt.want) {
7150
t.Errorf("GenerateSQL() output doesn't contain expected string\nwant substring: %v\ngot: %v", tt.want, got)
7251
}
@@ -79,30 +58,64 @@ func TestGenerateSQL(t *testing.T) {
7958

8059
func TestRenderTemplate(t *testing.T) {
8160
tests := []struct {
82-
name string
83-
template string
84-
want string
61+
name string
62+
template string
63+
columnType string
64+
want string
8565
}{
8666
{
87-
name: "simple template with index",
88-
template: "user_${index}@example.com",
89-
want: "'user_' || numbered_rows._row_num || '@example.com'",
67+
name: "simple text template with index",
68+
template: "user_${index}@example.com",
69+
columnType: "text",
70+
want: "'user_' || numbered_rows._row_num || '@example.com'",
71+
},
72+
{
73+
name: "text template with multiple placeholders",
74+
template: "User ${index}",
75+
columnType: "text",
76+
want: "'User ' || numbered_rows._row_num",
77+
},
78+
{
79+
name: "text template without placeholder",
80+
template: "static_value",
81+
columnType: "text",
82+
want: "'static_value'",
83+
},
84+
{
85+
name: "integer template without placeholder",
86+
template: "2222",
87+
columnType: "integer",
88+
want: "2222",
89+
},
90+
{
91+
name: "integer template with index",
92+
template: "${index}",
93+
columnType: "integer",
94+
want: "(numbered_rows._row_num::text)::integer",
95+
},
96+
{
97+
name: "boolean true",
98+
template: "true",
99+
columnType: "boolean",
100+
want: "true",
90101
},
91102
{
92-
name: "template with multiple placeholders",
93-
template: "User ${index}",
94-
want: "'User ' || numbered_rows._row_num",
103+
name: "boolean false",
104+
template: "false",
105+
columnType: "boolean",
106+
want: "false",
95107
},
96108
{
97-
name: "template without placeholder",
98-
template: "static_value",
99-
want: "'static_value'",
109+
name: "null type",
110+
template: "",
111+
columnType: "null",
112+
want: "NULL",
100113
},
101114
}
102115

103116
for _, tt := range tests {
104117
t.Run(tt.name, func(t *testing.T) {
105-
if got := renderTemplate(tt.template); got != tt.want {
118+
if got := renderTemplate(tt.template, tt.columnType); got != tt.want {
106119
t.Errorf("renderTemplate() = %v, want %v", got, tt.want)
107120
}
108121
})

‎internal/cli/client/client.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -271,9 +271,10 @@ func (c *Client) UpdateServer(serverIP string) error {
271271

272272
// AnonRule represents an anonymization rule
273273
type AnonRule struct {
274-
Table string `json:"table"`
275-
Column string `json:"column"`
276-
Template string `json:"template"`
274+
Table string `json:"table"`
275+
Column string `json:"column"`
276+
Template json.RawMessage `json:"template"`
277+
Type string `json:"type,omitempty"` // Optional: "text", "integer", "boolean", "null"
277278
}
278279

279280
// UpdateAnonRulesRequest represents the bulk update request

‎internal/cli/commands/update_anon_rules.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ func runUpdateAnonRules() error {
4444
Table: rule.Table,
4545
Column: rule.Column,
4646
Template: rule.Template,
47+
Type: rule.Type, // Pass through optional type field
4748
})
4849
}
4950

‎internal/cli/config/config.go‎

Lines changed: 120 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,126 @@ type Server struct {
1717

1818
// AnonRule represents an anonymization rule
1919
type AnonRule struct {
20-
Table string `json:"table"`
21-
Column string `json:"column"`
22-
Template string `json:"template"`
20+
Table string `json:"table"`
21+
Column string `json:"column"`
22+
Template json.RawMessage `json:"template"`
23+
Type string `json:"type,omitempty"` // Optional: "text", "integer", "boolean", "null" - overrides auto-detection
24+
}
25+
26+
// ParsedAnonRule represents a parsed anonymization rule with type information
27+
type ParsedAnonRule struct {
28+
Table string
29+
Column string
30+
Template string // String representation of the template value
31+
ColumnType string // "text", "integer", "boolean", "null"
32+
}
33+
34+
// Parse parses the JSON template and returns type information
35+
func (r *AnonRule) Parse() (ParsedAnonRule, error) {
36+
parsed := ParsedAnonRule{
37+
Table: r.Table,
38+
Column: r.Column,
39+
}
40+
41+
// Try to unmarshal as different types to detect the JSON type
42+
if len(r.Template) == 0 {
43+
return parsed, fmt.Errorf("template is empty")
44+
}
45+
46+
// If type is explicitly specified, use it and extract the template value
47+
if r.Type != "" {
48+
// Validate the type
49+
validTypes := map[string]bool{"text": true, "integer": true, "boolean": true, "null": true}
50+
if !validTypes[r.Type] {
51+
return parsed, fmt.Errorf("invalid type '%s', must be one of: text, integer, boolean, null", r.Type)
52+
}
53+
54+
parsed.ColumnType = r.Type
55+
56+
// For null type, template is ignored
57+
if r.Type == "null" {
58+
parsed.Template = ""
59+
return parsed, nil
60+
}
61+
62+
// Extract template value as string
63+
var strVal string
64+
if err := json.Unmarshal(r.Template, &strVal); err == nil {
65+
parsed.Template = strVal
66+
return parsed, nil
67+
}
68+
69+
// If string unmarshal fails, try other JSON types and convert to string
70+
// This handles cases like: template is number but type is "text"
71+
72+
// Try boolean
73+
var boolVal bool
74+
if err := json.Unmarshal(r.Template, &boolVal); err == nil {
75+
if boolVal {
76+
parsed.Template = "true"
77+
} else {
78+
parsed.Template = "false"
79+
}
80+
return parsed, nil
81+
}
82+
83+
// Try number
84+
var numVal float64
85+
if err := json.Unmarshal(r.Template, &numVal); err == nil {
86+
if numVal == float64(int64(numVal)) {
87+
parsed.Template = fmt.Sprintf("%d", int64(numVal))
88+
} else {
89+
parsed.Template = fmt.Sprintf("%f", numVal)
90+
}
91+
return parsed, nil
92+
}
93+
94+
return parsed, fmt.Errorf("failed to parse template with explicit type '%s'", r.Type)
95+
}
96+
97+
// Auto-detect type from JSON
98+
99+
// Check for null
100+
if string(r.Template) == "null" {
101+
parsed.ColumnType = "null"
102+
parsed.Template = ""
103+
return parsed, nil
104+
}
105+
106+
// Try boolean
107+
var boolVal bool
108+
if err := json.Unmarshal(r.Template, &boolVal); err == nil {
109+
parsed.ColumnType = "boolean"
110+
if boolVal {
111+
parsed.Template = "true"
112+
} else {
113+
parsed.Template = "false"
114+
}
115+
return parsed, nil
116+
}
117+
118+
// Try number (integer or float)
119+
var numVal float64
120+
if err := json.Unmarshal(r.Template, &numVal); err == nil {
121+
parsed.ColumnType = "integer"
122+
// Convert to string, handle both int and float
123+
if numVal == float64(int64(numVal)) {
124+
parsed.Template = fmt.Sprintf("%d", int64(numVal))
125+
} else {
126+
parsed.Template = fmt.Sprintf("%f", numVal)
127+
}
128+
return parsed, nil
129+
}
130+
131+
// Try string (must be last, as it's the most permissive)
132+
var strVal string
133+
if err := json.Unmarshal(r.Template, &strVal); err == nil {
134+
parsed.ColumnType = "text"
135+
parsed.Template = strVal
136+
return parsed, nil
137+
}
138+
139+
return parsed, fmt.Errorf("unsupported template type: %s", string(r.Template))
23140
}
24141

25142
// Config represents the CLI configuration file

0 commit comments

Comments
 (0)