DRY up input processing

mpchadwick · Nov 11, 2020 · 1db195b · 1db195b
1 parent 932d557
commit 1db195b
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 123 deletions.
diff --git a/main.go b/main.go
@@ -89,34 +89,33 @@ func main() {
 	// We don't want to hear about it
 	log.SetOutput(ioutil.Discard)
 	reader := bufio.NewReader(os.Stdin)
-
+	
 	args := flag.Args()
+	mode := "anonymize"
 	if len(args) > 0 && args[0] == "map-eav" {
-		eav := dbanon.NewEav(config)
-
-		for {
-			text, err := reader.ReadString('\n')
-			eav.ProcessLine(text)
-
-			if err != nil {
-				break
-			}
-		}
-
-		out, _ := eav.Config.String()
-		fmt.Print(string(out))
-		os.Exit(0)
+		mode = "map-eav"
 	}
 
 	provider := dbanon.NewProvider()
-	processor := dbanon.NewLineProcessor(config, provider)
+	eav := dbanon.NewEav(config)
+	processor := dbanon.NewLineProcessor(mode, config, provider, eav)
+
 
 	for {
 		text, err := reader.ReadString('\n')
-		fmt.Print(processor.ProcessLine(text))
+		result := processor.ProcessLine(text)
+		if mode == "anonymize" {
+			fmt.Print(result)
+		}
 
 		if err != nil {
 			break
 		}
 	}
+
+	if mode == "map-eav" {
+		out, _ := eav.Config.String()
+		fmt.Print(string(out))
+		os.Exit(0)
+	}
 }
diff --git a/src/eav.go b/src/eav.go
@@ -1,10 +1,5 @@
 package dbanon
 
-import (
-	"github.com/blastrain/vitess-sqlparser/sqlparser"
-	"strings"
-)
-
 type Eav struct{
 	Config *Config
 	entityMap map[string]string
@@ -15,53 +10,4 @@ func NewEav(c *Config) *Eav {
 	e := &Eav{Config: c, entityMap: made}
 
 	return e
-}
-
-func (eav Eav) ProcessLine(s string) {
-	// TODO: DRY up duplicated code from LineProcessor.ProcessLine
-	i := strings.Index(s, "INSERT")
-	if i == 0 {
-		eav.processInsert(s)
-		return
-	}
-
-	findNextTable(s)
-}
-
-func (eav Eav) processInsert (s string) {
-	stmt, _ := sqlparser.Parse(s)
-	insert := stmt.(*sqlparser.Insert)
-
-	var entityTypeId string
-	var attributeId string
-
-	rows := insert.Rows.(sqlparser.Values)
-	for _, vt := range rows {
-		for i, e := range vt {
-			column := currentTable[i]
-			switch v := e.(type) {
-			case *sqlparser.SQLVal:
-				if column == "attribute_id" {
-					attributeId = string(v.Val)
-				}
-				if column == "entity_type_id" {
-					entityTypeId = string(v.Val)
-				}
-				if column == "entity_type_code" {
-					eav.entityMap[string(v.Val)] = entityTypeId
-				}
-				if column == "attribute_code" {
-					for _, eavConfig := range eav.Config.Eav {
-						if eav.entityMap[eavConfig.Name] == entityTypeId {
-							for eavK, eavV := range eavConfig.Attributes {
-								if eavK == string(v.Val) {
-									eavConfig.Attributes[attributeId] = eavV
-								}
-							}
-						}
-					}
-				}
-			}
-		}
-	}
 }
diff --git a/src/eav_test.go b/src/eav_test.go
diff --git a/src/processor.go b/src/processor.go
@@ -6,12 +6,14 @@ import (
 )
 
 type LineProcessor struct {
+	Mode string
 	Config *Config
 	Provider ProviderInterface
+	Eav *Eav
 }
 
-func NewLineProcessor(c *Config, p ProviderInterface) *LineProcessor {
-	return &LineProcessor{Config: c, Provider: p}
+func NewLineProcessor(m string, c *Config, p ProviderInterface, e *Eav) *LineProcessor {
+	return &LineProcessor{Mode: m, Config: c, Provider: p, Eav: e}
 }
 
 func (p LineProcessor) ProcessLine(s string) string {
@@ -33,19 +35,22 @@ func (p LineProcessor) processInsert(s string) string {
 
 	processor := p.Config.ProcessTable(table)
 
-	if processor == "" {
+	if processor == "" && p.Mode == "anonymize" {
 		return s
 	}
 
 	var attributeId string
 	var result bool
 	var dataType string
+
+	var entityTypeId string
+
 	rows := insert.Rows.(sqlparser.Values)
 	for _, vt := range rows {
 		for i, e := range vt {
 			column := currentTable[i]
 
-			if processor == "table" {
+			if processor == "table" && p.Mode == "anonymize" {
 				result, dataType = p.Config.ProcessColumn(table, column)
 
 				if !result {
@@ -60,11 +65,32 @@ func (p LineProcessor) processInsert(s string) string {
 				} else {
 					if column == "attribute_id" {
 						attributeId = string(v.Val)
-						result, dataType = p.Config.ProcessEav(table, attributeId)
+						if (p.Mode == "anonymize") {
+							result, dataType = p.Config.ProcessEav(table, attributeId)
+						}
 					}
 					if column == "value" && result {
 						v.Val = []byte(p.Provider.Get(dataType))
 					}
+					if p.Mode == "map-eav" {
+						if column == "entity_type_id" {
+							entityTypeId = string(v.Val)
+						}
+						if column == "entity_type_code" {
+							p.Eav.entityMap[string(v.Val)] = entityTypeId
+						}
+						if column == "attribute_code" {
+							for _, eavConfig := range p.Eav.Config.Eav {
+								if p.Eav.entityMap[eavConfig.Name] == entityTypeId {
+									for eavK, eavV := range eavConfig.Attributes {
+										if eavK == string(v.Val) {
+											eavConfig.Attributes[attributeId] = eavV
+										}
+									}
+								}
+							}
+						}
+					}
 				}
 			}
 		}

diff --git a/src/processor_test.go b/src/processor_test.go
@@ -20,7 +20,9 @@ func (p TestProvider) Get(s string) string {
 func TestProcessLine(t *testing.T) {
 	config, _ := NewConfig("magento2")
 	provider := NewTestProvider()
-	processor := NewLineProcessor(config, provider)
+	mode := "anonymize"
+	eav := NewEav(config)
+	processor := NewLineProcessor(mode, config, provider, eav)
 
 	r1 := processor.ProcessLine("foobar")
 	if r1 != "foobar" {
@@ -78,4 +80,50 @@ func TestProcessLine(t *testing.T) {
 	if !strings.Contains(r4c, "jane") {
 		t.Error("Got no jane wanted jane")
 	}
+}
+
+func TestEavProcessLine(t *testing.T) {
+	config, _ := NewConfig("magento2")
+	provider := NewTestProvider()
+	mode := "map-eav"
+	eav := NewEav(config)
+	processor := NewLineProcessor(mode, config, provider, eav)
+
+
+	processor.ProcessLine("CREATE TABLE `eav_entity_type` (")
+	processor.ProcessLine("  `entity_type_id` smallint(5) unsigned NOT NULL AUTO_INCREMENT COMMENT 'Entity Type ID',")
+	processor.ProcessLine("  `entity_type_code` varchar(50) NOT NULL COMMENT 'Entity Type Code'")
+	processor.ProcessLine(") ENGINE=InnoDB AUTO_INCREMENT=10 DEFAULT CHARSET=utf8 COMMENT='Eav Entity Type';")
+	processor.ProcessLine("/*!40101 SET character_set_client = @saved_cs_client */;")
+	processor.ProcessLine("INSERT INTO `eav_entity_type` (`entity_type_id`, `entity_type_code`) VALUES (1, 'customer');")
+
+	processor.ProcessLine("CREATE TABLE `eav_attribute` (")
+	processor.ProcessLine("  `attribute_id` smallint(5) unsigned NOT NULL AUTO_INCREMENT COMMENT 'Attribute ID',")
+	processor.ProcessLine("  `entity_type_id` smallint(5) unsigned NOT NULL DEFAULT '0' COMMENT 'Entity Type ID',")
+	processor.ProcessLine("  `attribute_code` varchar(255) NOT NULL COMMENT 'Attribute Code'")
+	processor.ProcessLine(") ENGINE=InnoDB AUTO_INCREMENT=180 DEFAULT CHARSET=utf8 COMMENT='Eav Attribute';")
+	processor.ProcessLine("/*!40101 SET character_set_client = @saved_cs_client */;")
+	processor.ProcessLine("INSERT INTO `eav_attribute` (`attribute_id`, `entity_type_id`, `attribute_code`) VALUES (1, 1, 'firstname');")
+	processor.ProcessLine("INSERT INTO `eav_attribute` VALUES (2, 1, 'lastname');")
+	r1 := false
+	r2 := false
+	for _, eavConfig := range eav.Config.Eav {
+		for k, v := range eavConfig.Attributes {
+			if k == "1" && v == "firstname" {
+				r1 = true
+			}
+			if k == "2" && v == "lastname" {
+				r2 = true
+			}
+		}
+	}
+
+	if !r1 {
+		t.Errorf("Got false wanted true")
+	}
+
+	if !r2 {
+		t.Errorf("Got false wanted true")
+	}			
+
 }