1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
|
package server
import (
"context"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/agent_tracker"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/agent_tracker/rpc"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modshared"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)
type server struct {
rpc.UnimplementedAgentTrackerServer
agentQuerier agent_tracker.Querier
}
func (s *server) GetConnectedAgentsByProjectIDs(ctx context.Context, req *rpc.GetConnectedAgentsByProjectIDsRequest) (*rpc.GetConnectedAgentsByProjectIDsResponse, error) {
rpcAPI := modshared.RPCAPIFromContext(ctx)
log := rpcAPI.Log()
var infos agent_tracker.ConnectedAgentInfoCollector
for _, projectID := range req.ProjectIds {
err := s.agentQuerier.GetConnectionsByProjectID(ctx, projectID, infos.Collect)
if err != nil {
rpcAPI.HandleProcessingError(log, modshared.NoAgentID, "GetConnectionsByProjectID() failed", err)
return nil, status.Error(codes.Unavailable, "GetConnectionsByProjectID() failed")
}
}
return &rpc.GetConnectedAgentsByProjectIDsResponse{
Agents: infos,
}, nil
}
// This method has been deprecated and will be removed in favor of GetConnectedAgentsByProjectIDs.
func (s *server) GetConnectedAgentsByProjectIds(ctx context.Context, req *rpc.GetConnectedAgentsByProjectIDsRequest) (*rpc.GetConnectedAgentsByProjectIDsResponse, error) {
return s.GetConnectedAgentsByProjectIDs(ctx, req)
}
func (s *server) GetConnectedAgentsByAgentIDs(ctx context.Context, req *rpc.GetConnectedAgentsByAgentIDsRequest) (*rpc.GetConnectedAgentsByAgentIDsResponse, error) {
rpcAPI := modshared.RPCAPIFromContext(ctx)
log := rpcAPI.Log()
var infos agent_tracker.ConnectedAgentInfoCollector
for _, agentID := range req.AgentIds {
err := s.agentQuerier.GetConnectionsByAgentID(ctx, agentID, infos.Collect)
if err != nil {
rpcAPI.HandleProcessingError(log, modshared.NoAgentID, "GetConnectionsByAgentID() failed", err)
return nil, status.Error(codes.Unavailable, "GetConnectionsByAgentID() failed")
}
}
return &rpc.GetConnectedAgentsByAgentIDsResponse{
Agents: infos,
}, nil
}
// This method has been deprecated and will be removed in favor of GetConnectedAgentsByAgentIDs.
func (s *server) GetConnectedAgentsByAgentIds(ctx context.Context, req *rpc.GetConnectedAgentsByAgentIDsRequest) (*rpc.GetConnectedAgentsByAgentIDsResponse, error) {
return s.GetConnectedAgentsByAgentIDs(ctx, req)
}
func (s *server) GetConnectedAgents(ctx context.Context, req *rpc.GetConnectedAgentsRequest) (*rpc.GetConnectedAgentsResponse, error) {
rpcAPI := modshared.RPCAPIFromContext(ctx)
log := rpcAPI.Log()
switch v := req.Request.(type) {
case *rpc.GetConnectedAgentsRequest_AgentId:
var infos agent_tracker.ConnectedAgentInfoCollector
err := s.agentQuerier.GetConnectionsByAgentID(ctx, v.AgentId, infos.Collect)
if err != nil {
rpcAPI.HandleProcessingError(log, modshared.NoAgentID, "GetConnectionsByAgentID() failed", err)
return nil, status.Error(codes.Unavailable, "GetConnectionsByAgentID() failed")
}
return &rpc.GetConnectedAgentsResponse{
Agents: infos,
}, nil
case *rpc.GetConnectedAgentsRequest_ProjectId:
var infos agent_tracker.ConnectedAgentInfoCollector
err := s.agentQuerier.GetConnectionsByProjectID(ctx, v.ProjectId, infos.Collect)
if err != nil {
rpcAPI.HandleProcessingError(log, modshared.NoAgentID, "GetConnectionsByProjectID() failed", err)
return nil, status.Error(codes.Unavailable, "GetConnectionsByProjectID() failed")
}
return &rpc.GetConnectedAgentsResponse{
Agents: infos,
}, nil
default:
// Should never happen
return nil, status.Errorf(codes.InvalidArgument, "unexpected field type: %T", req.Request)
}
}
func (s *server) CountAgentsByAgentVersions(ctx context.Context, _ *rpc.CountAgentsByAgentVersionsRequest) (*rpc.CountAgentsByAgentVersionsResponse, error) {
rpcAPI := modshared.RPCAPIFromContext(ctx)
log := rpcAPI.Log()
counts, err := s.agentQuerier.CountAgentsByAgentVersions(ctx)
if err != nil {
rpcAPI.HandleProcessingError(log, modshared.NoAgentID, "CountAgentsByAgentVersions() failed", err)
return nil, status.Error(codes.Unavailable, "CountAgentsByAgentVersions() failed")
}
return &rpc.CountAgentsByAgentVersionsResponse{
AgentVersions: counts,
}, nil
}
|
