QueryMachineByID(ctx...)

mmetc · Sep 18, 2024 · a2358e7 · a2358e7
1 parent 1d54af3
commit a2358e7
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git a/cmd/crowdsec-cli/climachine/machines.go b/cmd/crowdsec-cli/climachine/machines.go
@@ -691,9 +691,11 @@ func (cli *cliMachines) newInspectCmd() *cobra.Command {
 		Args:              cobra.ExactArgs(1),
 		DisableAutoGenTag: true,
 		ValidArgsFunction: cli.validMachineID,
-		RunE: func(_ *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
 			machineID := args[0]
-			machine, err := cli.db.QueryMachineByID(machineID)
+
+			machine, err := cli.db.QueryMachineByID(ctx, machineID)
 			if err != nil {
 				return fmt.Errorf("unable to read machine data '%s': %w", machineID, err)
 			}

diff --git a/pkg/apiserver/usage_metrics_test.go b/pkg/apiserver/usage_metrics_test.go
@@ -199,7 +199,7 @@ func TestLPMetrics(t *testing.T) {
 			assert.Equal(t, tt.expectedStatusCode, w.Code)
 			assert.Contains(t, w.Body.String(), tt.expectedResponse)
 
-			machine, _ := dbClient.QueryMachineByID("test")
+			machine, _ := dbClient.QueryMachineByID(ctx, "test")
 			metrics, _ := dbClient.GetLPUsageMetricsByMachineID(ctx, "test")
 
 			assert.Len(t, metrics, tt.expectedMetricsCount)

diff --git a/pkg/database/alerts.go b/pkg/database/alerts.go
@@ -687,8 +687,10 @@ func (c *Client) CreateAlert(machineID string, alertList []*models.Alert) ([]str
 		err   error
 	)
 
+	ctx := context.TODO()
+
 	if machineID != "" {
-		owner, err = c.QueryMachineByID(machineID)
+		owner, err = c.QueryMachineByID(ctx, machineID)
 		if err != nil {
 			if !errors.Is(err, UserNotExists) {
 				return nil, fmt.Errorf("machine '%s': %w", machineID, err)

diff --git a/pkg/database/machines.go b/pkg/database/machines.go
@@ -95,7 +95,7 @@ func (c *Client) CreateMachine(ctx context.Context, machineID *string, password
 				return nil, errors.Wrapf(UpdateFail, "machine '%s'", *machineID)
 			}
 
-			machine, err := c.QueryMachineByID(*machineID)
+			machine, err := c.QueryMachineByID(ctx, *machineID)
 			if err != nil {
 				return nil, errors.Wrapf(QueryFail, "machine '%s': %s", *machineID, err)
 			}
@@ -122,11 +122,11 @@ func (c *Client) CreateMachine(ctx context.Context, machineID *string, password
 	return machine, nil
 }
 
-func (c *Client) QueryMachineByID(machineID string) (*ent.Machine, error) {
+func (c *Client) QueryMachineByID(ctx context.Context, machineID string) (*ent.Machine, error) {
 	machine, err := c.Ent.Machine.
 		Query().
 		Where(machine.MachineIdEQ(machineID)).
-		Only(c.CTX)
+		Only(ctx)
 	if err != nil {
 		c.Log.Warningf("QueryMachineByID : %s", err)
 		return &ent.Machine{}, errors.Wrapf(UserNotExists, "user '%s'", machineID)