From 4cc94063b2f39725bfdf094a37776a410ef31ad4 Mon Sep 17 00:00:00 2001
From: marco <marco@crowdsec.net>
Date: Wed, 18 Sep 2024 14:10:15 +0200
Subject: [PATCH] alerts wip/4

---
 pkg/apiserver/controllers/v1/alerts.go |  4 +++-
 pkg/database/alerts.go                 | 20 ++++++++++----------
 pkg/database/flush.go                  |  4 ++--
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/pkg/apiserver/controllers/v1/alerts.go b/pkg/apiserver/controllers/v1/alerts.go
index 84b3094865c..fd81614cec2 100644
--- a/pkg/apiserver/controllers/v1/alerts.go
+++ b/pkg/apiserver/controllers/v1/alerts.go
@@ -261,7 +261,9 @@ func (c *Controller) CreateAlert(gctx *gin.Context) {
 
 // FindAlerts: returns alerts from the database based on the specified filter
 func (c *Controller) FindAlerts(gctx *gin.Context) {
-	result, err := c.DBClient.QueryAlertWithFilter(gctx.Request.URL.Query())
+	ctx := gctx.Request.Context()
+
+	result, err := c.DBClient.QueryAlertWithFilter(ctx, gctx.Request.URL.Query())
 	if err != nil {
 		c.HandleDBErrors(gctx, err)
 		return
diff --git a/pkg/database/alerts.go b/pkg/database/alerts.go
index ece5e31fade..aa5f5d044b0 100644
--- a/pkg/database/alerts.go
+++ b/pkg/database/alerts.go
@@ -996,11 +996,11 @@ func (c *Client) AlertsCountPerScenario(ctx context.Context, filters map[string]
 	return counts, nil
 }
 
-func (c *Client) TotalAlerts() (int, error) {
-	return c.Ent.Alert.Query().Count(c.CTX)
+func (c *Client) TotalAlerts(ctx context.Context) (int, error) {
+	return c.Ent.Alert.Query().Count(ctx)
 }
 
-func (c *Client) QueryAlertWithFilter(filter map[string][]string) ([]*ent.Alert, error) {
+func (c *Client) QueryAlertWithFilter(ctx context.Context, filter map[string][]string) ([]*ent.Alert, error) {
 	sort := "DESC" // we sort by desc by default
 
 	if val, ok := filter["sort"]; ok {
@@ -1047,7 +1047,7 @@ func (c *Client) QueryAlertWithFilter(filter map[string][]string) ([]*ent.Alert,
 			WithOwner()
 
 		if limit == 0 {
-			limit, err = alerts.Count(c.CTX)
+			limit, err = alerts.Count(ctx)
 			if err != nil {
 				return nil, fmt.Errorf("unable to count nb alerts: %w", err)
 			}
@@ -1059,7 +1059,7 @@ func (c *Client) QueryAlertWithFilter(filter map[string][]string) ([]*ent.Alert,
 			alerts = alerts.Order(ent.Desc(alert.FieldCreatedAt), ent.Desc(alert.FieldID))
 		}
 
-		result, err := alerts.Limit(paginationSize).Offset(offset).All(c.CTX)
+		result, err := alerts.Limit(paginationSize).Offset(offset).All(ctx)
 		if err != nil {
 			return nil, errors.Wrapf(QueryFail, "pagination size: %d, offset: %d: %s", paginationSize, offset, err)
 		}
@@ -1088,35 +1088,35 @@ func (c *Client) QueryAlertWithFilter(filter map[string][]string) ([]*ent.Alert,
 	return ret, nil
 }
 
-func (c *Client) DeleteAlertGraphBatch(alertItems []*ent.Alert) (int, error) {
+func (c *Client) DeleteAlertGraphBatch(ctx context.Context, alertItems []*ent.Alert) (int, error) {
 	idList := make([]int, 0)
 	for _, alert := range alertItems {
 		idList = append(idList, alert.ID)
 	}
 
 	_, err := c.Ent.Event.Delete().
-		Where(event.HasOwnerWith(alert.IDIn(idList...))).Exec(c.CTX)
+		Where(event.HasOwnerWith(alert.IDIn(idList...))).Exec(ctx)
 	if err != nil {
 		c.Log.Warningf("DeleteAlertGraphBatch : %s", err)
 		return 0, errors.Wrapf(DeleteFail, "alert graph delete batch events")
 	}
 
 	_, err = c.Ent.Meta.Delete().
-		Where(meta.HasOwnerWith(alert.IDIn(idList...))).Exec(c.CTX)
+		Where(meta.HasOwnerWith(alert.IDIn(idList...))).Exec(ctx)
 	if err != nil {
 		c.Log.Warningf("DeleteAlertGraphBatch : %s", err)
 		return 0, errors.Wrapf(DeleteFail, "alert graph delete batch meta")
 	}
 
 	_, err = c.Ent.Decision.Delete().
-		Where(decision.HasOwnerWith(alert.IDIn(idList...))).Exec(c.CTX)
+		Where(decision.HasOwnerWith(alert.IDIn(idList...))).Exec(ctx)
 	if err != nil {
 		c.Log.Warningf("DeleteAlertGraphBatch : %s", err)
 		return 0, errors.Wrapf(DeleteFail, "alert graph delete batch decisions")
 	}
 
 	deleted, err := c.Ent.Alert.Delete().
-		Where(alert.IDIn(idList...)).Exec(c.CTX)
+		Where(alert.IDIn(idList...)).Exec(ctx)
 	if err != nil {
 		c.Log.Warningf("DeleteAlertGraphBatch : %s", err)
 		return deleted, errors.Wrapf(DeleteFail, "alert graph delete batch")
diff --git a/pkg/database/flush.go b/pkg/database/flush.go
index 46c8edfa308..2ec88ce34b8 100644
--- a/pkg/database/flush.go
+++ b/pkg/database/flush.go
@@ -239,7 +239,7 @@ func (c *Client) FlushAlerts(ctx context.Context, MaxAge string, MaxItems int) e
 	c.FlushOrphans(ctx)
 	c.Log.Debug("Done flushing orphan alerts")
 
-	totalAlerts, err = c.TotalAlerts()
+	totalAlerts, err = c.TotalAlerts(ctx)
 	if err != nil {
 		c.Log.Warningf("FlushAlerts (max items count): %s", err)
 		return fmt.Errorf("unable to get alerts count: %w", err)
@@ -268,7 +268,7 @@ func (c *Client) FlushAlerts(ctx context.Context, MaxAge string, MaxItems int) e
 		// This gives us the oldest alert that we want to keep
 		// We then delete all the alerts with an id lower than this one
 		// We can do this because the id is auto-increment, and the database won't reuse the same id twice
-		lastAlert, err := c.QueryAlertWithFilter(map[string][]string{
+		lastAlert, err := c.QueryAlertWithFilter(ctx, map[string][]string{
 			"sort":  {"DESC"},
 			"limit": {"1"},
 			// we do not care about fetching the edges, we just want the id