ListMachines(ctx)

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

Hide file tree

Showing 5 changed files with 20 additions and 17 deletions.
diff --git a/cmd/crowdsec-cli/climachine/machines.go b/cmd/crowdsec-cli/climachine/machines.go
@@ -211,11 +211,11 @@ func (cli *cliMachines) listCSV(out io.Writer, machines ent.Machines) error {
 	return nil
 }
 
-func (cli *cliMachines) List(out io.Writer, db *database.Client) error {
+func (cli *cliMachines) List(ctx context.Context, out io.Writer, db *database.Client) error {
 	// XXX: must use the provided db object, the one in the struct might be nil
 	// (calling List directly skips the PersistentPreRunE)
 
-	machines, err := db.ListMachines()
+	machines, err := db.ListMachines(ctx)
 	if err != nil {
 		return fmt.Errorf("unable to list machines: %w", err)
 	}
@@ -252,8 +252,8 @@ func (cli *cliMachines) newListCmd() *cobra.Command {
 		Example:           `cscli machines list`,
 		Args:              cobra.NoArgs,
 		DisableAutoGenTag: true,
-		RunE: func(_ *cobra.Command, _ []string) error {
-			return cli.List(color.Output, cli.db)
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			return cli.List(cmd.Context(), color.Output, cli.db)
 		},
 	}
 
@@ -400,6 +400,7 @@ func (cli *cliMachines) validMachineID(cmd *cobra.Command, args []string, toComp
 	var err error
 
 	cfg := cli.cfg()
+	ctx := cmd.Context()
 
 	// need to load config and db because PersistentPreRunE is not called for completions
 
@@ -408,13 +409,13 @@ func (cli *cliMachines) validMachineID(cmd *cobra.Command, args []string, toComp
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}
 
-	cli.db, err = require.DBClient(cmd.Context(), cfg.DbConfig)
+	cli.db, err = require.DBClient(ctx, cfg.DbConfig)
 	if err != nil {
 		cobra.CompError("unable to list machines " + err.Error())
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}
 
-	machines, err := cli.db.ListMachines()
+	machines, err := cli.db.ListMachines(ctx)
 	if err != nil {
 		cobra.CompError("unable to list machines " + err.Error())
 		return nil, cobra.ShellCompDirectiveNoFileComp

diff --git a/cmd/crowdsec-cli/clisupport/support.go b/cmd/crowdsec-cli/clisupport/support.go
@@ -210,7 +210,7 @@ func (cli *cliSupport) dumpBouncers(ctx context.Context, zw *zip.Writer, db *dat
 	return nil
 }
 
-func (cli *cliSupport) dumpAgents(zw *zip.Writer, db *database.Client) error {
+func (cli *cliSupport) dumpAgents(ctx context.Context, zw *zip.Writer, db *database.Client) error {
 	log.Info("Collecting agents")
 
 	if db == nil {
@@ -220,7 +220,7 @@ func (cli *cliSupport) dumpAgents(zw *zip.Writer, db *database.Client) error {
 	out := new(bytes.Buffer)
 	cm := climachine.New(cli.cfg)
 
-	if err := cm.List(out, db); err != nil {
+	if err := cm.List(ctx, out, db); err != nil {
 		return err
 	}
 
@@ -529,7 +529,7 @@ func (cli *cliSupport) dump(ctx context.Context, outFile string) error {
 		log.Warnf("could not collect bouncers information: %s", err)
 	}
 
-	if err = cli.dumpAgents(zipWriter, db); err != nil {
+	if err = cli.dumpAgents(ctx, zipWriter, db); err != nil {
 		log.Warnf("could not collect agents information: %s", err)
 	}
 

diff --git a/pkg/apiserver/apic.go b/pkg/apiserver/apic.go
@@ -85,7 +85,9 @@ func randomDuration(d time.Duration, delta time.Duration) time.Duration {
 func (a *apic) FetchScenariosListFromDB() ([]string, error) {
 	scenarios := make([]string, 0)
 
-	machines, err := a.dbClient.ListMachines()
+	ctx := context.TODO()
+
+	machines, err := a.dbClient.ListMachines(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("while listing machines: %w", err)
 	}

diff --git a/pkg/apiserver/apic_metrics.go b/pkg/apiserver/apic_metrics.go
@@ -27,7 +27,7 @@ func (a *apic) GetUsageMetrics(ctx context.Context) (*models.AllMetrics, []int,
 	allMetrics := &models.AllMetrics{}
 	metricsIds := make([]int, 0)
 
-	lps, err := a.dbClient.ListMachines()
+	lps, err := a.dbClient.ListMachines(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -186,7 +186,7 @@ func (a *apic) MarkUsageMetricsAsSent(ctx context.Context, ids []int) error {
 }
 
 func (a *apic) GetMetrics(ctx context.Context) (*models.Metrics, error) {
-	machines, err := a.dbClient.ListMachines()
+	machines, err := a.dbClient.ListMachines(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -230,8 +230,8 @@ func (a *apic) GetMetrics(ctx context.Context) (*models.Metrics, error) {
 	}, nil
 }
 
-func (a *apic) fetchMachineIDs() ([]string, error) {
-	machines, err := a.dbClient.ListMachines()
+func (a *apic) fetchMachineIDs(ctx context.Context) ([]string, error) {
+	machines, err := a.dbClient.ListMachines(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -277,7 +277,7 @@ func (a *apic) SendMetrics(stop chan (bool)) {
 	machineIDs := []string{}
 
 	reloadMachineIDs := func() {
-		ids, err := a.fetchMachineIDs()
+		ids, err := a.fetchMachineIDs(ctx)
 		if err != nil {
 			log.Debugf("unable to get machines (%s), will retry", err)
 

diff --git a/pkg/database/machines.go b/pkg/database/machines.go
@@ -135,8 +135,8 @@ func (c *Client) QueryMachineByID(ctx context.Context, machineID string) (*ent.M
 	return machine, nil
 }
 
-func (c *Client) ListMachines() ([]*ent.Machine, error) {
-	machines, err := c.Ent.Machine.Query().All(c.CTX)
+func (c *Client) ListMachines(ctx context.Context) ([]*ent.Machine, error) {
+	machines, err := c.Ent.Machine.Query().All(ctx)
 	if err != nil {
 		return nil, errors.Wrapf(QueryFail, "listing machines: %s", err)
 	}