Plan section access improvement

Author: flotter

internals/plan/extensions_test.go

@@ -266,7 +266,40 @@ var extensionTests = []struct {
 		},
 		errorString: "cannot validate plan section .* cannot find .* as required by .*",
 	},
-	// Index 8: Load file layers
+	// Index 8: Load nothing
+	{
+		extensions: []extension{
+			extension{
+				field: "x-field",
+				ext:   &xExtension{},
+			},
+			extension{
+				field: "y-field",
+				ext:   &yExtension{},
+			},
+		},
+		files: []*planInput{
+			&planInput{
+				order: 1,
+				label: "layer-x",
+				yaml: `
+					summary: x
+					description: desc-x
+				`,
+			},
+			&planInput{
+				order: 2,
+				label: "layer-y",
+				yaml: `
+					summary: y
+					description: desc-y
+				`,
+			},
+		},
+		result:     &planResult{},
+		resultYaml: string("{}\n"),
+	},
+	// Index 9: Load file layers
 	{
 		extensions: []extension{
 			extension{
@@ -369,12 +402,26 @@ func (ps *planSuite) TestPlanExtensions(c *C) {
 		} else {
 			if planTest.result != nil {
 				// Verify "x-field" data.
-				x := p.Section(xField).(*xSection)
-				c.Assert(x.Entries, DeepEquals, planTest.result.x.Entries)
+				var x *xSection
+				err = p.Section(xField, &x)
+				c.Assert(err, IsNil)
+				if planTest.result.x != nil {
+					c.Assert(x.Entries, DeepEquals, planTest.result.x.Entries)
+				} else {
+					// No section must result in nil.
+					c.Assert(x, IsNil)
+				}
 
 				// Verify "y-field" data.
-				y := p.Section(yField).(*ySection)
-				c.Assert(y.Entries, DeepEquals, planTest.result.y.Entries)
+				var y *ySection
+				err = p.Section(yField, &y)
+				c.Assert(err, IsNil)
+				if planTest.result.y != nil {
+					c.Assert(y.Entries, DeepEquals, planTest.result.y.Entries)
+				} else {
+					// No section must result in nil.
+					c.Assert(x, IsNil)
+				}
 			}
 
 			// Verify combined plan YAML.
@@ -430,16 +477,22 @@ func (x xExtension) CombineSections(sections ...plan.LayerSection) (plan.LayerSe
 }
 
 func (x xExtension) ValidatePlan(p *plan.Plan) error {
-	planXSection := p.Section(xField)
-	if planXSection != nil {
-		planYSection := p.Section(yField)
-		if planYSection == nil {
+	var xs *xSection
+	err := p.Section(xField, &xs)
+	if err != nil {
+		return err
+	}
+	if xs != nil {
+		var ys *ySection
+		err = p.Section(yField, &ys)
+		if err != nil {
+			return err
+		}
+		if ys == nil {
 			return fmt.Errorf("cannot validate %v field without %v field", xField, yField)
 		}
-		ys := planYSection.(*ySection)
 
 		// Make sure every Y field in X refer to an existing Y entry.
-		xs := planXSection.(*xSection)
 		for xEntryField, xEntryValue := range xs.Entries {
 			for _, yReference := range xEntryValue.Y {
 				found := false

internals/plan/plan.go

@@ -1,4 +1,4 @@
-// Copyright (c) 2021 Canonical Ltd
+// Copyright (c) 2024 Canonical Ltd
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -90,10 +90,28 @@ type Plan struct {
 	Sections map[string]LayerSection `yaml:",inline,omitempty"`
 }
 
-// Section retrieves a section from the plan. Returns nil if
-// the field does not exist.
-func (p *Plan) Section(field string) LayerSection {
-	return p.Sections[field]
+// Section retrieves a section from the plan, if it exists.
+func (p *Plan) Section(field string, out any) error {
+	if _, found := layerExtensions[field]; !found {
+		return fmt.Errorf("cannot find registered extension for field %q", field)
+	}
+
+	outVal := reflect.ValueOf(out)
+	if outVal.Kind() != reflect.Ptr || outVal.IsNil() {
+		return fmt.Errorf("cannot read non pointer to section interface type %q", outVal.Kind())
+	}
+
+	section, exists := p.Sections[field]
+	if exists {
+		sectionVal := reflect.ValueOf(section)
+		sectionType := sectionVal.Type()
+		outValPtrType := outVal.Elem().Type()
+		if !sectionType.AssignableTo(outValPtrType) {
+			return fmt.Errorf("cannot assign value of type %s to out argument of type %s", sectionType, outValPtrType)
+		}
+		outVal.Elem().Set(sectionVal)
+	}
+	return nil
 }
 
 type Layer struct {
@@ -108,17 +126,6 @@ type Layer struct {
 	Sections map[string]LayerSection `yaml:",inline,omitempty"`
 }
 
-// addSection adds a new section to the layer.
-func (layer *Layer) addSection(field string, section LayerSection) {
-	layer.Sections[field] = section
-}
-
-// Section retrieves a layer section from a layer. Returns nil if
-// the field does not exist.
-func (layer *Layer) Section(field string) LayerSection {
-	return layer.Sections[field]
-}
-
 type Service struct {
 	// Basic details
 	Name        string         `yaml:"-"`
@@ -708,7 +715,7 @@ func CombineLayers(layers ...*Layer) (*Layer, error) {
 	for field, extension := range layerExtensions {
 		var sections []LayerSection
 		for _, layer := range layers {
-			if section := layer.Section(field); section != nil {
+			if section := layer.Sections[field]; section != nil {
 				sections = append(sections, section)
 			}
 		}
@@ -723,7 +730,7 @@ func CombineLayers(layers ...*Layer) (*Layer, error) {
 			}
 			// We support the ability for a valid combine to result in an omitted section.
 			if combinedSection != nil {
-				combined.addSection(field, combinedSection)
+				combined.Sections[field] = combinedSection
 			}
 		}
 	}
@@ -1182,7 +1189,7 @@ func ParseLayer(order int, label string, data []byte) (*Layer, error) {
 					}
 				}
 				if extendedSection != nil {
-					layer.addSection(field, extendedSection)
+					layer.Sections[field] = extendedSection
 				}
 			} else {
 				// At the top level we do not ignore keys we do not understand.