Add 'observed' relation between source and assets.

Author: clems4ever

internal/knowledge/graph_updater.go

@@ -25,24 +25,66 @@ func NewGraphUpdater(graphDB GraphDB, schemaPersistor schema.Persistor) *GraphUp
 	return &GraphUpdater{graphDB, schemaPersistor}
 }
 
-func (sl *GraphUpdater) doUpdate(updates SourceSubGraphUpdates) error {
-	previousSchema, err := sl.schemaPersistor.LoadSchema(context.Background(), updates.Source)
+// Augment the graph of the user with "observed" relation from the source to the each asset
+func (sl *GraphUpdater) appendObservedRelations(source string, updates *GraphUpdatesBulk) {
+	assetsToAdd := []Asset{Asset{Type: "source", Key: source}}
+	observedRelationsToRemove := []Relation{}
+	observedRelationsToAdd := []Relation{}
+
+	for _, a := range updates.AssetUpserts {
+		observedRelationsToAdd = append(observedRelationsToAdd, Relation{
+			Type: "observed",
+			From: AssetKey(assetsToAdd[0]),
+			To:   AssetKey(a),
+		})
+	}
+
+	for _, a := range updates.AssetRemovals {
+		observedRelationsToRemove = append(observedRelationsToRemove, Relation{
+			Type: "observed",
+			From: AssetKey(assetsToAdd[0]),
+			To:   AssetKey(a),
+		})
+	}
+
+	updates.AssetUpserts = append(updates.AssetUpserts, assetsToAdd...)
+	updates.RelationUpserts = append(updates.RelationUpserts, observedRelationsToAdd...)
+	updates.RelationRemovals = append(updates.RelationRemovals, observedRelationsToRemove...)
+}
+
+func (sl *GraphUpdater) updateSchema(source string, sg *schema.SchemaGraph) error {
+	for _, a := range sg.Assets() {
+		sg.AddRelation(schema.AssetType("source"), "observed", a)
+	}
+	sg.AddAsset("source")
+
+	previousSchema, err := sl.schemaPersistor.LoadSchema(context.Background(), source)
 	if err != nil {
 		fmt.Printf("[ERROR] Unable to read schema from DB: %v.\n", err)
 		fmt.Println("[WARNING] The graph has not been updated.")
 		return err
 	}
 
-	schemaEqual := previousSchema.Equal(updates.Schema)
+	schemaEqual := previousSchema.Equal(*sg)
 
 	if !schemaEqual {
 		fmt.Println("The schema needs an update")
-		if err := sl.schemaPersistor.SaveSchema(context.Background(), updates.Source, updates.Schema); err != nil {
+		if err := sl.schemaPersistor.SaveSchema(context.Background(), source, *sg); err != nil {
 			fmt.Printf("[ERROR] Unable to write schema in DB: %v.\n", err)
 			fmt.Println("[WARNING] The graph has not been updated.")
 			return err
 		}
 	}
+	return nil
+}
+
+func (sl *GraphUpdater) doUpdate(updates SourceSubGraphUpdates) error {
+	if err := sl.updateSchema(updates.Source, &updates.Schema); err != nil {
+		return err
+	}
+
+	sl.appendObservedRelations(updates.Source, &updates.Updates)
+
 	if err := sl.graphDB.UpdateGraph(updates.Source, &updates.Updates); err != nil {
 		fmt.Printf("[ERROR] Unable to write schema in graph DB: %v\n", err)
 		return err

internal/server/server.go

@@ -311,6 +311,20 @@ func postGraphUpdates(graphUpdatesC chan knowledge.SourceSubGraphUpdates) http.H
 	}
 }
 
+func flushDatabase(graphDB knowledge.GraphDB) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if err := graphDB.FlushAll(); err != nil {
+			replyWithInternalError(w, err)
+			return
+		}
+
+		if err := graphDB.InitializeSchema(); err != nil {
+			replyWithInternalError(w, err)
+			return
+		}
+	}
+}
+
 // Secret is the secret provider function for basic auth
 func Secret(user, realm string) string {
 	if user == "admin" {
@@ -335,6 +349,7 @@ func StartServer(database knowledge.GraphDB, schemaPersistor schema.Persistor,
 	getSourceGraphHandler := getSourceGraph(schemaPersistor)
 	getDatabaseDetailsHandler := getDatabaseDetails(database)
 	postQueryHandler := postQuery(database)
+	flushDatabaseHandler := flushDatabase(database)
 
 	if viper.GetString("password") != "" {
 		authenticator := auth.NewBasicAuthenticator("example.com", Secret)
@@ -349,12 +364,15 @@ func StartServer(database knowledge.GraphDB, schemaPersistor schema.Persistor,
 		getSourceGraphHandler = AuthMiddleware(getSourceGraphHandler)
 		getDatabaseDetailsHandler = AuthMiddleware(getDatabaseDetailsHandler)
 		postQueryHandler = AuthMiddleware(postQueryHandler)
+		flushDatabaseHandler = AuthMiddleware(flushDatabaseHandler)
 	}
 
 	r.HandleFunc("/api/sources", listSourcesHandler).Methods("GET")
 	r.HandleFunc("/api/schema", getSourceGraphHandler).Methods("GET")
 	r.HandleFunc("/api/database", getDatabaseDetailsHandler).Methods("GET")
 
+	r.HandleFunc("/api/admin/flush", flushDatabaseHandler).Methods("POST")
+
 	r.HandleFunc("/api/graph/read", getGraphRead(database)).Methods("GET")
 	r.HandleFunc("/api/graph/update", postGraphUpdates(graphUpdatesC)).Methods("POST")