Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

context propagation: pkg/database/bouncers #130

Open
wants to merge 16 commits into
base: master
Choose a base branch
from
Open

context propagation: pkg/database/bouncers #130

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
965e1c6
context propagation: explicit ctx parameter in unit tests
mmetc Sep 13, 2024
4f10b2d
context propagation: pass context to NewAPIC()
mmetc Sep 13, 2024
66c3f0f
context propagation: drop field S3Source.ctx, pass explicitly
mmetc Sep 13, 2024
77c269a
context propagation: pkg/database/flush
mmetc Sep 13, 2024
f8ac861
context propagation: bouncer list
mmetc Sep 13, 2024
aba6dfd
context propagation: pkg/database/config
mmetc Sep 18, 2024
37c0109
context propagation: pkg/database/metrics
mmetc Sep 18, 2024
1d54af3
CreateMachine(ctx...)
mmetc Sep 18, 2024
a2358e7
QueryMachineByID(ctx...)
mmetc Sep 18, 2024
ed0646e
ListMachines(ctx)
mmetc Sep 18, 2024
2654d2c
ValidateMachine(ctx...)
mmetc Sep 18, 2024
9d9fc23
QueryPendingMachine(ctx)
mmetc Sep 18, 2024
45873cf
DeleteWatcher(ctx...)
mmetc Sep 18, 2024
e261ae6
BulkDeleteWatchers(ctx...)
mmetc Sep 18, 2024
94f3abe
rest of machines.go
mmetc Sep 18, 2024
b3109de
context propagation: pkg/database/bouncers
mmetc Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions cmd/crowdsec-cli/clialert/alerts.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -575,15 +575,17 @@ func (cli *cliAlerts) newFlushCmd() *cobra.Command {
DisableAutoGenTag: true,
RunE: func(cmd *cobra.Command, _ []string) error {
cfg := cli.cfg()
ctx := cmd.Context()

if err := require.LAPI(cfg); err != nil {
return err
}
db, err := require.DBClient(cmd.Context(), cfg.DbConfig)
db, err := require.DBClient(ctx, cfg.DbConfig)
if err != nil {
return err
}
log.Info("Flushing alerts. !! This may take a long time !!")
err = db.FlushAlerts(maxAge, maxItems)
err = db.FlushAlerts(ctx, maxAge, maxItems)
if err != nil {
return fmt.Errorf("unable to flush alerts: %w", err)
}
Expand Down
40 changes: 21 additions & 19 deletions cmd/crowdsec-cli/clibouncer/bouncers.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package clibouncer

import (
"context"
"encoding/csv"
"encoding/json"
"errors"
Expand Down Expand Up @@ -159,11 +160,11 @@ func (cli *cliBouncers) listCSV(out io.Writer, bouncers ent.Bouncers) error {
return nil
}

func (cli *cliBouncers) List(out io.Writer, db *database.Client) error {
func (cli *cliBouncers) 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)

bouncers, err := db.ListBouncers()
bouncers, err := db.ListBouncers(ctx)
if err != nil {
return fmt.Errorf("unable to list bouncers: %w", err)
}
Expand Down Expand Up @@ -199,15 +200,15 @@ func (cli *cliBouncers) newListCmd() *cobra.Command {
Example: `cscli bouncers list`,
Args: cobra.ExactArgs(0),
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)
},
}

return cmd
}

func (cli *cliBouncers) add(bouncerName string, key string) error {
func (cli *cliBouncers) add(ctx context.Context, bouncerName string, key string) error {
var err error

keyLength := 32
Expand All @@ -219,7 +220,7 @@ func (cli *cliBouncers) add(bouncerName string, key string) error {
}
}

_, err = cli.db.CreateBouncer(bouncerName, "", middlewares.HashSHA512(key), types.ApiKeyAuthType)
_, err = cli.db.CreateBouncer(ctx, bouncerName, "", middlewares.HashSHA512(key), types.ApiKeyAuthType)
if err != nil {
return fmt.Errorf("unable to create bouncer: %w", err)
}
Expand Down Expand Up @@ -253,8 +254,8 @@ func (cli *cliBouncers) newAddCmd() *cobra.Command {
cscli bouncers add MyBouncerName --key <random-key>`,
Args: cobra.ExactArgs(1),
DisableAutoGenTag: true,
RunE: func(_ *cobra.Command, args []string) error {
return cli.add(args[0], key)
RunE: func(cmd *cobra.Command, args []string) error {
return cli.add(cmd.Context(), args[0], key)
},
}

Expand All @@ -271,6 +272,7 @@ func (cli *cliBouncers) validBouncerID(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

Expand All @@ -279,13 +281,13 @@ func (cli *cliBouncers) validBouncerID(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 bouncers " + err.Error())
return nil, cobra.ShellCompDirectiveNoFileComp
}

bouncers, err := cli.db.ListBouncers()
bouncers, err := cli.db.ListBouncers(ctx)
if err != nil {
cobra.CompError("unable to list bouncers " + err.Error())
return nil, cobra.ShellCompDirectiveNoFileComp
Expand All @@ -302,9 +304,9 @@ func (cli *cliBouncers) validBouncerID(cmd *cobra.Command, args []string, toComp
return ret, cobra.ShellCompDirectiveNoFileComp
}

func (cli *cliBouncers) delete(bouncers []string, ignoreMissing bool) error {
func (cli *cliBouncers) delete(ctx context.Context, bouncers []string, ignoreMissing bool) error {
for _, bouncerID := range bouncers {
if err := cli.db.DeleteBouncer(bouncerID); err != nil {
if err := cli.db.DeleteBouncer(ctx, bouncerID); err != nil {
var notFoundErr *database.BouncerNotFoundError
if ignoreMissing && errors.As(err, &notFoundErr) {
return nil
Expand All @@ -330,8 +332,8 @@ func (cli *cliBouncers) newDeleteCmd() *cobra.Command {
Aliases: []string{"remove"},
DisableAutoGenTag: true,
ValidArgsFunction: cli.validBouncerID,
RunE: func(_ *cobra.Command, args []string) error {
return cli.delete(args, ignoreMissing)
RunE: func(cmd *cobra.Command, args []string) error {
return cli.delete(cmd.Context(), args, ignoreMissing)
},
}

Expand All @@ -341,7 +343,7 @@ func (cli *cliBouncers) newDeleteCmd() *cobra.Command {
return cmd
}

func (cli *cliBouncers) prune(duration time.Duration, force bool) error {
func (cli *cliBouncers) prune(ctx context.Context, duration time.Duration, force bool) error {
if duration < 2*time.Minute {
if yes, err := ask.YesNo(
"The duration you provided is less than 2 minutes. "+
Expand All @@ -353,7 +355,7 @@ func (cli *cliBouncers) prune(duration time.Duration, force bool) error {
}
}

bouncers, err := cli.db.QueryBouncersInactiveSince(time.Now().UTC().Add(-duration))
bouncers, err := cli.db.QueryBouncersInactiveSince(ctx, time.Now().UTC().Add(-duration))
if err != nil {
return fmt.Errorf("unable to query bouncers: %w", err)
}
Expand All @@ -376,7 +378,7 @@ func (cli *cliBouncers) prune(duration time.Duration, force bool) error {
}
}

deleted, err := cli.db.BulkDeleteBouncers(bouncers)
deleted, err := cli.db.BulkDeleteBouncers(ctx, bouncers)
if err != nil {
return fmt.Errorf("unable to prune bouncers: %w", err)
}
Expand All @@ -401,8 +403,8 @@ func (cli *cliBouncers) newPruneCmd() *cobra.Command {
DisableAutoGenTag: true,
Example: `cscli bouncers prune -d 45m
cscli bouncers prune -d 45m --force`,
RunE: func(_ *cobra.Command, _ []string) error {
return cli.prune(duration, force)
RunE: func(cmd *cobra.Command, _ []string) error {
return cli.prune(cmd.Context(), duration, force)
},
}

Expand Down
56 changes: 30 additions & 26 deletions cmd/crowdsec-cli/climachine/machines.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package climachine

import (
"context"
"encoding/csv"
"encoding/json"
"errors"
Expand Down Expand Up @@ -210,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)
}
Expand Down Expand Up @@ -251,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)
},
}

Expand All @@ -278,8 +279,8 @@ func (cli *cliMachines) newAddCmd() *cobra.Command {
cscli machines add MyTestMachine --auto
cscli machines add MyTestMachine --password MyPassword
cscli machines add -f- --auto > /tmp/mycreds.yaml`,
RunE: func(_ *cobra.Command, args []string) error {
return cli.add(args, string(password), dumpFile, apiURL, interactive, autoAdd, force)
RunE: func(cmd *cobra.Command, args []string) error {
return cli.add(cmd.Context(), args, string(password), dumpFile, apiURL, interactive, autoAdd, force)
},
}

Expand All @@ -294,7 +295,7 @@ cscli machines add -f- --auto > /tmp/mycreds.yaml`,
return cmd
}

func (cli *cliMachines) add(args []string, machinePassword string, dumpFile string, apiURL string, interactive bool, autoAdd bool, force bool) error {
func (cli *cliMachines) add(ctx context.Context, args []string, machinePassword string, dumpFile string, apiURL string, interactive bool, autoAdd bool, force bool) error {
var (
err error
machineID string
Expand Down Expand Up @@ -353,7 +354,7 @@ func (cli *cliMachines) add(args []string, machinePassword string, dumpFile stri

password := strfmt.Password(machinePassword)

_, err = cli.db.CreateMachine(&machineID, &password, "", true, force, types.PasswordAuthType)
_, err = cli.db.CreateMachine(ctx, &machineID, &password, "", true, force, types.PasswordAuthType)
if err != nil {
return fmt.Errorf("unable to create machine: %w", err)
}
Expand Down Expand Up @@ -399,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

Expand All @@ -407,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
Expand All @@ -430,9 +432,9 @@ func (cli *cliMachines) validMachineID(cmd *cobra.Command, args []string, toComp
return ret, cobra.ShellCompDirectiveNoFileComp
}

func (cli *cliMachines) delete(machines []string, ignoreMissing bool) error {
func (cli *cliMachines) delete(ctx context.Context, machines []string, ignoreMissing bool) error {
for _, machineID := range machines {
if err := cli.db.DeleteWatcher(machineID); err != nil {
if err := cli.db.DeleteWatcher(ctx, machineID); err != nil {
var notFoundErr *database.MachineNotFoundError
if ignoreMissing && errors.As(err, &notFoundErr) {
return nil
Expand Down Expand Up @@ -460,8 +462,8 @@ func (cli *cliMachines) newDeleteCmd() *cobra.Command {
Aliases: []string{"remove"},
DisableAutoGenTag: true,
ValidArgsFunction: cli.validMachineID,
RunE: func(_ *cobra.Command, args []string) error {
return cli.delete(args, ignoreMissing)
RunE: func(cmd *cobra.Command, args []string) error {
return cli.delete(cmd.Context(), args, ignoreMissing)
},
}

Expand All @@ -471,7 +473,7 @@ func (cli *cliMachines) newDeleteCmd() *cobra.Command {
return cmd
}

func (cli *cliMachines) prune(duration time.Duration, notValidOnly bool, force bool) error {
func (cli *cliMachines) prune(ctx context.Context, duration time.Duration, notValidOnly bool, force bool) error {
if duration < 2*time.Minute && !notValidOnly {
if yes, err := ask.YesNo(
"The duration you provided is less than 2 minutes. "+
Expand All @@ -484,12 +486,12 @@ func (cli *cliMachines) prune(duration time.Duration, notValidOnly bool, force b
}

machines := []*ent.Machine{}
if pending, err := cli.db.QueryPendingMachine(); err == nil {
if pending, err := cli.db.QueryPendingMachine(ctx); err == nil {
machines = append(machines, pending...)
}

if !notValidOnly {
if pending, err := cli.db.QueryMachinesInactiveSince(time.Now().UTC().Add(-duration)); err == nil {
if pending, err := cli.db.QueryMachinesInactiveSince(ctx, time.Now().UTC().Add(-duration)); err == nil {
machines = append(machines, pending...)
}
}
Expand All @@ -512,7 +514,7 @@ func (cli *cliMachines) prune(duration time.Duration, notValidOnly bool, force b
}
}

deleted, err := cli.db.BulkDeleteWatchers(machines)
deleted, err := cli.db.BulkDeleteWatchers(ctx, machines)
if err != nil {
return fmt.Errorf("unable to prune machines: %w", err)
}
Expand Down Expand Up @@ -540,8 +542,8 @@ cscli machines prune --duration 1h
cscli machines prune --not-validated-only --force`,
Args: cobra.NoArgs,
DisableAutoGenTag: true,
RunE: func(_ *cobra.Command, _ []string) error {
return cli.prune(duration, notValidOnly, force)
RunE: func(cmd *cobra.Command, _ []string) error {
return cli.prune(cmd.Context(), duration, notValidOnly, force)
},
}

Expand All @@ -553,8 +555,8 @@ cscli machines prune --not-validated-only --force`,
return cmd
}

func (cli *cliMachines) validate(machineID string) error {
if err := cli.db.ValidateMachine(machineID); err != nil {
func (cli *cliMachines) validate(ctx context.Context, machineID string) error {
if err := cli.db.ValidateMachine(ctx, machineID); err != nil {
return fmt.Errorf("unable to validate machine '%s': %w", machineID, err)
}

Expand All @@ -571,8 +573,8 @@ func (cli *cliMachines) newValidateCmd() *cobra.Command {
Example: `cscli machines validate "machine_name"`,
Args: cobra.ExactArgs(1),
DisableAutoGenTag: true,
RunE: func(_ *cobra.Command, args []string) error {
return cli.validate(args[0])
RunE: func(cmd *cobra.Command, args []string) error {
return cli.validate(cmd.Context(), args[0])
},
}

Expand Down Expand Up @@ -690,9 +692,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)
}
Expand Down
Loading
Loading