Add CSV source.

Author: clems4ever

cmd/go-graphkb/main.go

@@ -27,21 +27,6 @@ var ConfigPath string
 
 func init() {
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
-
-	Sources = []sources.Source{}
-
-	for _, s := range Sources {
-		sources.Registry.Add(s)
-		g := s.Graph()
-
-		for _, a := range g.Assets() {
-			knowledge.SchemaRegistrySingleton.AddAssetType(a)
-		}
-
-		for _, r := range g.Relations() {
-			knowledge.SchemaRegistrySingleton.AddRelationType(r.Type)
-		}
-	}
 }
 
 func main() {
@@ -117,6 +102,27 @@ func onInit() {
 		viper.GetString("mariadb.password"),
 		viper.GetString("mariadb.host"),
 		dbName)
+
+	Sources = []sources.Source{
+		sources.NewCSVSource(),
+	}
+
+	for _, s := range Sources {
+		sources.Registry.Add(s)
+		g, err := s.Graph()
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		for _, a := range g.Assets() {
+			knowledge.SchemaRegistrySingleton.AddAssetType(a)
+		}
+
+		for _, r := range g.Relations() {
+			knowledge.SchemaRegistrySingleton.AddRelationType(r.Type)
+		}
+	}
 }
 
 func start(cmd *cobra.Command, args []string) {
@@ -179,6 +185,7 @@ func flush(cmd *cobra.Command, args []string) {
 	if err := Database.FlushAll(); err != nil {
 		log.Fatal(err)
 	}
+	fmt.Println("Successul flush")
 }
 
 func getSource(cmd *cobra.Command, args []string) {
@@ -194,7 +201,11 @@ func getSource(cmd *cobra.Command, args []string) {
 	assets := make(map[string]bool)
 	relations := make(map[string]bool)
 	for _, s := range selectedSources {
-		g := s.Graph()
+		g, err := s.Graph()
+		if err != nil {
+			log.Fatal(err)
+		}
+
 		for _, a := range g.Assets() {
 			assets[string(a)] = true
 		}

config.yml

@@ -1,7 +1,12 @@
 port: 8090
 
 mariadb:
-  username: root
-  password: example
+  username: graphkb
+  password: password
   host:
   database: graphkb
+
+sources:
+  csv:
+    schema: ./examples/example-schema.csv
+    data: ./examples/example-data.csv

docker-compose.yml

@@ -10,6 +10,9 @@ services:
       - 3306:3306
     environment:
       MYSQL_ROOT_PASSWORD: example
+      MYSQL_DATABASE: graphkb
+      MYSQL_USER: graphkb
+      MYSQL_PASSWORD: password
     volumes:
       - /var/lib/go-graphkb:/var/lib/mysql
 

examples/example-data.csv

@@ -0,0 +1,17 @@
+from_type,from_value,relation_type,to_type,to_value
+employee,bob,is_in,team,Team-a
+employee,charlie,is_in,team,Team-a
+employee,alice,is_in,team,Team-b
+employee,harry,is_in,team,Team-b
+employee,john,is_in,team,Team-c
+employee,james,is_in,team,Team-b
+employee,john,work_at,office,P1
+employee,alice,work_at,office,P1
+employee,bob,work_at,office,B2
+employee,charlie,work_at,office,B2
+employee,james,work_at,office,B2
+employee,harry,work_at,office,B2
+office,P1,is_in,country,Paris
+office,B2,is_in,country,Barcelona
+employee,harry,is_manager_of,employee,alice
+employee,harry,is_manager_of,employee,james

examples/example-schema.csv

@@ -0,0 +1,5 @@
+from_type,relation_type,to_type
+employee,is_in,team
+employee,is_manager_of,employee
+employee,work_at,office
+office,is_in,country

internal/knowledge/fieldbinding.go

@@ -49,7 +49,7 @@ func NewGraph() *Graph {
 // AddAsset add an asset to the graph
 func (g *Graph) AddAsset(assetType AssetType, assetKey string) AssetKey {
 	if !SchemaRegistrySingleton.AssetExists(assetType) {
-		log.Fatal(fmt.Errorf("Asset type '%s' does not exist to add asset %s", assetType, assetKey))
+		log.Fatal(fmt.Errorf("Asset type '%s' does not exist in schema registry. Cannot add asset value '%s'", assetType, assetKey))
 	}
 	asset := Asset{Type: assetType, Key: assetKey}
 	g.assets[asset] = true
@@ -59,7 +59,7 @@ func (g *Graph) AddAsset(assetType AssetType, assetKey string) AssetKey {
 // AddRelation add a relation to the graph
 func (g *Graph) AddRelation(from AssetKey, relationType RelationKeyType, to AssetKey) Relation {
 	if !SchemaRegistrySingleton.RelationExists(relationType) {
-		log.Fatal(fmt.Errorf("Relation type '%s' does not exist", relationType))
+		log.Fatal(fmt.Errorf("Relation type '%s' does not exist in schema registry", relationType))
 	}
 	relation := Relation{
 		Type: relationType,

internal/knowledge/graph.go

@@ -56,7 +56,12 @@ func getSourceGraph(w http.ResponseWriter, r *http.Request) {
 	for _, sname := range sourceNames {
 		if sources.Registry.Exist(sname) {
 			s := sources.Registry.Get(sname)
-			sg.Merge(s.Graph())
+			g, err := s.Graph()
+			if err != nil {
+				w.WriteHeader(http.StatusInternalServerError)
+				return
+			}
+			sg.Merge(*g)
 		}
 	}
 	replyWithSourceGraph(w, &sg)

internal/server/server.go

@@ -0,0 +1,111 @@
+package sources
+
+import (
+	"encoding/csv"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/clems4ever/go-graphkb/internal/knowledge"
+	"github.com/spf13/viper"
+)
+
+type CSVSource struct {
+	schemaPath string
+	dataPath   string
+}
+
+func NewCSVSource() *CSVSource {
+	csvSource := new(CSVSource)
+	csvSource.schemaPath = viper.GetString("sources.csv.schema")
+	csvSource.dataPath = viper.GetString("sources.csv.data")
+	return csvSource
+}
+
+func (cs *CSVSource) Name() string {
+	return "csv"
+}
+
+func (cs *CSVSource) Graph() (*knowledge.SchemaGraph, error) {
+	file, err := os.Open(cs.schemaPath)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	r := csv.NewReader(file)
+
+	graph := knowledge.NewSchemaGraph()
+
+	header := true
+
+	for {
+		record, err := r.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, fmt.Errorf("Unable to read schema in CSV file: %v", err)
+		}
+
+		// Skip header line
+		if header {
+			header = false
+			continue
+		}
+
+		fromType := graph.AddAsset(record[0])
+		toType := graph.AddAsset(record[2])
+		graph.AddRelation(fromType, record[1], toType)
+	}
+	return &graph, nil
+}
+
+func (cs *CSVSource) Start(emitter *knowledge.GraphEmitter) error {
+	file, err := os.Open(cs.dataPath)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	r := csv.NewReader(file)
+
+	previousGraph, err := emitter.Read()
+	if err != nil {
+		return err
+	}
+
+	tx := emitter.CreateCompleteGraphTransaction(previousGraph)
+
+	header := true
+
+	for {
+		record, err := r.Read()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return fmt.Errorf("Unable to read schema in CSV file: %v", err)
+		}
+
+		// Skip header line
+		if header {
+			header = false
+			continue
+		}
+
+		relationType := knowledge.RelationType{
+			FromType: knowledge.AssetType(record[0]),
+			ToType:   knowledge.AssetType(record[3]),
+			Type:     knowledge.RelationKeyType(record[2]),
+		}
+
+		tx.Relate(record[1], relationType, record[4])
+	}
+	tx.Commit()
+	return nil
+}
+
+func (cs *CSVSource) Stop() error {
+	return nil
+}

internal/sources/csv_source.go

@@ -18,8 +18,11 @@ func (os *ObservableSource) Name() string {
 	return os.source.Name()
 }
 
-func (os *ObservableSource) Graph() knowledge.SchemaGraph {
-	sg := os.source.Graph()
+func (os *ObservableSource) Graph() (*knowledge.SchemaGraph, error) {
+	sg, err := os.source.Graph()
+	if err != nil {
+		return nil, err
+	}
 
 	source := sg.AddAsset("source")
 	for _, a := range sg.Assets() {
@@ -28,7 +31,7 @@ func (os *ObservableSource) Graph() knowledge.SchemaGraph {
 		}
 		sg.AddRelation(source, "observed", a)
 	}
-	return sg
+	return sg, nil
 }
 
 func (os *ObservableSource) Start(e *knowledge.GraphEmitter) error {