Merge pull request #51 from nixys/feat/44

 feat(#44): Delete rows during anonymization (PgSQL)

README.md

@@ -302,6 +302,7 @@ To anonymize a database fields you may use a Go template with the [Sprig templat
 Additional filter functions:
 - `null`: set a field value to `NULL`
 - `isNull`: compare a field value with `NULL`
+- `drop`: drop whole row. If table has filters for several columns and at least one of them returns drop value, whole row will be skipped during the anonymization process
 
 You may also use the following data in a templates:
 - Current table name. Statement: `{{ .TableName }}`
@@ -504,10 +505,10 @@ It's easy.
 ## Roadmap
 
 Following features are already in backlog for our development team and will be released soon:
-- Global variables with the templated values you may use through the filters for all tables and columns
-- Ability to delete tables and rows from faked dump 
-- Ability to output into log a custom messages. It’s quite useful it order to obtain some generated data like admin passwords, etc
-- Support of a big variety of databases
+- [x] Global variables with the templated values you may use through the filters for all tables and columns
+- [x] Ability to delete tables and rows from faked dump 
+- [ ] Ability to output into log a custom messages. It’s quite useful it order to obtain some generated data like admin passwords, etc
+- [ ] Support of a big variety of databases
 
 ## Feedback
 

modules/anonymizers/pgsql/dh.go

@@ -14,18 +14,23 @@ func dhSecurityCopy(usrCtx any, deferred, token []byte) ([]byte, error) {
 	uctx := usrCtx.(*userCtx)
 
 	uctx.security.tmpBuf = append(uctx.security.tmpBuf, token...)
+	uctx.insertIntoBuf = nil
 
 	return deferred, nil
 }
 
-func dhSecurityNil(usrCtx any, deferred, token []byte) ([]byte, error) {
+func dhCopyValuesEnd(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx := usrCtx.(*userCtx)
 
 	if uctx.security.isSkip == true {
 		return []byte{}, nil
 	}
 
+	if uctx.insertIntoBuf != nil {
+		return []byte{}, nil
+	}
+
 	return append(deferred, token...), nil
 }
 
@@ -88,12 +93,12 @@ func dhTableName(usrCtx any, deferred, token []byte) ([]byte, error) {
 
 	uctx.filter.TableCreate(tname)
 
-	d := append(uctx.security.tmpBuf, append(deferred, token...)...)
+	uctx.insertIntoBuf = append(uctx.security.tmpBuf, append(deferred, token...)...)
 
 	uctx.security.isSkip = false
 	uctx.security.tmpBuf = []byte{}
 
-	return d, nil
+	return []byte{}, nil
 }
 
 func dhFieldName(usrCtx any, deferred, token []byte) ([]byte, error) {
@@ -111,7 +116,22 @@ func dhFieldName(usrCtx any, deferred, token []byte) ([]byte, error) {
 		uctx.tables[uctx.filter.TableNameGet()][string(fname)],
 	)
 
-	return append(deferred, token...), nil
+	uctx.insertIntoBuf = append(uctx.insertIntoBuf, append(deferred, token...)...)
+
+	return []byte{}, nil
+}
+
+func dhTableCopyTail(usrCtx any, deferred, token []byte) ([]byte, error) {
+
+	uctx := usrCtx.(*userCtx)
+
+	if uctx.security.isSkip == true {
+		return []byte{}, nil
+	}
+
+	uctx.insertIntoBuf = append(uctx.insertIntoBuf, append(deferred, token...)...)
+
+	return []byte{}, nil
 }
 
 func dhValue(usrCtx any, deferred, token []byte) ([]byte, error) {
@@ -150,14 +170,27 @@ func dhValueEnd(usrCtx any, deferred, token []byte) ([]byte, error) {
 		return []byte{}, err
 	}
 
-	return rowDataGen(uctx.filter), nil
+	b := rowDataGen(uctx.filter)
+	if b == nil {
+		return []byte{}, nil
+	} else {
+		if uctx.insertIntoBuf != nil {
+			b = append(uctx.insertIntoBuf, b...)
+			uctx.insertIntoBuf = nil
+		}
+	}
+
+	return b, nil
 }
 
 func rowDataGen(filter *relfilter.Filter) []byte {
 
 	var out string
 
 	row := filter.ValuePop()
+	if row.Values == nil {
+		return nil
+	}
 
 	for i, v := range row.Values {
 

modules/anonymizers/pgsql/pgsql.go

@@ -37,10 +37,11 @@ type SecurityOpts struct {
 }
 
 type userCtx struct {
-	filter   *relfilter.Filter
-	tn       *string
-	security securityCtx
-	tables   map[string]map[string]string
+	filter        *relfilter.Filter
+	tn            *string
+	security      securityCtx
+	tables        map[string]map[string]string
+	insertIntoBuf []byte
 }
 
 type securityCtx struct {
@@ -201,7 +202,7 @@ func (p *PgSQL) Run(ctx context.Context, w io.Writer) error {
 							Switch: fsm.Switch{
 								Trigger: []byte(";\n"),
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: dhTableCopyTail,
 						},
 					},
 				},
@@ -217,7 +218,7 @@ func (p *PgSQL) Run(ctx context.Context, w io.Writer) error {
 								},
 								Escape: false,
 							},
-							DataHandler: dhSecurityNil,
+							DataHandler: dhCopyValuesEnd,
 						},
 						{
 							Name: stateTableValues,

modules/anonymizers/pgsql/pgsql_test.go

@@ -0,0 +1,82 @@
+package pgsql_anonymize
+
+import (
+	"bytes"
+	"context"
+	"os"
+	"testing"
+
+	"github.com/nixys/nxs-data-anonymizer/misc"
+	"github.com/nixys/nxs-data-anonymizer/modules/filters/relfilter"
+)
+
+func TestPgSQL(t *testing.T) {
+
+	var r, e bytes.Buffer
+
+	fin, err := os.Open("pgsql_test.in.sql")
+	if err != nil {
+		t.Fatal("open input SQL:", err)
+	}
+
+	m, err := Init(
+		fin,
+		InitOpts{
+			Rules: RulesOpts{
+				TableRules: map[string]map[string]relfilter.ColumnRuleOpts{
+
+					// Delete only row with id `2`
+					"public.list_types": {
+						"integer_type": relfilter.ColumnRuleOpts{
+							Type:   misc.ValueTypeTemplate,
+							Value:  "{{ if eq .Values.integer_type \"8765542\" }}{{ drop }}{{ else }}{{ .Values.integer_type }}{{ end }}",
+							Unique: false,
+						},
+					},
+
+					// Delete all rows from table
+					"public.list_types2": {
+						"integer_type": relfilter.ColumnRuleOpts{
+							Type:   misc.ValueTypeTemplate,
+							Value:  "{{ drop }}",
+							Unique: false,
+						},
+					},
+
+					// Delete no rows
+					"public.list_types3": {
+						"integer_type": relfilter.ColumnRuleOpts{
+							Type:   misc.ValueTypeTemplate,
+							Value:  "{{ if eq .Values.integer_type \"0\" }}{{ drop }}{{ else }}{{ .Values.integer_type }}{{ end }}",
+							Unique: false,
+						},
+					},
+				},
+			},
+		},
+	)
+	if err != nil {
+		t.Fatal("init PgSQL:", err)
+	}
+
+	if err := m.Run(context.Background(), &r); err != nil {
+		t.Fatal("run PgSQL:", err)
+	}
+
+	fout, err := os.Open("pgsql_test.out.sql")
+	if err != nil {
+		t.Fatal("open output SQL:", err)
+	}
+
+	if _, err := e.ReadFrom(fout); err != nil {
+		t.Fatal("read output SQL:", err)
+	}
+
+	//os.WriteFile("pgsql_test.out.sql", r.Bytes(), 0644)
+
+	if r.String() != e.String() {
+		t.Fatal("incorrect anonymization result")
+	}
+
+	t.Logf("success")
+}