historical_uptime: modularize queryByIndex and updateByIndex

Signed-off-by: bingyuyap <bingyu.yap.21@gmail.com>

Author: bingyuyap

fly/pkg/db/message.go

@@ -154,58 +154,6 @@ func (db *Database) AppendObservationIfNotExist(messageID string, observation Ob
 	})
 }
 
-// QueryMessagesByIndex retrieves messages based on indexed attributes.
-func (db *Database) QueryMessagesByIndex(metricsChecked bool, cutOffTime time.Duration) ([]*Message, error) {
-	var messagesToUpdate []*Message
-	now := time.Now()
-
-	// Start a read-only transaction
-	err := db.db.View(func(txn *badger.Txn) error {
-		opts := badger.DefaultIteratorOptions
-		opts.PrefetchValues = false // Only keys are needed
-		it := txn.NewIterator(opts)
-		defer it.Close()
-
-		metricsCheckedPrefix := fmt.Sprintf("metricsChecked|%t|", metricsChecked)
-
-		// Iterate over the metricsChecked index
-		for it.Seek([]byte(metricsCheckedPrefix)); it.ValidForPrefix([]byte(metricsCheckedPrefix)); it.Next() {
-			item := it.Item()
-			key := item.Key()
-
-			// Extract MessageID from the key and query lastObservedAt index
-			_, messageID, err := parseMetricsCheckedIndexKey(key)
-			if err != nil {
-				return fmt.Errorf("failed to parse index key: %w", err)
-			}
-
-			message, err := db.GetMessage(messageID)
-			if err != nil {
-				return fmt.Errorf("failed to get message by ID: %w", err)
-			}
-
-			// Check if the last observed timestamp is before the specified hours
-			if message.LastObservedAt.Before(now.Add(-cutOffTime)) {
-				message.MetricsChecked = true
-				messagesToUpdate = append(messagesToUpdate, message)
-			}
-		}
-
-		return nil
-	})
-
-	if err != nil {
-		return nil, fmt.Errorf("failed to query messages: %w", err)
-	}
-
-	// Perform batch updates in smaller transactions
-	if err := db.batchUpdateMessages(messagesToUpdate); err != nil {
-		return nil, fmt.Errorf("failed to batch update messages: %w", err)
-	}
-
-	return messagesToUpdate, nil
-}
-
 // batchUpdateMessages performs batch updates on a slice of messages.
 func (db *Database) batchUpdateMessages(messages []*Message) error {
 	batchSize := 100 // Adjust the batch size as needed
@@ -237,50 +185,107 @@ func (db *Database) updateMessagesBatch(messagesBatch []*Message) error {
 	})
 }
 
-func (db *Database) RemoveObservationsByIndex(metricsChecked bool, cutOffTime time.Duration) error {
-	now := time.Now()
-	var messagesToUpdate []*Message
-
-	// Start a read-write transaction
-	err := db.db.Update(func(txn *badger.Txn) error {
-		opts := badger.DefaultIteratorOptions
-		opts.PrefetchValues = false // Only keys are needed
+// iterateIndex iterates over a metricsChecked index and applies a callback function to each item.
+func (db *Database) iterateIndex(metricsChecked bool, callback func(item *badger.Item) error) error {
+	opts := badger.DefaultIteratorOptions
+	opts.PrefetchValues = false // Only keys are needed
+	return db.db.View(func(txn *badger.Txn) error {
 		it := txn.NewIterator(opts)
 		defer it.Close()
 
 		metricsCheckedPrefix := fmt.Sprintf("metricsChecked|%t|", metricsChecked)
-
-		// Iterate over the metricsChecked index
 		for it.Seek([]byte(metricsCheckedPrefix)); it.ValidForPrefix([]byte(metricsCheckedPrefix)); it.Next() {
-			item := it.Item()
-			key := item.Key()
-
-			// Extract MessageID from the key
-			_, messageID, err := parseMetricsCheckedIndexKey(key)
-			if err != nil {
-				return fmt.Errorf("failed to parse index key: %w", err)
+			if err := callback(it.Item()); err != nil {
+				return err
 			}
+		}
+		return nil
+	})
+}
 
-			message, err := db.GetMessage(messageID)
-			if err != nil {
-				return fmt.Errorf("failed to get message by ID: %w", err)
-			}
+// processMessage applies the provided logic to the message if it meets the specified conditions.
+func (db *Database) processMessage(messageID string, now time.Time, cutOffTime time.Duration, updateFunc func(*Message) bool) (*Message, error) {
+	message, err := db.GetMessage(messageID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get message by ID: %w", err)
+	}
 
-			// Check if the last observed timestamp is before the specified cutOffTime
-			if message.LastObservedAt.Before(now.Add(-cutOffTime)) {
-				// Remove observations array from the message
-				message.Observations = nil
-				messagesToUpdate = append(messagesToUpdate, message)
-			}
+	if message.LastObservedAt.Before(now.Add(-cutOffTime)) && updateFunc(message) {
+		return message, nil
+	}
+	return nil, nil
+}
+
+// QueryMessagesByIndex retrieves messages based on indexed attributes.
+func (db *Database) QueryMessagesByIndex(metricsChecked bool, cutOffTime time.Duration) ([]*Message, error) {
+	var messagesToUpdate []*Message
+	now := time.Now()
+
+	err := db.iterateIndex(metricsChecked, func(item *badger.Item) error {
+		key := item.Key()
+		_, messageID, err := parseMetricsCheckedIndexKey(key)
+		if err != nil {
+			return fmt.Errorf("failed to parse index key: %w", err)
+		}
+
+		message, err := db.processMessage(messageID, now, cutOffTime, func(m *Message) bool {
+			m.MetricsChecked = true
+			return true
+		})
+
+		if err != nil {
+			return err
+		}
+
+		if message != nil {
+			messagesToUpdate = append(messagesToUpdate, message)
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to query messages: %w", err)
+	}
+
+	if err := db.batchUpdateMessages(messagesToUpdate); err != nil {
+		return nil, fmt.Errorf("failed to batch update messages: %w", err)
+	}
+
+	return messagesToUpdate, nil
+}
+
+// RemoveObservationsByIndex removes observations from messages based on indexed attributes.
+func (db *Database) RemoveObservationsByIndex(metricsChecked bool, cutOffTime time.Duration) error {
+	var messagesToUpdate []*Message
+	now := time.Now()
+
+	err := db.iterateIndex(metricsChecked, func(item *badger.Item) error {
+		key := item.Key()
+		_, messageID, err := parseMetricsCheckedIndexKey(key)
+		if err != nil {
+			return fmt.Errorf("failed to parse index key: %w", err)
+		}
+
+		message, err := db.processMessage(messageID, now, cutOffTime, func(m *Message) bool {
+			m.Observations = nil
+			return true
+		})
+
+		if err != nil {
+			return err
 		}
 
+		if message != nil {
+			messagesToUpdate = append(messagesToUpdate, message)
+		}
+		
 		return nil
 	})
 
 	if err != nil {
 		return err
 	}
 
-	// Perform batch updates in smaller transactions
 	return db.batchUpdateMessages(messagesToUpdate)
 }