feat: include latest probe in lineage response (#180)

- Return node attributes in Lineage response for GetGraph call. Node
  attributes include latest probe for asset nodes in the Graph. This
  enables showing live lineage.
- Use context for gRPC dial and compass service handler in server.
- Add support for fetching probes with filters from asset repository.
  Supported filters:
  - AssetURNs - zero or more asset URNs.
  - MaxRows - Number of rows for each Asset URN ordered by latest first.
  - OlderThan,NewerThan - Data range.
  - Fix some lint errors

Author: sudo-suhas

Diff for: cli/utils.go

@@ -4,7 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"os"
 	"path/filepath"
 	"strings"
 
@@ -13,7 +13,7 @@ import (
 )
 
 func parseFile(filePath string, v protoreflect.ProtoMessage) error {
-	b, err := ioutil.ReadFile(filePath)
+	b, err := os.ReadFile(filePath)
 	if err != nil {
 		return err
 	}

Diff for: core/asset/asset.go

@@ -22,6 +22,7 @@ type Repository interface {
 	DeleteByURN(ctx context.Context, urn string) error
 	AddProbe(ctx context.Context, assetURN string, probe *Probe) error
 	GetProbes(ctx context.Context, assetURN string) ([]Probe, error)
+	GetProbesWithFilter(ctx context.Context, flt ProbesFilter) (map[string][]Probe, error)
 }
 
 // Asset is a model that wraps arbitrary data with Compass' context

Diff for: core/asset/lineage.go

@@ -33,6 +33,11 @@ type LineageRepository interface {
 
 type LineageGraph []LineageEdge
 
+type Lineage struct {
+	Edges     []LineageEdge             `json:"edges"`
+	NodeAttrs map[string]NodeAttributes `json:"node_attrs"`
+}
+
 type LineageEdge struct {
 	// Source represents source's node ID
 	Source string `json:"source"`
@@ -43,3 +48,11 @@ type LineageEdge struct {
 	// Prop is a map containing extra information about the edge
 	Prop map[string]interface{} `json:"prop"`
 }
+
+type NodeAttributes struct {
+	Probes ProbesInfo `json:"probes"`
+}
+
+type ProbesInfo struct {
+	Latest Probe `json:"latest"`
+}

Diff for: core/asset/mocks/asset_repository.go

@@ -0,0 +1,10 @@
+package asset
+
+import "time"
+
+type ProbesFilter struct {
+	AssetURNs []string
+	MaxRows   int
+	NewerThan time.Time
+	OlderThan time.Time
+}

Diff for: core/asset/probes_filter.go

@@ -108,8 +108,30 @@ func (s *Service) AddProbe(ctx context.Context, assetURN string, probe *Probe) e
 	return s.assetRepository.AddProbe(ctx, assetURN, probe)
 }
 
-func (s *Service) GetLineage(ctx context.Context, urn string, query LineageQuery) (LineageGraph, error) {
-	return s.lineageRepository.GetGraph(ctx, urn, query)
+func (s *Service) GetLineage(ctx context.Context, urn string, query LineageQuery) (Lineage, error) {
+	edges, err := s.lineageRepository.GetGraph(ctx, urn, query)
+	if err != nil {
+		return Lineage{}, fmt.Errorf("get lineage: get graph edges: %w", err)
+	}
+
+	urns := newUniqueStrings(len(edges))
+	urns.add(urn)
+	for _, edge := range edges {
+		urns.add(edge.Source, edge.Target)
+	}
+
+	assetProbes, err := s.assetRepository.GetProbesWithFilter(ctx, ProbesFilter{
+		AssetURNs: urns.list(),
+		MaxRows:   1,
+	})
+	if err != nil {
+		return Lineage{}, fmt.Errorf("get lineage: get latest probes: %w", err)
+	}
+
+	return Lineage{
+		Edges:     edges,
+		NodeAttrs: buildNodeAttrs(assetProbes),
+	}, nil
 }
 
 func (s *Service) GetTypes(ctx context.Context, flt Filter) (map[Type]int, error) {
@@ -131,3 +153,18 @@ func isValidUUID(u string) bool {
 	_, err := uuid.Parse(u)
 	return err == nil
 }
+
+func buildNodeAttrs(assetProbes map[string][]Probe) map[string]NodeAttributes {
+	nodeAttrs := make(map[string]NodeAttributes, len(assetProbes))
+	for urn, probes := range assetProbes {
+		if len(probes) == 0 {
+			continue
+		}
+
+		nodeAttrs[urn] = NodeAttributes{
+			Probes: ProbesInfo{Latest: probes[0]},
+		}
+	}
+
+	return nodeAttrs
+}