Report custom labels set by agent admins back (#4141)

server/grpc/rpc.go

@@ -441,7 +441,7 @@ func (s *RPC) Log(c context.Context, stepUUID string, rpcLogEntries []*rpc.LogEn
 	return nil
 }
 
-func (s *RPC) RegisterAgent(ctx context.Context, platform, backend, version string, capacity int32) (int64, error) {
+func (s *RPC) RegisterAgent(ctx context.Context, info rpc.AgentInfo) (int64, error) {
 	agent, err := s.getAgentFromContext(ctx)
 	if err != nil {
 		return -1, err
@@ -453,10 +453,11 @@ func (s *RPC) RegisterAgent(ctx context.Context, platform, backend, version stri
 		}
 	}
 
-	agent.Backend = backend
-	agent.Platform = platform
-	agent.Capacity = capacity
-	agent.Version = version
+	agent.Backend = info.Backend
+	agent.Platform = info.Platform
+	agent.Capacity = int32(info.Capacity)
+	agent.Version = info.Version
+	agent.CustomLabels = info.CustomLabels
 
 	err = s.store.AgentUpdate(agent)
 	if err != nil {

server/grpc/rpc_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc/metadata"
 
+	"go.woodpecker-ci.org/woodpecker/v2/pipeline/rpc"
 	"go.woodpecker-ci.org/woodpecker/v2/server/model"
 	mocks_store "go.woodpecker-ci.org/woodpecker/v2/server/store/mocks"
 )
@@ -52,15 +53,19 @@ func TestRegisterAgent(t *testing.T) {
 
 			store.On("AgentFind", int64(1337)).Once().Return(storeAgent, nil)
 			store.On("AgentUpdate", &updatedAgent).Once().Return(nil)
-			rpc := RPC{
+			grpc := RPC{
 				store: store,
 			}
 			ctx := metadata.NewIncomingContext(
 				context.Background(),
 				metadata.Pairs("hostname", "hostname", "agent_id", "1337"),
 			)
-			capacity := int32(2)
-			agentID, err := rpc.RegisterAgent(ctx, "platform", "backend", "version", capacity)
+			agentID, err := grpc.RegisterAgent(ctx, rpc.AgentInfo{
+				Version:  "version",
+				Platform: "platform",
+				Backend:  "backend",
+				Capacity: 2,
+			})
 			if !assert.NoError(t, err) {
 				return
 			}
@@ -92,15 +97,19 @@ func TestRegisterAgent(t *testing.T) {
 
 			store.On("AgentFind", int64(1337)).Once().Return(storeAgent, nil)
 			store.On("AgentUpdate", &updatedAgent).Once().Return(nil)
-			rpc := RPC{
+			grpc := RPC{
 				store: store,
 			}
 			ctx := metadata.NewIncomingContext(
 				context.Background(),
 				metadata.Pairs("hostname", "newHostname", "agent_id", "1337"),
 			)
-			capacity := int32(2)
-			agentID, err := rpc.RegisterAgent(ctx, "platform", "backend", "version", capacity)
+			agentID, err := grpc.RegisterAgent(ctx, rpc.AgentInfo{
+				Version:  "version",
+				Platform: "platform",
+				Backend:  "backend",
+				Capacity: 2,
+			})
 			if !assert.NoError(t, err) {
 				return
 			}

server/grpc/server.go

@@ -172,7 +172,14 @@ func (s *WoodpeckerServer) Log(c context.Context, req *proto.LogRequest) (*proto
 
 func (s *WoodpeckerServer) RegisterAgent(c context.Context, req *proto.RegisterAgentRequest) (*proto.RegisterAgentResponse, error) {
 	res := new(proto.RegisterAgentResponse)
-	agentID, err := s.peer.RegisterAgent(c, req.GetPlatform(), req.GetBackend(), req.GetVersion(), req.GetCapacity())
+	agentInfo := req.GetInfo()
+	agentID, err := s.peer.RegisterAgent(c, rpc.AgentInfo{
+		Version:      agentInfo.GetVersion(),
+		Platform:     agentInfo.GetPlatform(),
+		Backend:      agentInfo.GetBackend(),
+		Capacity:     int(agentInfo.GetCapacity()),
+		CustomLabels: agentInfo.GetCustomLabels(),
+	})
 	res.AgentId = agentID
 	return res, err
 }

server/model/agent.go

@@ -22,19 +22,20 @@ import (
 )
 
 type Agent struct {
-	ID          int64  `json:"id"            xorm:"pk autoincr 'id'"`
-	Created     int64  `json:"created"       xorm:"created"`
-	Updated     int64  `json:"updated"       xorm:"updated"`
-	Name        string `json:"name"          xorm:"name"`
-	OwnerID     int64  `json:"owner_id"      xorm:"'owner_id'"`
-	Token       string `json:"token"         xorm:"token"`
-	LastContact int64  `json:"last_contact"  xorm:"last_contact"`
-	LastWork    int64  `json:"last_work"     xorm:"last_work"` // last time the agent did something, this value is used to determine if the agent is still doing work used by the autoscaler
-	Platform    string `json:"platform"      xorm:"VARCHAR(100) 'platform'"`
-	Backend     string `json:"backend"       xorm:"VARCHAR(100) 'backend'"`
-	Capacity    int32  `json:"capacity"      xorm:"capacity"`
-	Version     string `json:"version"       xorm:"'version'"`
-	NoSchedule  bool   `json:"no_schedule"   xorm:"no_schedule"`
+	ID           int64             `json:"id"            xorm:"pk autoincr 'id'"`
+	Created      int64             `json:"created"       xorm:"created"`
+	Updated      int64             `json:"updated"       xorm:"updated"`
+	Name         string            `json:"name"          xorm:"name"`
+	OwnerID      int64             `json:"owner_id"      xorm:"'owner_id'"`
+	Token        string            `json:"token"         xorm:"token"`
+	LastContact  int64             `json:"last_contact"  xorm:"last_contact"`
+	LastWork     int64             `json:"last_work"     xorm:"last_work"` // last time the agent did something, this value is used to determine if the agent is still doing work used by the autoscaler
+	Platform     string            `json:"platform"      xorm:"VARCHAR(100) 'platform'"`
+	Backend      string            `json:"backend"       xorm:"VARCHAR(100) 'backend'"`
+	Capacity     int32             `json:"capacity"      xorm:"capacity"`
+	Version      string            `json:"version"       xorm:"'version'"`
+	NoSchedule   bool              `json:"no_schedule"   xorm:"no_schedule"`
+	CustomLabels map[string]string `json:"custom_labels" xorm:"JSON 'custom_labels'"`
 	// OrgID is counted as unset if set to -1, this is done to ensure a new(Agent) still enforce the OrgID check by default
 	OrgID int64 `json:"org_id"        xorm:"INDEX 'org_id'"`
 } //	@name Agent

server/store/datastore/migration/016_add_custom_labels_to_agent.go

@@ -0,0 +1,41 @@
+// Copyright 2024 Woodpecker Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package migration
+
+import (
+	"fmt"
+
+	"src.techknowlogick.com/xormigrate"
+	"xorm.io/xorm"
+)
+
+type agentV016 struct {
+	ID           int64             `xorm:"pk autoincr 'id'"`
+	CustomLabels map[string]string `xorm:"JSON 'custom_labels'"`
+}
+
+func (agentV016) TableName() string {
+	return "agents"
+}
+
+var addCustomLabelsToAgent = xormigrate.Migration{
+	ID: "add-custom-labels-to-agent",
+	MigrateSession: func(sess *xorm.Session) (err error) {
+		if err := sess.Sync(new(agentV016)); err != nil {
+			return fmt.Errorf("sync models failed: %w", err)
+		}
+		return nil
+	},
+}

server/store/datastore/migration/migration.go

@@ -44,6 +44,7 @@ var migrationTasks = []*xormigrate.Migration{
 	&fixV31Registries,
 	&removeOldMigrationsOfV1,
 	&addOrgAgents,
+	&addCustomLabelsToAgent,
 }
 
 var allBeans = []any{