use Batch helper for decisions

Author: mmetc

pkg/database/decisions.go

@@ -9,8 +9,6 @@ import (
 	"entgo.io/ent/dialect/sql"
 	"github.com/pkg/errors"
 
-	"github.com/crowdsecurity/go-cs-lib/slicetools"
-
 	"github.com/crowdsecurity/crowdsec/pkg/csnet"
 	"github.com/crowdsecurity/crowdsec/pkg/database/ent"
 	"github.com/crowdsecurity/crowdsec/pkg/database/ent/decision"
@@ -298,68 +296,79 @@ func decisionIDs(decisions []*ent.Decision) []int {
 	return ids
 }
 
-// ExpireDecisions sets the expiration of a list of decisions to now()
-// It returns the number of impacted decisions for the CAPI/PAPI
-func (c *Client) ExpireDecisions(ctx context.Context, decisions []*ent.Decision) (int, error) {
-	if len(decisions) <= decisionDeleteBulkSize {
-		ids := decisionIDs(decisions)
+// expireDecisionBatch expires the decisions as a single operation.
+func (c *Client) expireDecisionBatch(ctx context.Context, batch []*ent.Decision, now time.Time) (int, error) {
+	ids := decisionIDs(batch)
 
-		rows, err := c.Ent.Decision.Update().Where(
-			decision.IDIn(ids...),
-		).SetUntil(time.Now().UTC()).Save(ctx)
-		if err != nil {
-			return 0, fmt.Errorf("expire decisions with provided filter: %w", err)
-		}
+	rows, err := c.Ent.Decision.
+		Update().
+		Where(decision.IDIn(ids...)).
+		SetUntil(now).
+		Save(ctx)
+	if err != nil {
+		return 0, fmt.Errorf("expire decisions with provided filter: %w", err)
+	}
 
-		return rows, nil
+	return rows, nil
+}
+
+// ExpireDecisions sets the expiration of a list of decisions to now(),
+// in multiple operations if len(decisions) > decisionDeleteBulkSize.
+// It returns the number of impacted decisions for the CAPI/PAPI, even in case of error.
+func (c *Client) ExpireDecisions(ctx context.Context, decisions []*ent.Decision) (int, error) {
+	if len(decisions) == 0 {
+		return 0, nil
 	}
 
-	// big batch, let's split it and recurse
+	now := time.Now().UTC()
 
 	total := 0
-
-	for _, chunk := range slicetools.Chunks(decisions, decisionDeleteBulkSize) {
-		rows, err := c.ExpireDecisions(ctx, chunk)
+	err := Batch(ctx, decisions, decisionDeleteBulkSize, func(ctx context.Context, batch []*ent.Decision) error {
+		rows, err := c.expireDecisionBatch(ctx, batch, now)
 		if err != nil {
-			return total, err
+			return err
 		}
-
 		total += rows
-	}
+		return nil
+	})
 
-	return total, nil
+	return total, err
 }
 
-// DeleteDecisions removes a list of decisions from the database
-// It returns the number of impacted decisions for the CAPI/PAPI
-func (c *Client) DeleteDecisions(ctx context.Context, decisions []*ent.Decision) (int, error) {
-	if len(decisions) < decisionDeleteBulkSize {
-		ids := decisionIDs(decisions)
-
-		rows, err := c.Ent.Decision.Delete().Where(
-			decision.IDIn(ids...),
-		).Exec(ctx)
-		if err != nil {
-			return 0, fmt.Errorf("hard delete decisions with provided filter: %w", err)
-		}
+// deleteDecisionBatch removes the decisions as a single operation.
+func (c *Client) deleteDecisionBatch(ctx context.Context, batch []*ent.Decision) (int, error) {
+	ids := decisionIDs(batch)
 
-		return rows, nil
+	rows, err := c.Ent.Decision.
+		Delete().
+		Where(decision.IDIn(ids...)).
+		Exec(ctx)
+	if err != nil {
+		return 0, fmt.Errorf("hard delete decisions with provided filter: %w", err)
 	}
 
-	// big batch, let's split it and recurse
+	return rows, nil
+}
 
-	tot := 0
+// DeleteDecisions removes a list of decisions from the database,
+// in multiple operations if len(decisions) > decisionDeleteBulkSize.
+// It returns the number of impacted decisions for the CAPI/PAPI, even in case of error.
+func (c *Client) DeleteDecisions(ctx context.Context, decisions []*ent.Decision) (int, error) {
+	if len(decisions) == 0 {
+		return 0, nil
+	}
 
-	for _, chunk := range slicetools.Chunks(decisions, decisionDeleteBulkSize) {
-		rows, err := c.DeleteDecisions(ctx, chunk)
+	total := 0
+	err := Batch(ctx, decisions, decisionDeleteBulkSize, func(ctx context.Context, batch []*ent.Decision) error {
+		rows, err := c.deleteDecisionBatch(ctx, batch)
 		if err != nil {
-			return tot, err
+			return err
 		}
+		total += rows
+		return nil
+	})
 
-		tot += rows
-	}
-
-	return tot, nil
+	return total, err
 }
 
 // ExpireDecisionByID set the expiration of a decision to now()