get schema with augmented/prefixed path

pkg/schema/container.go

@@ -108,6 +108,16 @@ func getMandatoryChildren(e *yang.Entry) []*sdcpb.MandatoryChild {
 			result = append(result, c)
 		}
 	}
+	// include mandatory children from augments
+	for _, v := range e.Augmented {
+		if v.Mandatory == yang.TSTrue {
+			c := &sdcpb.MandatoryChild{
+				Name:    v.Name,
+				IsState: isState(v),
+			}
+			result = append(result, c)
+		}
+	}
 	return result
 }
 
@@ -157,6 +167,16 @@ func getChoiceInfo(e *yang.Entry) *sdcpb.ChoiceInfo {
 
 		processChoice(de, ci)
 	}
+	// also consider choices introduced via augments
+	for _, de := range e.Augmented {
+		if !de.IsChoice() {
+			continue
+		}
+		if ci == nil {
+			ci = &sdcpb.ChoiceInfo{Choice: map[string]*sdcpb.ChoiceInfoChoice{}}
+		}
+		processChoice(de, ci)
+	}
 	return ci
 }
 

pkg/schema/expand.go

@@ -41,7 +41,7 @@ func (sc *Schema) ExpandPath(p *sdcpb.Path, dt sdcpb.DataType) ([]*sdcpb.Path, e
 	for _, k := range strings.Fields(e.Key) {
 		keys[k] = struct{}{}
 	}
-	for _, c := range e.Dir {
+	for _, c := range getChildren(e) {
 		// skip keys
 		if _, ok := keys[c.Name]; ok {
 			continue
@@ -109,7 +109,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 			kmap[k] = struct{}{}
 		}
 
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			if _, ok := kmap[c.Name]; ok {
 				continue
 			}
@@ -126,7 +126,7 @@ func (sc *Schema) getPathElems(e *yang.Entry, dt sdcpb.DataType) [][]*sdcpb.Path
 	case e.IsContainer():
 		log.Debugf("got container: %s", e.Name)
 		containerPE := &sdcpb.PathElem{Name: e.Name, Key: make(map[string]string)}
-		for _, c := range e.Dir {
+		for _, c := range getChildren(e) {
 			log.Debugf("container parent adding child: %s", c.Name)
 			childrenPE := sc.getPathElems(c, dt)
 

pkg/schema/object.go

@@ -182,7 +182,7 @@ func getEntry(e *yang.Entry, pe []string) (*yang.Entry, error) {
 		}
 		return e, nil
 	default:
-		if e.Dir == nil {
+		if e.Dir == nil && len(e.Augmented) == 0 {
 			return nil, errors.New("not found")
 		}
 		for _, ee := range getChildren(e) {
@@ -261,6 +261,22 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		return nil
 	}
 
+	// helper to find a direct child by name from Dir or Augmented
+	childByName := func(parent *yang.Entry, name string) *yang.Entry {
+		if parent == nil {
+			return nil
+		}
+		if ce, ok := parent.Dir[name]; ok {
+			return ce
+		}
+		for _, ae := range parent.Augmented {
+			if ae.Name == name {
+				return ae
+			}
+		}
+		return nil
+	}
+
 	switch {
 	case e.IsList():
 		if cpe.GetKey() == nil {
@@ -281,7 +297,7 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 			return nil
 		}
 		nxt := pe[count]
-		if ee, ok := e.Dir[nxt]; ok {
+		if ee := childByName(e, nxt); ee != nil {
 			return sc.buildPath(pe[count:], p, ee)
 		}
 		// find choices/cases
@@ -316,20 +332,20 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 		return fmt.Errorf("case %s - unknown element %s", e.Name, pe[0])
 	case e.IsContainer():
 		// implicit case: child with same name which is a choice
-		if ee, ok := e.Dir[pe[0]]; ee != nil && ok {
+		if ee := childByName(e, pe[0]); ee != nil {
 			if ee.IsChoice() {
 				return sc.buildPath(pe[1:], p, ee)
 			}
 		}
 
 		p.Elem = append(p.Elem, cpe)
-		if ee, ok := e.Dir[pe[0]]; ok {
+		if ee := childByName(e, pe[0]); ee != nil {
 			return sc.buildPath(pe, p, ee)
 		}
 		if lpe == 1 {
 			return nil
 		}
-		if ee, ok := e.Dir[pe[1]]; ok {
+		if ee := childByName(e, pe[1]); ee != nil {
 			return sc.buildPath(pe[1:], p, ee)
 		}
 		// find choice/case
@@ -359,14 +375,22 @@ func (sc *Schema) buildPath(pe []string, p *sdcpb.Path, e *yang.Entry) error {
 func getChildren(e *yang.Entry) []*yang.Entry {
 	switch {
 	case e.IsChoice(), e.IsCase(), e.IsContainer(), e.IsList():
-		rs := make([]*yang.Entry, 0, len(e.Dir))
+		rs := make([]*yang.Entry, 0, len(e.Dir)+len(e.Augmented))
 		for _, ee := range e.Dir {
 			if ee.IsChoice() || ee.IsCase() {
 				rs = append(rs, getChildren(ee)...)
 				continue
 			}
 			rs = append(rs, ee)
 		}
+		// add augmented children as well
+		for _, ee := range e.Augmented {
+			if ee.IsChoice() || ee.IsCase() {
+				rs = append(rs, getChildren(ee)...)
+				continue
+			}
+			rs = append(rs, ee)
+		}
 		//sort.Slice(rs, sortFn(rs))
 		return rs
 		// case e.IsCase():
@@ -562,21 +586,41 @@ func (sc *Schema) findChoiceCase(e *yang.Entry, pe []string) (*yang.Entry, error
 	if len(pe) == 0 {
 		return e, nil
 	}
-	for _, ee := range e.Dir {
-		if !ee.IsChoice() {
-			continue
+	// pe is expected to contain at least the current element name at index 0 and
+	// the sought child element name at index 1.
+	//
+	// This is used from container/list resolution paths where the next element may
+	// live under a choice/case. Case nodes do not exist in the data tree, so we
+	// must search through cases to find the actual schema node.
+	if len(pe) < 2 {
+		return nil, fmt.Errorf("unknown element %s", pe[0])
+	}
+
+	// scan choice nodes in both direct and augmented children
+	choices := make([]*yang.Entry, 0)
+	for _, child := range e.Dir {
+		if child != nil && child.IsChoice() {
+			choices = append(choices, child)
 		}
-		if eee, ok := ee.Dir[pe[1]]; ok && !eee.IsCase() {
-			return eee, nil
+	}
+	for _, child := range e.Augmented {
+		if child != nil && child.IsChoice() {
+			choices = append(choices, child)
 		}
-		// assume there was a case obj,
-		// search one step deeper
-		for _, eee := range ee.Dir {
-			if !eee.IsCase() {
+	}
+
+	for _, choice := range choices {
+		// implicit case: choice directly contains the data node
+		if direct, ok := choice.Dir[pe[1]]; ok && direct != nil && !direct.IsCase() {
+			return direct, nil
+		}
+		// explicit cases: the data node is under a case
+		for _, cc := range choice.Dir {
+			if cc == nil || !cc.IsCase() {
 				continue
 			}
-			if eeee, ok := eee.Dir[pe[1]]; ok {
-				return eeee, nil
+			if target, ok := cc.Dir[pe[1]]; ok && target != nil {
+				return target, nil
 			}
 		}
 	}

pkg/schema/references.go

@@ -70,6 +70,21 @@ func (sc *Schema) buildReferences(e *yang.Entry) error {
 			return err
 		}
 	}
+	// also recurse into augmented children
+	for _, ce := range e.Augmented {
+		if ce.IsCase() || ce.IsChoice() {
+			for _, cce := range ce.Dir {
+				err := sc.buildReferences(cce)
+				if err != nil {
+					return err
+				}
+			}
+		}
+		err := sc.buildReferences(ce)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 

pkg/schema/schema.go

@@ -141,14 +141,28 @@ func (s *Schema) Walk(e *yang.Entry, fn func(ec *yang.Entry) error) error {
 				return err
 			}
 		}
+		// also walk augments merged into this entry
+		for _, ee := range e.Augmented {
+			err = s.Walk(ee, fn)
+			if err != nil {
+				return err
+			}
+		}
 		return nil
 	}
 	err = fn(e)
 	if err != nil {
 		return err
 	}
-	for _, e := range e.Dir {
-		err = s.Walk(e, fn)
+	for _, ce := range e.Dir {
+		err = s.Walk(ce, fn)
+		if err != nil {
+			return err
+		}
+	}
+	// include augmented children at every level
+	for _, ce := range e.Augmented {
+		err = s.Walk(ce, fn)
 		if err != nil {
 			return err
 		}