Skip to content

Commit d955524

Browse files
PaamickyPaamicky
committed
Add workload detection for postgreSQL
1 parent 706474e commit d955524

7 files changed

Lines changed: 406 additions & 0 deletions

File tree

‎cmd/workload-discovery/discovery.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/aws/amazon-cloudwatch-agent/internal/detector/filter"
2222
"github.com/aws/amazon-cloudwatch-agent/internal/detector/java"
2323
"github.com/aws/amazon-cloudwatch-agent/internal/detector/nvidia"
24+
"github.com/aws/amazon-cloudwatch-agent/internal/detector/postgresql"
2425
"github.com/aws/amazon-cloudwatch-agent/internal/detector/util"
2526
"github.com/aws/amazon-cloudwatch-agent/tool/paths"
2627
)
@@ -47,6 +48,7 @@ func NewDiscoverer(cfg Config, logger *slog.Logger) *Discoverer {
4748
logger: logger,
4849
processDetectors: []detector.ProcessDetector{
4950
java.NewDetector(logger, filters.Process.Name),
51+
postgresql.NewDetector(logger),
5052
},
5153
deviceDetectors: []detector.DeviceDetector{
5254
nvidia.NewDetector(logger),

‎internal/detector/metadata.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ const (
3030
CategoryKafkaBroker Category = "KAFKA/BROKER"
3131
CategoryKafkaClient Category = "KAFKA/CLIENT"
3232
CategoryNvidiaGPU Category = "NVIDIA_GPU"
33+
CategoryPostgreSQL Category = "POSTGRESQL"
3334
)
3435

3536
// Status represents whether the resource requires more actions before telemetry is available.

‎internal/detector/postgresql/README.md‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# PostgreSQL Process Detector
2+
3+
Detects PostgreSQL database server processes running on the system.
4+
5+
## Overview
6+
7+
The PostgreSQL detector identifies PostgreSQL server instances by checking for the `postgres` executable name. It is used by the `workload-discovery` command to automatically discover PostgreSQL workloads.
8+
9+
## Detection Method
10+
11+
The detector examines the executable path of each process and checks if the base name is `postgres`.
12+
13+
## Status Results
14+
- `READY`: PostgreSQL process detected with a port (explicit via `-p` flag or `PGPORT`, otherwise defaults to 5432).
15+
16+
## Sample Metadata Result
17+
```json
18+
{
19+
"categories": ["POSTGRESQL"],
20+
"name": "postgresql",
21+
"status": "READY",
22+
"telemetryPort": 5432
23+
}
24+
```

‎internal/detector/postgresql/extract/port.go‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: MIT
3+
4+
package extract
5+
6+
import (
7+
"context"
8+
"strconv"
9+
"strings"
10+
11+
"github.com/aws/amazon-cloudwatch-agent/internal/detector"
12+
"github.com/aws/amazon-cloudwatch-agent/internal/detector/util"
13+
)
14+
15+
const (
16+
portFlag = "-p"
17+
portEnvVar = "PGPORT"
18+
defaultPostgresPort = 5432
19+
)
20+
21+
type portExtractor struct {
22+
subExtractors []detector.PortExtractor
23+
}
24+
25+
// NewPortExtractor creates a port extractor that attempts to find the PostgreSQL port
26+
// from command line arguments (-p flag) or environment variables (PGPORT).
27+
// Falls back to the default PostgreSQL port 5432.
28+
func NewPortExtractor() detector.PortExtractor {
29+
return &portExtractor{
30+
subExtractors: []detector.PortExtractor{
31+
&cmdlinePortExtractor{},
32+
&envPortExtractor{},
33+
},
34+
}
35+
}
36+
37+
func (e *portExtractor) Extract(ctx context.Context, process detector.Process) (int, error) {
38+
for _, sub := range e.subExtractors {
39+
port, err := sub.Extract(ctx, process)
40+
if err == nil {
41+
return port, nil
42+
}
43+
}
44+
return defaultPostgresPort, nil
45+
}
46+
47+
// cmdlinePortExtractor extracts port from -p flag
48+
type cmdlinePortExtractor struct{}
49+
50+
func (e *cmdlinePortExtractor) Extract(ctx context.Context, process detector.Process) (int, error) {
51+
args, err := process.CmdlineSliceWithContext(ctx)
52+
if err != nil {
53+
return 0, err
54+
}
55+
56+
for i, arg := range args {
57+
if arg == portFlag && i+1 < len(args) {
58+
port, err := strconv.Atoi(args[i+1])
59+
if err == nil && util.IsValidPort(port) {
60+
return port, nil
61+
}
62+
}
63+
if strings.HasPrefix(arg, portFlag) && len(arg) > len(portFlag) {
64+
port, err := strconv.Atoi(arg[len(portFlag):])
65+
if err == nil && util.IsValidPort(port) {
66+
return port, nil
67+
}
68+
}
69+
}
70+
71+
return 0, detector.ErrExtractPort
72+
}
73+
74+
// envPortExtractor extracts port from PGPORT environment variable
75+
type envPortExtractor struct{}
76+
77+
func (e *envPortExtractor) Extract(ctx context.Context, process detector.Process) (int, error) {
78+
env, err := process.EnvironWithContext(ctx)
79+
if err != nil {
80+
return 0, err
81+
}
82+
83+
for _, entry := range env {
84+
parts := strings.SplitN(entry, "=", 2)
85+
if len(parts) == 2 && parts[0] == portEnvVar {
86+
port, err := strconv.Atoi(strings.TrimSpace(parts[1]))
87+
if err == nil && util.IsValidPort(port) {
88+
return port, nil
89+
}
90+
}
91+
}
92+
93+
return 0, detector.ErrExtractPort
94+
}

‎internal/detector/postgresql/extract/port_test.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: MIT
3+
4+
package extract
5+
6+
import (
7+
"context"
8+
"testing"
9+
10+
"github.com/stretchr/testify/assert"
11+
"github.com/stretchr/testify/require"
12+
13+
"github.com/aws/amazon-cloudwatch-agent/internal/detector/detectortest"
14+
)
15+
16+
func TestPortExtractor(t *testing.T) {
17+
ctx := context.Background()
18+
extractor := NewPortExtractor()
19+
20+
tests := map[string]struct {
21+
setup func(*detectortest.MockProcess)
22+
wantPort int
23+
}{
24+
"Success/PortFromFlag": {
25+
setup: func(mp *detectortest.MockProcess) {
26+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres", "-p", "5433"}, nil)
27+
},
28+
wantPort: 5433,
29+
},
30+
"Success/PortFromFlagNoSpace": {
31+
setup: func(mp *detectortest.MockProcess) {
32+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres", "-p5434"}, nil)
33+
},
34+
wantPort: 5434,
35+
},
36+
"Success/PortFromEnv": {
37+
setup: func(mp *detectortest.MockProcess) {
38+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres"}, nil)
39+
mp.On("EnvironWithContext", ctx).Return([]string{"PATH=/usr/bin", "PGPORT=5435"}, nil)
40+
},
41+
wantPort: 5435,
42+
},
43+
// cmdline is tried first; when it finds a port, env is never called
44+
"Success/CmdlineTakesPrecedence": {
45+
setup: func(mp *detectortest.MockProcess) {
46+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres", "-p", "5433"}, nil)
47+
},
48+
wantPort: 5433,
49+
},
50+
"Success/DefaultPort": {
51+
setup: func(mp *detectortest.MockProcess) {
52+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres"}, nil)
53+
mp.On("EnvironWithContext", ctx).Return([]string{"PATH=/usr/bin"}, nil)
54+
},
55+
wantPort: 5432,
56+
},
57+
"Success/DefaultPortWithOtherFlags": {
58+
setup: func(mp *detectortest.MockProcess) {
59+
mp.On("CmdlineSliceWithContext", ctx).Return([]string{"postgres", "-D", "/var/lib/postgresql/data"}, nil)
60+
mp.On("EnvironWithContext", ctx).Return([]string{}, nil)
61+
},
62+
wantPort: 5432,
63+
},
64+
"Success/DefaultOnAllSourcesFail": {
65+
setup: func(mp *detectortest.MockProcess) {
66+
mp.On("CmdlineSliceWithContext", ctx).Return(nil, assert.AnError)
67+
mp.On("EnvironWithContext", ctx).Return(nil, assert.AnError)
68+
},
69+
wantPort: 5432,
70+
},
71+
}
72+
73+
for name, tt := range tests {
74+
t.Run(name, func(t *testing.T) {
75+
mp := new(detectortest.MockProcess)
76+
tt.setup(mp)
77+
78+
port, err := extractor.Extract(ctx, mp)
79+
80+
require.NoError(t, err)
81+
assert.Equal(t, tt.wantPort, port)
82+
mp.AssertExpectations(t)
83+
})
84+
}
85+
}

‎internal/detector/postgresql/postgresql.go‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: MIT
3+
4+
package postgresql
5+
6+
import (
7+
"context"
8+
"log/slog"
9+
"strings"
10+
11+
"github.com/aws/amazon-cloudwatch-agent/internal/detector"
12+
"github.com/aws/amazon-cloudwatch-agent/internal/detector/postgresql/extract"
13+
"github.com/aws/amazon-cloudwatch-agent/internal/detector/util"
14+
)
15+
16+
const (
17+
exeName = "postgres"
18+
)
19+
20+
type postgresqlDetector struct {
21+
logger *slog.Logger
22+
portExtractor detector.PortExtractor
23+
}
24+
25+
var _ detector.ProcessDetector = (*postgresqlDetector)(nil)
26+
27+
// NewDetector creates a new process detector that identifies PostgreSQL processes.
28+
func NewDetector(logger *slog.Logger) detector.ProcessDetector {
29+
return &postgresqlDetector{
30+
logger: logger,
31+
portExtractor: extract.NewPortExtractor(),
32+
}
33+
}
34+
35+
// Detect identifies PostgreSQL processes and returns metadata.
36+
// Only detects the main postgres process, not worker processes.
37+
func (d *postgresqlDetector) Detect(ctx context.Context, process detector.Process) (*detector.Metadata, error) {
38+
exe, err := process.ExeWithContext(ctx)
39+
if err != nil {
40+
return nil, err
41+
}
42+
43+
base := util.BaseExe(exe)
44+
if base != exeName {
45+
return nil, detector.ErrIncompatibleDetector
46+
}
47+
48+
// Check if this is the main postgres process or a worker
49+
args, err := process.CmdlineSliceWithContext(ctx)
50+
if err != nil {
51+
return nil, err
52+
}
53+
54+
if len(args) > 0 && strings.HasPrefix(strings.TrimSpace(args[0]), exeName+":") {
55+
return nil, detector.ErrIncompatibleDetector
56+
}
57+
58+
d.logger.Debug("PostgreSQL process detected", "pid", process.PID())
59+
60+
md := &detector.Metadata{
61+
Name: "postgresql",
62+
Categories: []detector.Category{detector.CategoryPostgreSQL},
63+
}
64+
65+
port, _ := d.portExtractor.Extract(ctx, process)
66+
67+
md.Status = detector.StatusReady
68+
md.TelemetryPort = port
69+
70+
return md, nil
71+
}

0 commit comments

Comments
 (0)