try to fix ci failures

Signed-off-by: Davanum Srinivas <[email protected]>

Author: dims

node-drainer/main.go

@@ -24,13 +24,13 @@ import (
 	"strconv"
 	"syscall"
 
+	"github.com/nvidia/nvsentinel/commons/pkg/eventutil"
 	"github.com/nvidia/nvsentinel/commons/pkg/flags"
 	"github.com/nvidia/nvsentinel/commons/pkg/logger"
 	"github.com/nvidia/nvsentinel/commons/pkg/server"
 	"github.com/nvidia/nvsentinel/data-models/pkg/model"
 	"github.com/nvidia/nvsentinel/node-drainer/pkg/initializer"
 	"github.com/nvidia/nvsentinel/store-client/pkg/client"
-	"github.com/nvidia/nvsentinel/store-client/pkg/datastore"
 	"github.com/nvidia/nvsentinel/store-client/pkg/query"
 	"golang.org/x/sync/errgroup"
 )
@@ -274,29 +274,37 @@ func handleColdStart(ctx context.Context, components *initializer.Components) er
 
 	slog.Info("Found events to re-process", "count", len(healthEvents))
 
-	// Convert to Event format for compatibility with existing code
-	events := make([]datastore.Event, len(healthEvents))
-	for i, he := range healthEvents {
-		events[i] = he.RawEvent
-	}
-
 	// Re-process each event
-	for _, event := range events {
-		// Extract health event (datastore.Event is map[string]interface{})
-		healthEvent, ok := event["healthevent"].(datastore.Event)
-		if !ok {
-			slog.Error("Failed to extract healthevent from cold start event")
+	for _, he := range healthEvents {
+		// Use the RawEvent from the database query which includes _id
+		// This is critical for status updates to work properly
+		event := he.RawEvent
+		if len(event) == 0 {
+			slog.Error("RawEvent is empty, skipping cold start event")
+			continue
+		}
+
+		// Parse the event to extract node name
+		parsedEvent, err := eventutil.ParseHealthEventFromEvent(event)
+		if err != nil {
+			slog.Error("Failed to parse health event from cold start event", "error", err)
 			continue
 		}
 
-		nodeName, ok := healthEvent["nodename"].(string)
-		if !ok {
-			slog.Error("Failed to extract node name from cold start event")
+		if parsedEvent.HealthEvent == nil {
+			slog.Error("Health event is nil in cold start event")
+			continue
+		}
+
+		nodeName := parsedEvent.HealthEvent.GetNodeName()
+		if nodeName == "" {
+			slog.Error("Node name is empty in cold start event")
 			continue
 		}
 
 		// Create adapter to bridge interface differences
 		dbAdapter := &dataStoreAdapter{DatabaseClient: components.DatabaseClient}
+
 		if err := components.QueueManager.EnqueueEventGeneric(ctx, nodeName, event, dbAdapter); err != nil {
 			slog.Error("Failed to enqueue cold start event", "error", err, "nodeName", nodeName)
 		} else {

store-client/pkg/datastore/providers/mongodb/health_store.go

@@ -18,6 +18,9 @@ import (
 	"context"
 	"fmt"
 
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+
 	"github.com/nvidia/nvsentinel/store-client/pkg/client"
 	"github.com/nvidia/nvsentinel/store-client/pkg/datastore"
 )
@@ -100,6 +103,9 @@ func (h *MongoHealthEventStore) FindHealthEventsByNode(ctx context.Context,
 		).WithMetadata("nodeName", nodeName)
 	}
 
+	// Normalize HealthEvent fields from bson.M to map[string]interface{}
+	normalizeHealthEvents(events)
+
 	return events, nil
 }
 
@@ -118,15 +124,55 @@ func (h *MongoHealthEventStore) FindHealthEventsByFilter(ctx context.Context,
 
 	var events []datastore.HealthEventWithStatus
 
-	err = cursor.All(ctx, &events)
-	if err != nil {
+	// Iterate manually to populate both the struct and RawEvent with _id
+	for cursor.Next(ctx) {
+		// First decode into a raw map to preserve _id
+		var rawDoc map[string]interface{}
+		if err := cursor.Decode(&rawDoc); err != nil {
+			return nil, datastore.NewQueryError(
+				datastore.ProviderMongoDB,
+				"failed to decode raw document",
+				err,
+			).WithMetadata("filter", filter)
+		}
+
+		// Then use bson to convert the raw map into our struct
+		var event datastore.HealthEventWithStatus
+
+		bsonBytes, err := bson.Marshal(rawDoc)
+		if err != nil {
+			return nil, datastore.NewQueryError(
+				datastore.ProviderMongoDB,
+				"failed to marshal raw document to BSON",
+				err,
+			).WithMetadata("filter", filter)
+		}
+
+		if err := bson.Unmarshal(bsonBytes, &event); err != nil {
+			return nil, datastore.NewQueryError(
+				datastore.ProviderMongoDB,
+				"failed to unmarshal BSON to health event",
+				err,
+			).WithMetadata("filter", filter)
+		}
+
+		// Store the raw document (with _id) in RawEvent
+		event.RawEvent = rawDoc
+
+		events = append(events, event)
+	}
+
+	if err := cursor.Err(); err != nil {
 		return nil, datastore.NewQueryError(
 			datastore.ProviderMongoDB,
-			"failed to decode health events by filter",
+			"cursor error while iterating health events",
 			err,
 		).WithMetadata("filter", filter)
 	}
 
+	// Normalize HealthEvent fields from bson.M to map[string]interface{}
+	normalizeHealthEvents(events)
+
 	return events, nil
 }
 
@@ -162,6 +208,9 @@ func (h *MongoHealthEventStore) FindHealthEventsByStatus(ctx context.Context,
 		).WithMetadata("status", string(status))
 	}
 
+	// Normalize HealthEvent fields from bson.M to map[string]interface{}
+	normalizeHealthEvents(events)
+
 	return events, nil
 }
 
@@ -281,3 +330,80 @@ func (h *MongoHealthEventStore) UpdateHealthEventsByQuery(ctx context.Context,
 
 	return nil
 }
+
+// normalizeHealthEvents converts bson.M types to map[string]interface{} in HealthEvent fields
+// This ensures consistency across database providers (MongoDB-specific types -> generic maps)
+func normalizeHealthEvents(events []datastore.HealthEventWithStatus) {
+	for i := range events {
+		if events[i].HealthEvent != nil {
+			events[i].HealthEvent = normalizeValue(events[i].HealthEvent)
+		}
+	}
+}
+
+// normalizeValue recursively converts MongoDB types (bson.M, primitive.D, primitive.A, etc.) to standard Go types
+func normalizeValue(v interface{}) interface{} {
+	switch val := v.(type) {
+	case primitive.D:
+		return normalizePrimitiveD(val)
+	case primitive.A:
+		return normalizeArray(val)
+	case map[string]interface{}:
+		return normalizeMap(val)
+	case []interface{}:
+		return normalizeArray(val)
+	default:
+		// Primitive types, return as-is
+		return val
+	}
+}
+
+// normalizePrimitiveD converts primitive.D to map[string]interface{} and normalizes nested values
+func normalizePrimitiveD(val primitive.D) interface{} {
+	bsonBytes, err := bson.Marshal(val)
+	if err != nil {
+		return val // Return as-is if marshal fails
+	}
+
+	var m map[string]interface{}
+	if err := bson.Unmarshal(bsonBytes, &m); err != nil {
+		return val // Return as-is if unmarshal fails
+	}
+
+	return normalizeMap(m)
+}
+
+// normalizeMap recursively normalizes all values in a map
+func normalizeMap(m map[string]interface{}) map[string]interface{} {
+	result := make(map[string]interface{}, len(m))
+	for k, v := range m {
+		result[k] = normalizeValue(v)
+	}
+
+	return result
+}
+
+// normalizeArray recursively normalizes all elements in an array
+func normalizeArray(arr interface{}) []interface{} {
+	var length int
+
+	var getValue func(int) interface{}
+
+	switch v := arr.(type) {
+	case primitive.A:
+		length = len(v)
+		getValue = func(i int) interface{} { return v[i] }
+	case []interface{}:
+		length = len(v)
+		getValue = func(i int) interface{} { return v[i] }
+	default:
+		return nil
+	}
+
+	result := make([]interface{}, length)
+	for i := 0; i < length; i++ {
+		result[i] = normalizeValue(getValue(i))
+	}
+
+	return result
+}

store-client/pkg/datastore/providers/mongodb/health_store_test.go

@@ -17,11 +17,13 @@ package mongodb
 import (
 	"context"
 	"errors"
+	"reflect"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 
 	"github.com/nvidia/nvsentinel/store-client/pkg/client"
 	"github.com/nvidia/nvsentinel/store-client/pkg/datastore"
@@ -287,3 +289,113 @@ func TestMongoHealthEventStore_FindLatestEventForNode(t *testing.T) {
 		mockResult.AssertExpectations(t)
 	})
 }
+
+func TestNormalizeValue(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    interface{}
+		expected interface{}
+	}{
+		{
+			name: "primitive.D with nested values",
+			input: primitive.D{
+				{Key: "nodename", Value: "test-node"},
+				{Key: "status", Value: "healthy"},
+				{Key: "metadata", Value: primitive.D{
+					{Key: "region", Value: "us-west"},
+					{Key: "zone", Value: "a"},
+				}},
+			},
+			expected: map[string]interface{}{
+				"nodename": "test-node",
+				"status":   "healthy",
+				"metadata": map[string]interface{}{
+					"region": "us-west",
+					"zone":   "a",
+				},
+			},
+		},
+		{
+			name: "primitive.D with array containing primitive.D",
+			input: primitive.D{
+				{Key: "items", Value: []interface{}{
+					primitive.D{
+						{Key: "name", Value: "item1"},
+						{Key: "value", Value: int32(100)},
+					},
+					primitive.D{
+						{Key: "name", Value: "item2"},
+						{Key: "value", Value: int32(200)},
+					},
+				}},
+			},
+			expected: map[string]interface{}{
+				"items": []interface{}{
+					map[string]interface{}{
+						"name":  "item1",
+						"value": int32(100),
+					},
+					map[string]interface{}{
+						"name":  "item2",
+						"value": int32(200),
+					},
+				},
+			},
+		},
+		{
+			name: "already normalized map",
+			input: map[string]interface{}{
+				"key1": "value1",
+				"key2": map[string]interface{}{
+					"nested": "value2",
+				},
+			},
+			expected: map[string]interface{}{
+				"key1": "value1",
+				"key2": map[string]interface{}{
+					"nested": "value2",
+				},
+			},
+		},
+		{
+			name: "array with mixed types",
+			input: []interface{}{
+				"string",
+				123,
+				primitive.D{{Key: "key", Value: "value"}},
+				map[string]interface{}{"already": "normalized"},
+			},
+			expected: []interface{}{
+				"string",
+				123,
+				map[string]interface{}{"key": "value"},
+				map[string]interface{}{"already": "normalized"},
+			},
+		},
+		{
+			name:     "primitive string",
+			input:    "test-string",
+			expected: "test-string",
+		},
+		{
+			name:     "primitive int",
+			input:    42,
+			expected: 42,
+		},
+		{
+			name:     "nil value",
+			input:    nil,
+			expected: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := normalizeValue(tt.input)
+			if !reflect.DeepEqual(result, tt.expected) {
+				t.Errorf("normalizeValue() = %v (type: %T), expected %v (type: %T)",
+					result, result, tt.expected, tt.expected)
+			}
+		})
+	}
+}

store-client/pkg/datastore/providers/postgresql/datastore.go

@@ -195,7 +195,8 @@ func (p *PostgreSQLDataStore) NewChangeStreamWatcher(
 	watcher := NewPostgreSQLChangeStreamWatcher(p.db, clientName, tableName)
 	watcher.pipelineFilter = pipelineFilter
 
-	return watcher, nil
+	// Wrap the watcher to provide Unwrap() support for backward compatibility
+	return NewPostgreSQLChangeStreamWatcherWithUnwrap(watcher), nil
 }
 
 // --- Backward Compatibility Methods for MongoDB-style Type Assertions ---