feat: add custom dictionaries for additional, customized flow hydration

2025-12-12 06:24:10 +01:00 · 2023-08-08 08:55:25 +02:00
parent c1b2008ee9
commit e6effd1335
16 changed files with 764 additions and 3 deletions
--- a/common/schema/root_test.go
+++ b/common/schema/root_test.go
@@ -6,6 +6,7 @@ package schema_test
 import (
 	"testing"

+	"akvorado/common/helpers"
 	"akvorado/common/schema"
 )

@@ -61,3 +62,180 @@ func TestDisableForbiddenColumns(t *testing.T) {
 		t.Fatalf("New() error:\n%+v", err)
 	}
 }
+
+func TestCustomDictionaries(t *testing.T) {
+	config := schema.DefaultConfiguration()
+	config.CustomDictionaries = make(map[string]schema.CustomDict)
+	config.CustomDictionaries["test"] = schema.CustomDict{
+		Keys: []schema.CustomDictKey{
+			{Name: "SrcAddr", Type: "string"},
+		},
+		Attributes: []schema.CustomDictAttribute{
+			{Name: "csv_col_name", Type: "string", Label: "DimensionAttribute"},
+			{Name: "role", Type: "string"},
+		},
+		Source:     "test.csv",
+		Dimensions: []string{"SrcAddr", "DstAddr"},
+	}
+
+	s, err := schema.New(config)
+	if err != nil {
+		t.Fatalf("New() error:\n%+v", err)
+	}
+
+	// Test if SrcAddrAttribute and DstAddrAttribute are in s.columns
+	srcFound := false
+	dstFound := false
+	srcRoleFound := false
+	dstRoleFound := false
+
+	// check if srcAddrAttribute and dstAddrAttribute are in s.columns, and have the correct type/generatefrom
+	for _, column := range s.Columns() {
+		if column.Name == "SrcAddrDimensionAttribute" {
+			srcFound = true
+			if column.ClickHouseType != "LowCardinality(string)" {
+				t.Fatalf("SrcAddrDimensionAttribute should be LowCardinality(string), is %s", column.ClickHouseType)
+			}
+			if column.ClickHouseGenerateFrom != "dictGet('custom_dict_test', 'csv_col_name', SrcAddr)" {
+				t.Fatalf("SrcAddrDimensionAttribute should be generated from `dictGet('custom_dict_test', 'csv_col_name', SrcAddr)`, is %s", column.ClickHouseGenerateFrom)
+			}
+		}
+		if column.Name == "DstAddrDimensionAttribute" {
+			dstFound = true
+			if column.ClickHouseType != "LowCardinality(string)" {
+				t.Fatalf("DstAddrDimensionAttribute should be LowCardinality(string), is %s", column.ClickHouseType)
+			}
+			if column.ClickHouseGenerateFrom != "dictGet('custom_dict_test', 'csv_col_name', DstAddr)" {
+				t.Fatalf("DstAddrDimensionAttribute should be generated from `dictGet('custom_dict_test', 'csv_col_name', DstAddr)`, is %s", column.ClickHouseGenerateFrom)
+			}
+		}
+		// this part only tests default dimension name generation
+		if column.Name == "SrcAddrRole" {
+			srcRoleFound = true
+		}
+		if column.Name == "DstAddrRole" {
+			dstRoleFound = true
+		}
+
+	}
+
+	if !srcFound {
+		t.Fatal("SrcAddrDimensionAttribute not found")
+	}
+	if !dstFound {
+		t.Fatal("DstAddrDimensionAttribute not found")
+	}
+	if !srcRoleFound {
+		t.Fatal("SrcAddrRole not found")
+	}
+	if !dstRoleFound {
+		t.Fatal("DstAddrRole not found")
+	}
+
+}
+
+func TestCustomDictionariesMatcher(t *testing.T) {
+	config := schema.DefaultConfiguration()
+	config.CustomDictionaries = make(map[string]schema.CustomDict)
+	config.CustomDictionaries["test"] = schema.CustomDict{
+		Keys: []schema.CustomDictKey{
+			{Name: "exporter", Type: "string", MatchDimension: "ExporterAddress"},
+			{Name: "interface", Type: "string", MatchDimensionSuffix: "Name"},
+		},
+		Attributes: []schema.CustomDictAttribute{
+			{Name: "csv_col_name", Type: "string", Label: "DimensionAttribute"},
+		},
+		Source:     "test.csv",
+		Dimensions: []string{"OutIf", "InIf"},
+		Layout:     "complex_key_hashed",
+	}
+
+	s, err := schema.New(config)
+	if err != nil {
+		t.Fatalf("New() error:\n%+v", err)
+	}
+
+	// Test if SrcAddrAttribute and DstAddrAttribute are in s.columns
+	outFound := false
+	inFound := false
+
+	// check if srcAddrAttribute and dstAddrAttribute are in s.columns, and have the correct type/generatefrom
+	for _, column := range s.Columns() {
+		if column.Name == "OutIfDimensionAttribute" {
+			outFound = true
+			if column.ClickHouseType != "LowCardinality(string)" {
+				t.Fatalf("OutIfDimensionAttribute should be LowCardinality(string), is %s", column.ClickHouseType)
+			}
+			if column.ClickHouseGenerateFrom != "dictGet('custom_dict_test', 'csv_col_name', (ExporterAddress,OutIfName))" {
+				t.Fatalf("OutIfDimensionAttribute should be generated from `dictGet('custom_dict_test', 'csv_col_name', (ExporterAddress,OutIfName))`, is %s", column.ClickHouseGenerateFrom)
+			}
+		}
+		if column.Name == "InIfDimensionAttribute" {
+			inFound = true
+			if column.ClickHouseType != "LowCardinality(string)" {
+				t.Fatalf("InIfDimensionAttribute should be LowCardinality(string), is %s", column.ClickHouseType)
+			}
+			if column.ClickHouseGenerateFrom != "dictGet('custom_dict_test', 'csv_col_name', (ExporterAddress,InIfName))" {
+				t.Fatalf("InIfDimensionAttribute should be generated from `dictGet('custom_dict_test', 'csv_col_name', (ExporterAddress,InIfName)), is %s", column.ClickHouseGenerateFrom)
+			}
+		}
+	}
+
+	if !outFound {
+		t.Fatal("OutIfDimensionAttribute not found")
+	}
+	if !inFound {
+		t.Fatal("InIfDimensionAttribute not found")
+	}
+}
+
+// we need MatchDimension or MatchDimensionSuffix for multiple keys
+func TestCustomDictMultiKeyErr(t *testing.T) {
+	config := schema.DefaultConfiguration()
+	config.CustomDictionaries = make(map[string]schema.CustomDict)
+	config.CustomDictionaries["test"] = schema.CustomDict{
+		Keys: []schema.CustomDictKey{
+			{Name: "exporter", Type: "string"},
+			{Name: "interface", Type: "string"},
+		},
+		Attributes: []schema.CustomDictAttribute{
+			{Name: "csv_col_name", Type: "string", Label: "DimensionAttribute"},
+		},
+		Source:     "test.csv",
+		Dimensions: []string{"OutIf", "InIf"},
+		Layout:     "complex_key_hashed",
+	}
+
+	_, err := schema.New(config)
+	if err == nil {
+		t.Fatal("New() did not error")
+	}
+
+	if diff := helpers.Diff(err.Error(), "custom dictionary test has more than one key, but key exporter has neither MatchDimension nor MatchDimensionSuffix set"); diff != "" {
+		t.Fatalf("New() did not error correctly\n %s", diff)
+	}
+}
+
+// a dict without key makes no sense, catch this
+func TestCustomDictNoKeyErr(t *testing.T) {
+	config := schema.DefaultConfiguration()
+	config.CustomDictionaries = make(map[string]schema.CustomDict)
+	config.CustomDictionaries["test"] = schema.CustomDict{
+		Keys: []schema.CustomDictKey{},
+		Attributes: []schema.CustomDictAttribute{
+			{Name: "csv_col_name", Type: "string", Label: "DimensionAttribute"},
+		},
+		Source:     "test.csv",
+		Dimensions: []string{"OutIf", "InIf"},
+		Layout:     "complex_key_hashed",
+	}
+
+	_, err := schema.New(config)
+	if err == nil {
+		t.Fatal("New() did not error")
+	}
+
+	if diff := helpers.Diff(err.Error(), "custom dictionary test has no keys, this is not supported"); diff != "" {
+		t.Fatalf("New() did not error correctly\n %s", diff)
+	}
+}