From 1484591aea6c4f2955ce1a417ee3fca93e3977d0 Mon Sep 17 00:00:00 2001
From: Patrick Pacher <patrick.pacher@gmail.com>
Date: Thu, 5 Oct 2023 10:28:18 +0200
Subject: [PATCH] Split netquery package files and update bandwidth chart
 handler

---
 netquery/bandwidth_chart_handler.go | 20 ++++++++++-------
 netquery/query.go                   | 33 ++++++++++++++++++++---------
 netquery/query_request.go           | 14 ++++++++++++
 3 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/netquery/bandwidth_chart_handler.go b/netquery/bandwidth_chart_handler.go
index 1d61e1de..921673d5 100644
--- a/netquery/bandwidth_chart_handler.go
+++ b/netquery/bandwidth_chart_handler.go
@@ -19,6 +19,7 @@ type BandwidthChartHandler struct {
 }
 
 type BandwidthChartRequest struct {
+	AllProfiles bool     `json:"allProfiles"`
 	Profiles    []string `json:"profiles"`
 	Connections []string `json:"connections"`
 }
@@ -107,18 +108,21 @@ func (req *BandwidthChartRequest) generateSQL(ctx context.Context, schema *orm.T
 	whereClause := ""
 	params := make(map[string]any)
 
-	if len(req.Profiles) > 0 {
+	if (len(req.Profiles) > 0) || (req.AllProfiles == true) {
 		groupBy = []string{"profile", "round(time/10, 0)*10"}
 		selects = append(selects, "profile")
-		clauses := make([]string, len(req.Profiles))
 
-		for idx, p := range req.Profiles {
-			key := fmt.Sprintf(":p%d", idx)
-			clauses[idx] = "profile = " + key
-			params[key] = p
+		if !req.AllProfiles {
+			clauses := make([]string, len(req.Profiles))
+
+			for idx, p := range req.Profiles {
+				key := fmt.Sprintf(":p%d", idx)
+				clauses[idx] = "profile = " + key
+				params[key] = p
+			}
+
+			whereClause = "WHERE " + strings.Join(clauses, " OR ")
 		}
-
-		whereClause = "WHERE " + strings.Join(clauses, " OR ")
 	} else if len(req.Connections) > 0 {
 		groupBy = []string{"conn_id", "round(time/10, 0)*10"}
 		selects = append(selects, "conn_id")
diff --git a/netquery/query.go b/netquery/query.go
index f62cd958..2ef943bc 100644
--- a/netquery/query.go
+++ b/netquery/query.go
@@ -71,12 +71,18 @@ type (
 		Distinct  bool   `json:"distinct"`
 	}
 
+	FieldSelect struct {
+		Field string `json:"field"`
+		As    string `json:"as"`
+	}
+
 	Select struct {
-		Field    string  `json:"field"`
-		Count    *Count  `json:"$count,omitempty"`
-		Sum      *Sum    `json:"$sum,omitempty"`
-		Min      *Min    `json:"$min,omitempty"`
-		Distinct *string `json:"$distinct,omitempty"`
+		Field       string       `json:"field"`
+		FieldSelect *FieldSelect `json:"$field"`
+		Count       *Count       `json:"$count,omitempty"`
+		Sum         *Sum         `json:"$sum,omitempty"`
+		Min         *Min         `json:"$min,omitempty"`
+		Distinct    *string      `json:"$distinct,omitempty"`
 	}
 
 	Selects []Select
@@ -538,11 +544,12 @@ func (sel *Select) UnmarshalJSON(blob []byte) error {
 	// directly
 	if blob[0] == '{' {
 		var res struct {
-			Field    string  `json:"field"`
-			Count    *Count  `json:"$count"`
-			Sum      *Sum    `json:"$sum"`
-			Min      *Min    `json:"$min"`
-			Distinct *string `json:"$distinct"`
+			Field       string       `json:"field"`
+			Count       *Count       `json:"$count"`
+			Sum         *Sum         `json:"$sum"`
+			Min         *Min         `json:"$min"`
+			Distinct    *string      `json:"$distinct"`
+			FieldSelect *FieldSelect `json:"$field"`
 		}
 
 		if err := json.Unmarshal(blob, &res); err != nil {
@@ -551,6 +558,7 @@ func (sel *Select) UnmarshalJSON(blob []byte) error {
 
 		sel.Count = res.Count
 		sel.Field = res.Field
+		sel.FieldSelect = res.FieldSelect
 		sel.Distinct = res.Distinct
 		sel.Sum = res.Sum
 		sel.Min = res.Min
@@ -570,6 +578,11 @@ func (sel *Select) UnmarshalJSON(blob []byte) error {
 				return fmt.Errorf("invalid characters in $min.as, value must match [a-zA-Z]+")
 			}
 		}
+		if sel.FieldSelect != nil && sel.FieldSelect.As != "" {
+			if !charOnlyRegexp.MatchString(sel.FieldSelect.As) {
+				return fmt.Errorf("invalid characters in $field.as, value must match [a-zA-Z]+")
+			}
+		}
 
 		return nil
 	}
diff --git a/netquery/query_request.go b/netquery/query_request.go
index b4a07041..ea5162a9 100644
--- a/netquery/query_request.go
+++ b/netquery/query_request.go
@@ -126,6 +126,8 @@ func (req *QueryRequestPayload) prepareSelectedFields(ctx context.Context, schem
 			} else {
 				field = "*"
 			}
+		case s.FieldSelect != nil:
+			field = s.FieldSelect.Field
 		default:
 			field = s.Field
 		}
@@ -141,6 +143,18 @@ func (req *QueryRequestPayload) prepareSelectedFields(ctx context.Context, schem
 		}
 
 		switch {
+		case s.FieldSelect != nil:
+			as := s.FieldSelect.As
+			if as == "" {
+				as = s.FieldSelect.Field
+			}
+
+			req.selectedFields = append(
+				req.selectedFields,
+				fmt.Sprintf("%s AS %s", s.FieldSelect.Field, as),
+			)
+			req.whitelistedFields = append(req.whitelistedFields, as)
+
 		case s.Count != nil:
 			as := s.Count.As
 			if as == "" {