Beispiel #1
0
func sortEdgeSets(sets map[string]*gpb.EdgeSet) map[string]*gpb.EdgeSet {
	for _, set := range sets {
		for _, g := range set.Groups {
			sort.Sort(xrefs.ByOrdinal(g.Edge))
		}
	}
	return sets
}
Beispiel #2
0
func (d *DB) edges(ctx context.Context, req *xpb.EdgesRequest, edgeFilter func(kind string) bool) (*xpb.EdgesReply, error) {
	tickets, err := xrefs.FixTickets(req.Ticket)
	if err != nil {
		return nil, err
	}

	pageSize := int(req.PageSize)
	if pageSize <= 0 {
		pageSize = defaultPageSize
	} else if pageSize > maxPageSize {
		pageSize = maxPageSize
	}

	var pageOffset int
	if req.PageToken != "" {
		rec, err := base64.StdEncoding.DecodeString(req.PageToken)
		if err != nil {
			return nil, fmt.Errorf("invalid page_token: %q", req.PageToken)
		}
		var t ipb.PageToken
		if err := proto.Unmarshal(rec, &t); err != nil || t.Index < 0 {
			return nil, fmt.Errorf("invalid page_token: %q", req.PageToken)
		}
		pageOffset = int(t.Index)
	}

	// Select only the edges from the given source tickets
	setQ, args := sqlSetQuery(1, tickets)
	query := fmt.Sprintf(`
SELECT * FROM AllEdges
WHERE source IN %s`, setQ)

	// Filter by edges kinds, if given
	if len(req.Kind) > 0 {
		kSetQ, kArgs := sqlSetQuery(1+len(args), req.Kind)
		query += fmt.Sprintf(`
AND kind IN %s`, kSetQ)
		args = append(args, kArgs...)
	}

	// Scan edge sets/groups in order; necessary for CrossReferences
	query += " ORDER BY source, kind, ordinal"

	// Seek to the requested page offset (req.PageToken.Index).  We don't use
	// LIMIT here because we don't yet know how many edges will be filtered by
	// edgeFilter.
	query += fmt.Sprintf(" OFFSET $%d", len(args)+1)
	args = append(args, pageOffset)

	rs, err := d.Query(query, args...)
	if err != nil {
		return nil, fmt.Errorf("error querying for edges: %v", err)
	}
	defer closeRows(rs)

	var scanned int
	// edges := map { source -> kind -> target -> ordinal set }
	edges := make(map[string]map[string]map[string]map[int32]struct{}, len(tickets))
	for count := 0; count < pageSize && rs.Next(); scanned++ {
		var source, kind, target string
		var ordinal int
		if err := rs.Scan(&source, &kind, &target, &ordinal); err != nil {
			return nil, fmt.Errorf("edges scan error: %v", err)
		}
		if edgeFilter != nil && !edgeFilter(kind) {
			continue
		}
		count++

		groups, ok := edges[source]
		if !ok {
			groups = make(map[string]map[string]map[int32]struct{})
			edges[source] = groups
		}
		targets, ok := groups[kind]
		if !ok {
			targets = make(map[string]map[int32]struct{})
			groups[kind] = targets
		}
		ordinals, ok := targets[target]
		if !ok {
			ordinals = make(map[int32]struct{})
			targets[target] = ordinals
		}
		ordinals[int32(ordinal)] = struct{}{}
	}

	reply := &xpb.EdgesReply{EdgeSets: make(map[string]*xpb.EdgeSet, len(edges))}
	nodeTickets := stringset.New()
	for src, groups := range edges {
		gs := make(map[string]*xpb.EdgeSet_Group, len(groups))
		nodeTickets.Add(src)
		for kind, targets := range groups {
			edges := make([]*xpb.EdgeSet_Group_Edge, 0, len(targets))
			for ticket, ordinals := range targets {
				for ordinal := range ordinals {
					edges = append(edges, &xpb.EdgeSet_Group_Edge{
						TargetTicket: ticket,
						Ordinal:      ordinal,
					})
				}
				nodeTickets.Add(ticket)
			}
			sort.Sort(xrefs.ByOrdinal(edges))
			gs[kind] = &xpb.EdgeSet_Group{
				Edge: edges,
			}
		}
		reply.EdgeSets[src] = &xpb.EdgeSet{
			Groups: gs,
		}
	}

	// If there is another row, there is a NextPageToken.
	if rs.Next() {
		rec, err := proto.Marshal(&ipb.PageToken{Index: int32(pageOffset + scanned)})
		if err != nil {
			return nil, fmt.Errorf("internal error: error marshalling page token: %v", err)
		}
		reply.NextPageToken = base64.StdEncoding.EncodeToString(rec)
	}

	// TODO(schroederc): faster node lookups
	if len(req.Filter) > 0 && len(nodeTickets) > 0 {
		nodes, err := d.Nodes(ctx, &xpb.NodesRequest{
			Ticket: nodeTickets.Slice(),
			Filter: req.Filter,
		})
		if err != nil {
			return nil, fmt.Errorf("error filtering nodes:%v", err)
		}
		reply.Nodes = nodes.Nodes
	}

	return reply, nil
}