Skip to content

Commit

Permalink
context propagation: pkg/database/metrics
Browse files Browse the repository at this point in the history
  • Loading branch information
mmetc committed Sep 18, 2024
1 parent c50fb39 commit 15bb09a
Show file tree
Hide file tree
Showing 3 changed files with 17 additions and 13 deletions.
10 changes: 5 additions & 5 deletions pkg/apiserver/apic_metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ func (a *apic) GetUsageMetrics(ctx context.Context) (*models.AllMetrics, []int,
}

for _, bouncer := range bouncers {
dbMetrics, err := a.dbClient.GetBouncerUsageMetricsByName(bouncer.Name)
dbMetrics, err := a.dbClient.GetBouncerUsageMetricsByName(ctx, bouncer.Name)
if err != nil {
log.Errorf("unable to get bouncer usage metrics: %s", err)
continue
Expand Down Expand Up @@ -81,7 +81,7 @@ func (a *apic) GetUsageMetrics(ctx context.Context) (*models.AllMetrics, []int,
}

for _, lp := range lps {
dbMetrics, err := a.dbClient.GetLPUsageMetricsByMachineID(lp.MachineId)
dbMetrics, err := a.dbClient.GetLPUsageMetricsByMachineID(ctx, lp.MachineId)
if err != nil {
log.Errorf("unable to get LP usage metrics: %s", err)
continue
Expand Down Expand Up @@ -181,8 +181,8 @@ func (a *apic) GetUsageMetrics(ctx context.Context) (*models.AllMetrics, []int,
return allMetrics, metricsIds, nil
}

func (a *apic) MarkUsageMetricsAsSent(ids []int) error {
return a.dbClient.MarkUsageMetricsAsSent(ids)
func (a *apic) MarkUsageMetricsAsSent(ctx context.Context, ids []int) error {
return a.dbClient.MarkUsageMetricsAsSent(ctx, ids)
}

func (a *apic) GetMetrics(ctx context.Context) (*models.Metrics, error) {
Expand Down Expand Up @@ -379,7 +379,7 @@ func (a *apic) SendUsageMetrics() {
}
}

err = a.MarkUsageMetricsAsSent(metricsId)
err = a.MarkUsageMetricsAsSent(ctx, metricsId)
if err != nil {
log.Errorf("unable to mark usage metrics as sent: %s", err)
continue
Expand Down
8 changes: 6 additions & 2 deletions pkg/apiserver/usage_metrics_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ import (
)

func TestLPMetrics(t *testing.T) {
ctx := context.Background()

tests := []struct {
name string
body string
Expand Down Expand Up @@ -198,7 +200,7 @@ func TestLPMetrics(t *testing.T) {
assert.Contains(t, w.Body.String(), tt.expectedResponse)

machine, _ := dbClient.QueryMachineByID("test")
metrics, _ := dbClient.GetLPUsageMetricsByMachineID("test")
metrics, _ := dbClient.GetLPUsageMetricsByMachineID(ctx, "test")

assert.Len(t, metrics, tt.expectedMetricsCount)
assert.Equal(t, tt.expectedOSName, machine.Osname)
Expand All @@ -214,6 +216,8 @@ func TestLPMetrics(t *testing.T) {
}

func TestRCMetrics(t *testing.T) {
ctx := context.Background()

tests := []struct {
name string
body string
Expand Down Expand Up @@ -368,7 +372,7 @@ func TestRCMetrics(t *testing.T) {
assert.Contains(t, w.Body.String(), tt.expectedResponse)

bouncer, _ := dbClient.SelectBouncerByName("test")
metrics, _ := dbClient.GetBouncerUsageMetricsByName("test")
metrics, _ := dbClient.GetBouncerUsageMetricsByName(ctx, "test")

assert.Len(t, metrics, tt.expectedMetricsCount)
assert.Equal(t, tt.expectedOSName, bouncer.Osname)
Expand Down
12 changes: 6 additions & 6 deletions pkg/database/metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,14 @@ func (c *Client) CreateMetric(ctx context.Context, generatedType metric.Generate
return metric, nil
}

func (c *Client) GetLPUsageMetricsByMachineID(machineId string) ([]*ent.Metric, error) {
func (c *Client) GetLPUsageMetricsByMachineID(ctx context.Context, machineId string) ([]*ent.Metric, error) {
metrics, err := c.Ent.Metric.Query().
Where(
metric.GeneratedTypeEQ(metric.GeneratedTypeLP),
metric.GeneratedByEQ(machineId),
metric.PushedAtIsNil(),
).
All(c.CTX)
All(ctx)
if err != nil {
c.Log.Warningf("GetLPUsageMetricsByOrigin: %s", err)
return nil, fmt.Errorf("getting LP usage metrics by origin %s: %w", machineId, err)
Expand All @@ -41,14 +41,14 @@ func (c *Client) GetLPUsageMetricsByMachineID(machineId string) ([]*ent.Metric,
return metrics, nil
}

func (c *Client) GetBouncerUsageMetricsByName(bouncerName string) ([]*ent.Metric, error) {
func (c *Client) GetBouncerUsageMetricsByName(ctx context.Context, bouncerName string) ([]*ent.Metric, error) {
metrics, err := c.Ent.Metric.Query().
Where(
metric.GeneratedTypeEQ(metric.GeneratedTypeRC),
metric.GeneratedByEQ(bouncerName),
metric.PushedAtIsNil(),
).
All(c.CTX)
All(ctx)
if err != nil {
c.Log.Warningf("GetBouncerUsageMetricsByName: %s", err)
return nil, fmt.Errorf("getting bouncer usage metrics by name %s: %w", bouncerName, err)
Expand All @@ -57,11 +57,11 @@ func (c *Client) GetBouncerUsageMetricsByName(bouncerName string) ([]*ent.Metric
return metrics, nil
}

func (c *Client) MarkUsageMetricsAsSent(ids []int) error {
func (c *Client) MarkUsageMetricsAsSent(ctx context.Context, ids []int) error {
_, err := c.Ent.Metric.Update().
Where(metric.IDIn(ids...)).
SetPushedAt(time.Now().UTC()).
Save(c.CTX)
Save(ctx)
if err != nil {
c.Log.Warningf("MarkUsageMetricsAsSent: %s", err)
return fmt.Errorf("marking usage metrics as sent: %w", err)
Expand Down

0 comments on commit 15bb09a

Please sign in to comment.