add Traffic interface

bxy4543 · Jan 9, 2024 · 24ceaff · 24ceaff
1 parent dcaf8c4
commit 24ceaff
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 0 deletions.
diff --git a/controllers/pkg/database/interface.go b/controllers/pkg/database/interface.go
@@ -29,6 +29,7 @@ import (
 type Interface interface {
 	Account
 	Auth
+	Traffic
 }
 
 type Auth interface {
@@ -51,11 +52,17 @@ type Account interface {
 	GetBillingCount(accountType common.Type, startTime, endTime time.Time) (count, amount int64, err error)
 	GenerateBillingData(startTime, endTime time.Time, prols *resources.PropertyTypeLS, namespaces []string, owner string) (orderID []string, amount int64, err error)
 	InsertMonitor(ctx context.Context, monitors ...*resources.Monitor) error
+	GetDistinctMonitorCombinations(startTime, endTime time.Time, namespace string) ([]resources.Monitor, error)
 	DropMonitorCollectionsOlderThan(days int) error
 	Disconnect(ctx context.Context) error
 	Creator
 }
 
+type Traffic interface {
+	GetTrafficSentBytes(startTime, endTime time.Time, namespace string, _type uint8, name string) (int64, error)
+	GetTrafficRecvBytes(startTime, endTime time.Time, namespace string, _type uint8, name string) (int64, error)
+}
+
 type Creator interface {
 	CreateBillingIfNotExist() error
 	//suffix by day, eg： monitor_20200101

diff --git a/controllers/pkg/database/mongo/account.go b/controllers/pkg/database/mongo/account.go
@@ -70,6 +70,16 @@ type mongoDB struct {
 	PropertiesConn    string
 }
 
+func (m *mongoDB) GetTrafficSentBytes(_, _ time.Time, _ string, _ uint8, _ string) (int64, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
+func (m *mongoDB) GetTrafficRecvBytes(_, _ time.Time, _ string, _ uint8, _ string) (int64, error) {
+	//TODO implement me
+	panic("implement me")
+}
+
 type AccountBalanceSpecBSON struct {
 	// Time    metav1.Time `json:"time" bson:"time"`
 	// If the Time field is of the time. time type, it cannot be converted to json crd, so use metav1.Time. However, metav1.Time cannot be inserted into mongo, so you need to convert it to time.Time
@@ -247,6 +257,44 @@ func (m *mongoDB) InsertMonitor(ctx context.Context, monitors ...*resources.Moni
 	return err
 }
 
+func (m *mongoDB) GetDistinctMonitorCombinations(startTime, endTime time.Time, namespace string) ([]resources.Monitor, error) {
+	matchStage := bson.D{
+		{Key: "$match", Value: bson.M{
+			"time": bson.M{
+				"$gte": startTime.UTC(),
+				"$lt":  endTime.UTC(),
+			},
+			"category": namespace,
+		}},
+	}
+	groupStage := bson.D{
+		{Key: "$group", Value: bson.M{
+			"_id": bson.M{
+				"category": "$category",
+				"name":     "$name",
+				"type":     "$type",
+			},
+		}},
+	}
+	cursor, err := m.getMonitorCollection(startTime).Aggregate(context.Background(), mongo.Pipeline{matchStage, groupStage})
+	if err != nil {
+		return nil, fmt.Errorf("aggregate error: %v", err)
+	}
+	defer cursor.Close(context.Background())
+	var monitors []resources.Monitor
+	for cursor.Next(context.Background()) {
+		var monitor resources.Monitor
+		if err := cursor.Decode(&monitor); err != nil {
+			return nil, fmt.Errorf("decode error: %v", err)
+		}
+		monitors = append(monitors, monitor)
+	}
+	if err := cursor.Err(); err != nil {
+		return nil, fmt.Errorf("cursor error: %v", err)
+	}
+	return monitors, nil
+}
+
 func (m *mongoDB) GetAllPricesMap() (map[string]resources.Price, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
@@ -445,6 +493,11 @@ func (m *mongoDB) GenerateBillingData(startTime, endTime time.Time, prols *resou
 			}}}
 			continue
 		}
+		if value.PriceType == resources.SUM {
+			groupStage = append(groupStage, primitive.E{Key: keyStr, Value: bson.D{{Key: "$sum", Value: "$used." + keyStr}}})
+			usedStage[keyStr] = primitive.E{Key: keyStr, Value: bson.D{{Key: "$sum", Value: "$used." + keyStr}}}
+			continue
+		}
 		groupStage = append(groupStage, primitive.E{Key: keyStr, Value: bson.D{{Key: "$sum", Value: "$used." + keyStr}}})
 		usedStage[keyStr] = bson.D{{Key: "$toInt", Value: bson.D{{Key: "$round", Value: bson.D{{Key: "$divide", Value: bson.A{
 			"$" + keyStr, bson.D{{Key: "$cond", Value: bson.A{bson.D{{Key: "$gt", Value: bson.A{"$count", minutes}}}, "$count", minutes}}}}}}}}}}