feat: add client customization option (#20)

* feat: add client customization

* chore: fix linting

* feat: allow excluding client function by resource name

codegen/client.go.tmpl

@@ -10,29 +10,10 @@ import (
 
 type {{ .Name }} interface {
 
-    /* custom method signatures */
-
-    {{ range $k, $v := .CustomFunctions }}
-    {{ $v.Signature }}
-    {{- end }}
-
-    /* client methods generated based on resource generation */
-
     {{- range $k, $v := .Functions }}
-    /* client methods for {{ $v.Name }} API */
 
-    // Get{{ $v.Name }} returns {{ $v.Name }} resource{{ if not $v.IsSetting }} by its ID{{ end }}
-    Get{{ $v.Name }}(ctx context.Context, site{{ if not $v.IsSetting }}, id{{ end }} string) (*{{ $v.Name }}, error)
-    // Update{{ $v.Name }} updates {{ $v.Name }} resource{{ if not $v.IsSetting }} by its ID{{ end }}
-    Update{{ $v.Name }}(ctx context.Context, site string, d *{{ $v.Name }}) (*{{ $v.Name }}, error)
-    {{- if not $v.IsSetting }}
-    // List{{ $v.Name }} returns list of {{ $v.Name }} resources
-    List{{ $v.Name }}(ctx context.Context, site string) ([]{{ $v.Name }}, error)
-    // Delete{{ $v.Name }} deletes {{ $v.Name }} resource by its ID
-    Delete{{ $v.Name }}(ctx context.Context, site, id string) error
-    // Create{{ $v.Name }} creates new {{ $v.Name }} resource
-    Create{{ $v.Name }}(ctx context.Context, site string, d *{{ $v.Name }}) (*{{ $v.Name }}, error)
-    {{ end }}
+    {{ if $v.Comment }}// {{ $v.Comment }}{{ end }}
+    {{ $v.Signature }}
 
     {{- end }}
 }

codegen/clients.go

@@ -3,13 +3,16 @@ package main
 import (
 	_ "embed"
 	"fmt"
+	"sort"
 	"strings"
 )
 
 // ClientFunction is the interface for client functions.
 type ClientFunction interface {
 	Name() string
-	IsSetting() bool
+	ResourceName() string
+	Comment() string
+	Signature() string
 }
 
 type FunctionParam struct {
@@ -17,21 +20,54 @@ type FunctionParam struct {
 	Type string
 }
 
+type Comment struct {
+	comment      string
+	resourceName string
+}
+
+func (c *Comment) Name() string {
+	return ""
+}
+
+func (c *Comment) Comment() string {
+	return c.comment
+}
+
+func (c *Comment) Signature() string {
+	return ""
+}
+
+func (c *Comment) ResourceName() string {
+	return c.resourceName
+}
+
 // CustomClientFunction represents a custom client function definition.
 type CustomClientFunction struct {
-	Name             string
-	Parameters       []FunctionParam
-	ReturnParameters []string
-	Comment          string
+	Resource         string          `yaml:"resourceName"`
+	FunctionName     string          `yaml:"name"`
+	Parameters       []FunctionParam `yaml:"params"`
+	ReturnParameters []string        `yaml:"returns"`
+	FunctionComment  string          `yaml:"comment"`
+}
+
+func (c *CustomClientFunction) Name() string {
+	return c.FunctionName
+}
+
+func (c *CustomClientFunction) ResourceName() string {
+	return c.Resource
 }
 
 // Signature returns the signature string for the custom client function.
 func (c *CustomClientFunction) Signature() string {
-	var b strings.Builder
-	if c.Comment != "" {
-		b.WriteString(fmt.Sprintf("// %s %s\n", c.Name, c.Comment))
+	if c.Name() == "" {
+		return ""
 	}
-	b.WriteString(c.Name)
+	var b strings.Builder
+	//if c.comment != "" {
+	//	b.WriteString(fmt.Sprintf("// %s %s\n", c.Name, c.Comment))
+	//}
+	b.WriteString(c.Name())
 	b.WriteString("(")
 
 	// Build parameters without trailing comma
@@ -52,20 +88,104 @@ func (c *CustomClientFunction) Signature() string {
 	return b.String()
 }
 
+func (c *CustomClientFunction) Comment() string {
+	return c.FunctionComment
+}
+
 // ClientInfo represents the client information used for code generation.
 type ClientInfo struct {
-	Imports         []string
-	Functions       []ClientFunction
-	CustomFunctions []CustomClientFunction
+	Imports   []string
+	Functions []ClientFunction
+}
+
+type ClientInfoBuilder struct {
+	imports   []string
+	functions []ClientFunction
+}
+
+func (c *ClientInfoBuilder) AddFunction(f ClientFunction) *ClientInfoBuilder { //nolint: unparam
+	c.functions = append(c.functions, f)
+	return c
+}
+
+func (c *ClientInfoBuilder) AddFunctions(f []CustomClientFunction) *ClientInfoBuilder {
+	for _, v := range f {
+		c.functions = append(c.functions, &v)
+	}
+	return c
+}
+
+func (c *ClientInfoBuilder) addResourceFunction(actionName, resourceName, comment string, additionalParams []FunctionParam, additionalReturns []string) {
+	fName := fmt.Sprintf("%s%s", actionName, resourceName)
+	params := []FunctionParam{
+		{"ctx", "context.Context"},
+		{"site", "string"},
+	}
+	params = append(params, additionalParams...)
+	returns := additionalReturns
+	returns = append(returns, "error")
+	f := CustomClientFunction{
+		FunctionName:     fName,
+		Resource:         resourceName,
+		Parameters:       params,
+		ReturnParameters: returns,
+		FunctionComment:  fmt.Sprintf("%s %s", fName, comment),
+	}
+	c.AddFunction(&f)
+}
+
+func singlePointerReturn(name string) []string {
+	return []string{"*" + name}
+}
+
+func singlePointerParam(name string) []FunctionParam {
+	return []FunctionParam{{strings.ToLower(name[0:1]), "*" + name}}
+}
+
+func (c *ClientInfoBuilder) AddResource(r *Resource) *ClientInfoBuilder {
+	c.AddFunction(&Comment{comment: fmt.Sprintf("client methods for %s resource", r.Name()), resourceName: r.Name()})
+	if r.IsSetting() {
+		c.addResourceFunction("Get", r.Name(), "retrieves the settings for a resource", nil, singlePointerReturn(r.Name()))
+		c.addResourceFunction("Update", r.Name(), "updates a resource", singlePointerParam(r.Name()), singlePointerReturn(r.Name()))
+		return c
+	}
+	c.addResourceFunction("Get", r.Name(), "retrieves a resource", []FunctionParam{{"id", "string"}}, singlePointerReturn(r.Name()))
+	c.addResourceFunction("List", r.Name(), "lists the resources", nil, []string{"[]*" + r.Name()})
+	c.addResourceFunction("Create", r.Name(), "creates a resource", singlePointerParam(r.Name()), singlePointerReturn(r.Name()))
+	c.addResourceFunction("Update", r.Name(), "updates a resource", singlePointerParam(r.Name()), singlePointerReturn(r.Name()))
+	c.addResourceFunction("Delete", r.Name(), "deletes a resource", []FunctionParam{{"id", "string"}}, nil)
+	return c
+}
+
+func (c *ClientInfoBuilder) AddImport(i string) *ClientInfoBuilder {
+	c.imports = append(c.imports, i)
+	return c
+}
+
+func (c *ClientInfoBuilder) AddImports(i []string) *ClientInfoBuilder {
+	c.imports = append(c.imports, i...)
+	return c
+}
+
+func (c *ClientInfoBuilder) Build() *ClientInfo {
+	// Sort the functions by resource name and then by name.
+	sort.Slice(c.functions, func(i, j int) bool {
+		if c.functions[i].ResourceName() == c.functions[j].ResourceName() {
+			return c.functions[i].Signature() < c.functions[j].Signature()
+		}
+		return c.functions[i].ResourceName() < c.functions[j].ResourceName()
+	})
+
+	return newClientInfo(c.imports, c.functions)
+}
+
+func NewClientInfoBuilder() *ClientInfoBuilder {
+	return &ClientInfoBuilder{}
 }
 
 // newClientInfo creates ClientInfo from the provided resources.
-func newClientInfo(resources []*Resource) *ClientInfo {
-	functions := make([]ClientFunction, 0)
-	for _, resource := range resources {
-		functions = append(functions, resource)
-	}
-	return &ClientInfo{Functions: functions}
+func newClientInfo(imports []string, functions []ClientFunction) *ClientInfo {
+	return &ClientInfo{imports, functions}
 }
 
 //go:embed client.go.tmpl

codegen/clients_test.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -19,23 +18,21 @@ func TestCustomClientFunctionSignature(t *testing.T) {
 		{
 			name: "no comment, no params, no returns",
 			fn: CustomClientFunction{
-				Name: "Foo",
+				FunctionName: "Foo",
 			},
 			wantFunc: "Foo()",
 		},
 		{
 			name: "with comment, no params, no returns",
 			fn: CustomClientFunction{
-				Name:    "Bar",
-				Comment: "does something",
+				FunctionName: "Bar",
 			},
-			wantComment: "// Bar does something",
-			wantFunc:    "Bar()",
+			wantFunc: "Bar()",
 		},
 		{
 			name: "with one param and one return",
 			fn: CustomClientFunction{
-				Name:             "Baz",
+				FunctionName:     "Baz",
 				Parameters:       []FunctionParam{{"a", "int"}},
 				ReturnParameters: []string{"error"},
 			},
@@ -44,7 +41,7 @@ func TestCustomClientFunctionSignature(t *testing.T) {
 		{
 			name: "with multiple returns",
 			fn: CustomClientFunction{
-				Name:             "Qux",
+				FunctionName:     "Qux",
 				Parameters:       []FunctionParam{{"x", "string"}},
 				ReturnParameters: []string{"int", "error"},
 			},
@@ -53,13 +50,11 @@ func TestCustomClientFunctionSignature(t *testing.T) {
 		{
 			name: "with multiple params",
 			fn: CustomClientFunction{
-				Name:             "MultiParams",
+				FunctionName:     "MultiParams",
 				Parameters:       []FunctionParam{{"x", "string"}, {"y", "int"}},
 				ReturnParameters: []string{},
-				Comment:          "function with multiple parameters",
 			},
-			wantComment: "// MultiParams function with multiple parameters",
-			wantFunc:    "MultiParams(x string, y int)",
+			wantFunc: "MultiParams(x string, y int)",
 		},
 	}
 
@@ -67,19 +62,7 @@ func TestCustomClientFunctionSignature(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
 			a := assert.New(t)
-
-			got := tt.fn.Signature()
-
-			parts := strings.Split(got, "\n")
-			var comment, funcSig string
-			if len(tt.wantComment) > 0 {
-				comment = parts[0]
-				funcSig = parts[1]
-			} else {
-				funcSig = parts[0]
-			}
-			a.Equal(tt.wantComment, comment)
-			a.Equal(tt.wantFunc, funcSig)
+			a.Equal(tt.wantFunc, tt.fn.Signature())
 		})
 	}
 }
@@ -87,17 +70,16 @@ func TestCustomClientFunctionSignature(t *testing.T) {
 func TestGenerateCode(t *testing.T) {
 	t.Parallel()
 	a := assert.New(t)
-	ci := &ClientInfo{
-		Imports: []string{"fmt"},
-		CustomFunctions: []CustomClientFunction{
-			{
-				Name:             "TestFunc",
-				Parameters:       []FunctionParam{{"x", "int"}},
-				ReturnParameters: []string{"error"},
-				Comment:          "This is a test function",
-			},
-		},
-	}
+
+	b := NewClientInfoBuilder()
+	b.AddImport("fmt")
+	b.AddFunction(&CustomClientFunction{
+		FunctionName:     "TestFunc",
+		Parameters:       []FunctionParam{{"x", "int"}},
+		ReturnParameters: []string{"error"},
+		FunctionComment:  "This is a test function",
+	})
+	ci := b.Build()
 	code, err := ci.GenerateCode()
 	require.NoError(t, err)
 	a.NotEmpty(code, "GenerateCode() returned empty code")

codegen/customizations.yml

@@ -1,5 +1,33 @@
 ---
 customizations:
+  client:
+    excludeResources:
+      - "Setting*"  # Exclude all resources that start with "Setting"
+    functions:
+      - name: "Login"
+        comment: "Login logs in to the controller. Useful only for user/password authentication."
+        returns:
+          - "error"
+      - name: "Logout"
+        comment: "Logout logs out from the controller."
+        returns:
+          - "error"
+      - name: "BaseURL"
+        comment: "BaseURL returns the base URL of the controller."
+        returns:
+          - "string"
+      - name: "AdoptDevice"
+        comment: "AdoptDevice adopts a device by MAC address."
+        resourceName: "Device"
+        params:
+          - name: "ctx"
+            type: "context.Context"
+          - name: "site"
+            type: "string"
+          - name: "mac"
+            type: "string"
+        returns:
+          - "error"
   resources:
     Account:
       fields:

codegen/customize.go

@@ -4,6 +4,7 @@ import (
 	_ "embed"
 	"fmt"
 	"os"
+	"strings"
 
 	"gopkg.in/yaml.v3"
 )
@@ -13,10 +14,13 @@ const (
 	defaultCustomizationsPath     = "customizations.yml"
 )
 
+type Customizations struct {
+	Resources map[string]*ResourceCustomization `yaml:"resources"`
+	Client    *ClientCustomization              `yaml:"client"`
+}
+
 type Generate struct {
-	Customizations struct {
-		Resources map[string]*ResourceCustomization `yaml:"resources"`
-	} `yaml:"customizations"`
+	Customizations *Customizations `yaml:"customizations"`
 }
 
 type ResourceCustomization struct {
@@ -24,6 +28,12 @@ type ResourceCustomization struct {
 	Fields       map[string]*FieldCustomization `yaml:"fields"`
 }
 
+type ClientCustomization struct {
+	Imports          []string               `yaml:"imports"`
+	Functions        []CustomClientFunction `yaml:"functions"`
+	ExcludeResources []string               `yaml:"excludeResources"`
+}
+
 type FieldCustomization struct {
 	FieldName   string             `yaml:"-"`
 	Overrides   *FieldInfoOverride `yaml:",inline"`
@@ -114,7 +124,7 @@ func unmarshalCustomizationYaml(customizationsPath string) (*Generate, error) {
 	if err != nil {
 		return nil, err
 	}
-	err = yaml.Unmarshal(customizationsYml, &generate)
+	err = yaml.Unmarshal(customizationsYml, &generate) //nolint: musttag
 	if err != nil {
 		return nil, fmt.Errorf("failed unmarshalling YAML to Generate structure: %w", err)
 	}
@@ -125,33 +135,54 @@ func unmarshalCustomizationYaml(customizationsPath string) (*Generate, error) {
 			field.FieldName = fieldName
 		}
 	}
+
 	return &generate, nil
 }
 
-type YamlConfigCodeCustomizer struct {
-	Customizations map[string]*ResourceCustomization
+type CodeCustomizer struct {
+	Customizations Customizations
 }
 
-type CodeCustomizer interface {
-	ApplyToResource(resource *Resource)
-}
-
-type noopCustomizer struct{}
-
-func (noopCustomizer) ApplyToResource(resource *Resource) {}
-
-func NewCodeCustomizer(customizationsPath string) (CodeCustomizer, error) { //nolint: ireturn
+func NewCodeCustomizer(customizationsPath string) (*CodeCustomizer, error) {
 	generate, err := unmarshalCustomizationYaml(customizationsPath)
 	if err != nil {
 		return nil, err
 	}
-	return &YamlConfigCodeCustomizer{generate.Customizations.Resources}, nil
+	if generate.Customizations == nil {
+		generate.Customizations = &Customizations{}
+	}
+	return &CodeCustomizer{*generate.Customizations}, nil
 }
 
-func (r *YamlConfigCodeCustomizer) ApplyToResource(resource *Resource) {
-	for resourceName, resourceCustomization := range r.Customizations {
+func (r *CodeCustomizer) IsExcludedFromClient(resourceName string) bool {
+	for _, excludedResource := range r.Customizations.Client.ExcludeResources {
+		prefixedAll := strings.HasPrefix(excludedResource, "*")
+		suffixedAll := strings.HasSuffix(excludedResource, "*")
+		if prefixedAll && suffixedAll && strings.Contains(resourceName, excludedResource[1:len(excludedResource)-1]) {
+			return true
+		} else if prefixedAll && strings.HasSuffix(resourceName, excludedResource[1:]) {
+			return true
+		} else if suffixedAll && strings.HasPrefix(resourceName, excludedResource[:len(excludedResource)-1]) {
+			return true
+		} else if resourceName == excludedResource {
+			return true
+		}
+	}
+	return false
+}
+
+func (r *CodeCustomizer) ApplyToResource(resource *Resource) {
+	for resourceName, resourceCustomization := range r.Customizations.Resources {
 		if resource.StructName == resourceName {
 			resourceCustomization.ApplyTo(resource)
 		}
 	}
 }
+
+func (r *CodeCustomizer) ApplyToClient(client *ClientInfoBuilder) {
+	if client == nil || r.Customizations.Client == nil {
+		return
+	}
+	client.AddFunctions(r.Customizations.Client.Functions)
+	client.AddImports(r.Customizations.Client.Imports)
+}

codegen/generator.go

@@ -51,12 +51,17 @@ func generateCode(fieldsDir string, outDir string, customizer CodeCustomizer) er
 	if err != nil {
 		return fmt.Errorf("failed to build resources from downloaded fields: %w", err)
 	}
-	client := newClientInfo(resources)
+	cb := NewClientInfoBuilder()
+	customizer.ApplyToClient(cb)
 	for _, resource := range resources {
+		if customizer.IsExcludedFromClient(resource.Name()) {
+			continue
+		}
+		cb.AddResource(resource)
 		customizer.ApplyToResource(resource)
 		generators = append(generators, resource)
 	}
-	generators = append(generators, client)
+	generators = append(generators, cb.Build())
 
 	for _, g := range generators {
 		var code string

codegen/generator_test.go

@@ -199,7 +199,7 @@ func TestGenerateCodeFromFields(t *testing.T) {
 				tt.setupMockFiles(tt.fieldsDir)
 			}
 
-			err := generateCode(tt.fieldsDir, tt.outDir, noopCustomizer{})
+			err := generateCode(tt.fieldsDir, tt.outDir, CodeCustomizer{})
 
 			if tt.expectedError {
 				require.Error(t, err)

codegen/main.go

@@ -120,7 +120,7 @@ func generate(opts options) error {
 	if err != nil {
 		return fmt.Errorf("unable to create code customizer: %w", err)
 	}
-	if err = generateCode(structuresDir, outDir, customizer); err != nil {
+	if err = generateCode(structuresDir, outDir, *customizer); err != nil {
 		return fmt.Errorf("unable to generate resources code: %w", err)
 	}
 

codegen/resources_test.go

@@ -354,7 +354,7 @@ func TestBuildResourcesFromDownloadedFields(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 			a := assert.New(t)
-			resources, err := buildResourcesFromDownloadedFields(tc.dir, noopCustomizer{})
+			resources, err := buildResourcesFromDownloadedFields(tc.dir, CodeCustomizer{})
 			if tc.errorContains != "" {
 				require.ErrorContains(t, err, tc.errorContains)
 				a.Nil(resources)