support different type

pkg/exprhelpers/exprlib.go

@@ -47,7 +47,7 @@ func Init() error {
 	return nil
 }
 
-func FileInit(fileFolder string, filename string) error {
+func FileInit(fileFolder string, filename string, fileType string) error {
 	filepath := path.Join(fileFolder, filename)
 	file, err := os.Open(filepath)
 	if err != nil {
@@ -58,17 +58,18 @@ func FileInit(fileFolder string, filename string) error {
 	if _, ok := dataFile[filename]; !ok {
 		dataFile[filename] = []string{}
 	}
-	fileType := "string"
-	scanner := bufio.NewScanner(file)
-	scanner.Scan()
-	if scanner.Text() == "#type: regex" { // if file contains regexps, it should have this header
-		fileType = "regex"
+	if fileType == "" {
+		fileType = "string"
 	}
+	scanner := bufio.NewScanner(file)
 	for scanner.Scan() {
-		if fileType == "regex" {
+		switch fileType {
+		case "regex":
 			dataFileRegex[filename] = append(dataFileRegex[filename], regexp.MustCompile(scanner.Text()))
-		} else {
+		case "string":
 			dataFile[filename] = append(dataFile[filename], scanner.Text())
+		default:
+			log.Errorf("unknown data type '%s' for : '%s'", fileType, filename)
 		}
 	}
 

pkg/leakybucket/manager.go

@@ -282,7 +282,7 @@ func LoadBucket(g *BucketFactory, dataFolder string) error {
 
 	if len(g.Data) > 0 {
 		for _, data := range g.Data {
-			err = exprhelpers.FileInit(dataFolder, data.DestPath)
+			err = exprhelpers.FileInit(dataFolder, data.DestPath, data.Type)
 			if err != nil {
 				log.Errorf(err.Error())
 			}

pkg/parser/stage.go

@@ -117,7 +117,7 @@ func LoadStages(stageFiles []Stagefile, pctx *UnixParserCtx) ([]Node, error) {
 
 			if len(node.Data) > 0 {
 				for _, data := range node.Data {
-					err = exprhelpers.FileInit(pctx.DataFolder, data.DestPath)
+					err = exprhelpers.FileInit(pctx.DataFolder, data.DestPath, data.Type)
 					if err != nil {
 						log.Errorf(err.Error())
 					}

pkg/types/dataset.go

@@ -13,6 +13,7 @@ import (
 type DataSource struct {
 	SourceURL string `yaml:"source_url"`
 	DestPath  string `yaml:"dest_file"`
+	Type      string `yaml:"type"`
 }
 
 type DataSet struct {