Add streaming support on read side.

Streaming was implemented only on the write side but when multiple data
sources requested a graph from the API, huge graph could have been
loaded in memory before being encoded and sent over the wire. With this
encoder, each relation or asset is itself encoded and sent independently
and the whole graph is reconstructed on client side.

The encoding format is definitely not optimized yet since it uses json
format and duplicated assets could be sent multiple times but that is
worth it considering the amount of memory saved. The deduplication is
done on the client side anyway.

Author: clems4ever

cmd/go-graphkb/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"strings"
@@ -149,11 +150,15 @@ func listen(cmd *cobra.Command, args []string) {
 
 func read(cmd *cobra.Command, args []string) {
 	g := knowledge.NewGraph()
-	err := Database.ReadGraph(context.Background(), args[0], g)
+	buff := bytes.NewBuffer(nil)
+
+	encoder := knowledge.NewGraphEncoder(buff)
+	err := Database.ReadGraph(context.Background(), args[0], encoder)
 	if err != nil {
 		logrus.Fatal(err)
 	}
 
+	knowledge.NewGraphDecoder(buff).Decode(g)
 	fmt.Printf("assets = %d\nrelations = %d\n", len(g.Assets()), len(g.Relations()))
 }
 

cmd/test/main.go

@@ -0,0 +1,40 @@
+package main
+
+import (
+	"database/sql"
+	"fmt"
+
+	_ "github.com/go-sql-driver/mysql"
+	"github.com/sirupsen/logrus"
+)
+
+var (
+	username     = "svc-criteokb"
+	password     = "@1wQ.WBE93~r1D~b5ia~"
+	host         = "security-exp01-pa4.central.criteo.preprod"
+	databaseName = "criteokb"
+)
+
+func main() {
+	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@(%s)/%s", username, password,
+		host, databaseName))
+	if err != nil {
+		logrus.Fatal(err)
+	}
+
+	tx, err := db.Begin()
+	if err != nil {
+		logrus.Fatal(err)
+	}
+
+	_, err = tx.Exec("INSERT INTO test (value) VALUES (?)", uint64(12435887700123278845))
+	if err != nil {
+		tx.Rollback()
+		logrus.Fatal(err)
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		logrus.Fatal(err)
+	}
+}

internal/client/graph_client.go

@@ -68,13 +68,9 @@ func (gc *GraphClient) ReadCurrentGraph() (*knowledge.Graph, error) {
 		return nil, handleUnexpectedResponse(res)
 	}
 
-	b, err := ioutil.ReadAll(res.Body)
-	if err != nil {
-		return nil, err
-	}
-
+	graphDecoder := knowledge.NewGraphDecoder(res.Body)
 	graph := knowledge.NewGraph()
-	err = json.Unmarshal(b, graph)
+	err = graphDecoder.Decode(graph)
 	if err != nil {
 		return nil, err
 	}

internal/database/mariadb.go

@@ -406,7 +406,7 @@ func (m *MariaDB) RemoveRelations(ctx context.Context, source string, relations
 }
 
 // ReadGraph read source subgraph
-func (m *MariaDB) ReadGraph(ctx context.Context, sourceName string, graph *knowledge.Graph) error {
+func (m *MariaDB) ReadGraph(ctx context.Context, sourceName string, encoder *knowledge.GraphEncoder) error {
 	logrus.Debugf("Start reading graph of data source with name %s", sourceName)
 	sourceID, err := m.resolveSourceID(sourceName)
 	if err != nil {
@@ -437,18 +437,25 @@ func (m *MariaDB) ReadGraph(ctx context.Context, sourceName string, graph *knowl
 				if err := rows.Scan(&FromType, &FromKey, &ToType, &ToKey, &Type); err != nil {
 					return err
 				}
-				fromAsset := knowledge.Asset{
+				fromAsset := knowledge.AssetKey{
 					Type: schema.AssetType(FromType),
 					Key:  FromKey,
 				}
-				toAsset := knowledge.Asset{
+				toAsset := knowledge.AssetKey{
 					Type: schema.AssetType(ToType),
 					Key:  ToKey,
 				}
-				graph.AddRelation(
-					graph.AddAsset(fromAsset.Type, fromAsset.Key),
-					schema.RelationKeyType(Type),
-					graph.AddAsset(toAsset.Type, toAsset.Key))
+
+				relation := knowledge.Relation{
+					Type: schema.RelationKeyType(Type),
+					From: fromAsset,
+					To:   toAsset,
+				}
+
+				err = encoder.EncodeRelation(relation)
+				if err != nil {
+					return fmt.Errorf("Unable to write relation %v: %v", relation, err)
+				}
 			}
 		}
 
@@ -472,7 +479,16 @@ func (m *MariaDB) ReadGraph(ctx context.Context, sourceName string, graph *knowl
 				if err := rows.Scan(&Type, &Key); err != nil {
 					return fmt.Errorf("Unable to read standalone asset: %v", err)
 				}
-				graph.AddAsset(schema.AssetType(Type), Key)
+
+				asset := knowledge.Asset{
+					Type: schema.AssetType(Type),
+					Key:  Key,
+				}
+
+				err := encoder.EncodeAsset(asset)
+				if err != nil {
+					return fmt.Errorf("Unable to write asset %v: %v", asset, err)
+				}
 			}
 		}
 

internal/handlers/handler_read_graph.go

@@ -23,7 +23,7 @@ func GetGraphRead(registry sources.Registry, graphDB knowledge.GraphDB) http.Han
 		}
 
 		g := knowledge.NewGraph()
-		if err := graphDB.ReadGraph(r.Context(), source, g); err != nil {
+		if err := graphDB.ReadGraph(r.Context(), source, knowledge.NewGraphEncoder(w)); err != nil {
 			ReplyWithInternalError(w, err)
 			return
 		}

internal/knowledge/graph_encoder_decoder.go

@@ -0,0 +1,95 @@
+package knowledge
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"io"
+)
+
+// ******************************* DECODER ***********************************
+
+// GraphEncoder represent a graph encoder
+type GraphEncoder struct {
+	writer io.Writer
+
+	jsonEncoder *json.Encoder
+}
+
+// NewGraphEncoder create an instance of a graph encoder
+func NewGraphEncoder(w io.Writer) *GraphEncoder {
+	return &GraphEncoder{
+		writer:      w,
+		jsonEncoder: json.NewEncoder(w),
+	}
+}
+
+// EncodeRelation encode a relation
+func (ge *GraphEncoder) EncodeRelation(relation Relation) error {
+	_, err := fmt.Fprint(ge.writer, "R")
+	if err != nil {
+		return err
+	}
+	err = ge.jsonEncoder.Encode(relation)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// EncodeAsset encode an asset
+func (ge *GraphEncoder) EncodeAsset(asset Asset) error {
+	_, err := fmt.Fprint(ge.writer, "A")
+	if err != nil {
+		return err
+	}
+	err = ge.jsonEncoder.Encode(asset)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// ******************************* DECODER ***********************************
+
+// GraphDecoder represent a graph decoder
+type GraphDecoder struct {
+	reader io.Reader
+
+	jsonDecoder *json.Decoder
+}
+
+// NewGraphDecoder create an instance of a graph encoder
+func NewGraphDecoder(r io.Reader) *GraphDecoder {
+	return &GraphDecoder{
+		reader:      r,
+		jsonDecoder: json.NewDecoder(r),
+	}
+}
+
+// Decode decode a graph
+func (ge *GraphDecoder) Decode(graph *Graph) error {
+	scanner := bufio.NewScanner(ge.reader)
+	for scanner.Scan() {
+		line := scanner.Text()
+		switch line[0] {
+		case 'A':
+			var asset Asset
+			err := json.Unmarshal([]byte(line[1:]), &asset)
+			if err != nil {
+				return err
+			}
+			graph.AddAsset(asset.Type, asset.Key)
+		case 'R':
+			var relation Relation
+			err := json.Unmarshal([]byte(line[1:]), &relation)
+			if err != nil {
+				return err
+			}
+			graph.AddAsset(relation.From.Type, relation.From.Key)
+			graph.AddAsset(relation.To.Type, relation.To.Key)
+			graph.AddRelation(relation.From, relation.Type, relation.To)
+		}
+	}
+	return nil
+}

internal/knowledge/graph_encoder_decoder_test.go

@@ -0,0 +1,86 @@
+package knowledge
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	Asset1 = Asset{Type: "type1", Key: "value1"}
+	Asset2 = Asset{Type: "type2", Key: "value2"}
+	Asset3 = Asset{Type: "type3", Key: "value3"}
+)
+
+var (
+	Relation1 = Relation{From: AssetKey(Asset1), Type: "is_linked_to", To: AssetKey(Asset2)}
+	Relation2 = Relation{From: AssetKey(Asset1), Type: "has_relation_with", To: AssetKey(Asset2)}
+	Relation3 = Relation{From: AssetKey(Asset1), Type: "has_weird_relation_with", To: AssetKey(Asset3)}
+)
+
+func TestEncodeAssets(t *testing.T) {
+	buff := bytes.NewBuffer(nil)
+
+	encoder := NewGraphEncoder(buff)
+
+	assert.NoError(t, encoder.EncodeAsset(Asset1))
+	assert.NoError(t, encoder.EncodeAsset(Asset2))
+
+	assert.Equal(t,
+		"A{\"type\":\"type1\",\"key\":\"value1\"}\nA{\"type\":\"type2\",\"key\":\"value2\"}\n",
+		string(buff.Bytes()))
+}
+
+func TestEncodeRelations(t *testing.T) {
+	buff := bytes.NewBuffer(nil)
+
+	encoder := NewGraphEncoder(buff)
+
+	assert.NoError(t, encoder.EncodeRelation(Relation1))
+	assert.NoError(t, encoder.EncodeRelation(Relation2))
+
+	assert.Equal(t,
+		"R{\"type\":\"is_linked_to\",\"from\":{\"type\":\"type1\",\"key\":\"value1\"},\"to\":{\"type\":\"type2\",\"key\":\"value2\"}}\nR{\"type\":\"has_relation_with\",\"from\":{\"type\":\"type1\",\"key\":\"value1\"},\"to\":{\"type\":\"type2\",\"key\":\"value2\"}}\n",
+		string(buff.Bytes()))
+}
+
+func TestEncodeAssetsAndRelations(t *testing.T) {
+	buff := bytes.NewBuffer(nil)
+
+	encoder := NewGraphEncoder(buff)
+
+	assert.NoError(t, encoder.EncodeRelation(Relation1))
+	assert.NoError(t, encoder.EncodeRelation(Relation2))
+
+	assert.Equal(t,
+		"A{\"type\":\"type3\",\"key\":\"value3\"}\nR{\"type\":\"has_relation_with\",\"from\":{\"type\":\"type1\",\"key\":\"value1\"},\"to\":{\"type\":\"type2\",\"key\":\"value2\"}}\n",
+		string(buff.Bytes()))
+}
+
+func TestEncodeAndDecode(t *testing.T) {
+	buff := bytes.NewBuffer(nil)
+
+	encoder := NewGraphEncoder(buff)
+
+	assert.NoError(t, encoder.EncodeRelation(Relation1))
+	assert.NoError(t, encoder.EncodeRelation(Relation2))
+	assert.NoError(t, encoder.EncodeRelation(Relation3))
+
+	decoder := NewGraphDecoder(buff)
+
+	graph := NewGraph()
+
+	assert.NoError(t, decoder.Decode(graph))
+
+	assert.Len(t, graph.Assets(), 3)
+	assert.Len(t, graph.Relations(), 3)
+
+	assert.True(t, graph.HasAsset(Asset1))
+	assert.True(t, graph.HasAsset(Asset2))
+	assert.True(t, graph.HasAsset(Asset3))
+
+	assert.True(t, graph.HasRelation(Relation1))
+	assert.True(t, graph.HasRelation(Relation2))
+	assert.True(t, graph.HasRelation(Relation3))
+}

internal/knowledge/graphdb.go

@@ -19,7 +19,7 @@ type GraphDB interface {
 
 	InitializeSchema() error
 
-	ReadGraph(ctx context.Context, source string, graph *Graph) error
+	ReadGraph(ctx context.Context, source string, encoder *GraphEncoder) error
 
 	// Atomic operations on the graph
 	InsertAssets(ctx context.Context, source string, assets []Asset) error