Merge branch 'main' into xrfxlp/device-count-labels

Author: XRFXLP

store-client/pkg/client/postgresql_client.go

@@ -20,7 +20,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"regexp"
 	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/XSAM/otelsql"
@@ -392,26 +394,26 @@ func (c *PostgreSQLClient) UpdateDocumentStatusFields(
 func (c *PostgreSQLClient) UpdateDocument(
 	ctx context.Context, filter interface{}, update interface{},
 ) (*UpdateResult, error) {
-	// Build WHERE clause from filter
-	whereClause, args, err := c.buildWhereClause(filter)
+	// Build SET clause first so its placeholders own the initial parameter range.
+	setClause, updateArgs, err := c.buildUpdateClause(update)
 	if err != nil {
 		return nil, err
 	}
 
-	// Build SET clause from update
-	setClause, updateArgs, err := c.buildUpdateClause(update)
+	whereClause, filterArgs, err := c.buildWhereClause(filter)
 	if err != nil {
 		return nil, err
 	}
 
-	// Combine args (WHERE args + SET args)
-	args = append(args, updateArgs...)
+	adjustedWhereClause := c.adjustParameterNumbers(whereClause, len(updateArgs))
+	args := updateArgs
+	args = append(args, filterArgs...)
 
 	// Build final query
 	//nolint:gosec // G201: table name from config, clauses built with parameterized queries
 	query := fmt.Sprintf(
 		"UPDATE %s SET %s, updated_at = NOW() WHERE %s",
-		c.table, setClause, whereClause,
+		c.table, setClause, adjustedWhereClause,
 	)
 
 	result, err := c.db.ExecContext(ctx, query, args...)
@@ -2622,17 +2624,16 @@ func (c *PostgreSQLClient) adjustParameterNumbers(clause string, offset int) str
 		return clause
 	}
 
-	// Replace parameter placeholders: $1 → $N, $2 → $N+1, etc.
-	// This is a simple implementation; for production, use a more robust parser
-	result := clause
+	paramRe := regexp.MustCompile(`\$(\d+)`)
 
-	for i := 20; i >= 1; i-- { // Process in reverse to avoid double replacement
-		oldParam := fmt.Sprintf("$%d", i)
-		newParam := fmt.Sprintf("$%d", i+offset)
-		result = strings.ReplaceAll(result, oldParam, newParam)
-	}
+	return paramRe.ReplaceAllStringFunc(clause, func(match string) string {
+		n, err := strconv.Atoi(match[1:])
+		if err != nil {
+			return match
+		}
 
-	return result
+		return fmt.Sprintf("$%d", n+offset)
+	})
 }
 
 // buildUpdateClause converts MongoDB-style update operators to PostgreSQL SET clause

store-client/pkg/client/postgresql_client_test.go

@@ -16,8 +16,11 @@ package client
 
 import (
 	"context"
+	"regexp"
 	"strings"
 	"testing"
+
+	"github.com/DATA-DOG/go-sqlmock"
 )
 
 // TestPostgreSQLClient_BasicOperations tests basic CRUD operations
@@ -207,6 +210,59 @@ func TestBuildUpdateClause(t *testing.T) {
 	}
 }
 
+func TestUpdateDocumentOffsetsWherePlaceholdersAfterUpdateArgs(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("failed to create sql mock: %v", err)
+	}
+	defer db.Close()
+
+	client := &PostgreSQLClient{
+		db:    db,
+		table: "health_events",
+	}
+
+	filter := map[string]interface{}{
+		"status": "old",
+	}
+	update := map[string]interface{}{
+		"$set": map[string]interface{}{
+			"status": "new",
+		},
+	}
+
+	expectedSQL := "UPDATE health_events SET document = jsonb_set(document, '{status}', $1), updated_at = NOW() " +
+		"WHERE document->>'status' = $2"
+	mock.ExpectExec(regexp.QuoteMeta(expectedSQL)).
+		WithArgs(`"new"`, "old").
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	result, err := client.UpdateDocument(context.Background(), filter, update)
+	if err != nil {
+		t.Fatalf("UpdateDocument returned error: %v", err)
+	}
+
+	if result.MatchedCount != 1 || result.ModifiedCount != 1 {
+		t.Fatalf("expected one matched and modified document, got matched=%d modified=%d",
+			result.MatchedCount, result.ModifiedCount)
+	}
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("unmet sql expectations: %v", err)
+	}
+}
+
+func TestAdjustParameterNumbersHandlesMultiDigitPlaceholders(t *testing.T) {
+	client := &PostgreSQLClient{}
+
+	clause := "document->>'field1' = $1 AND document->>'field10' = $10"
+	expected := "document->>'field1' = $3 AND document->>'field10' = $12"
+
+	if got := client.adjustParameterNumbers(clause, 2); got != expected {
+		t.Fatalf("expected %q, got %q", expected, got)
+	}
+}
+
 // TestAggregationPipelineConversion tests aggregation pipeline parsing
 func TestAggregationPipelineConversion(t *testing.T) {
 	client := &PostgreSQLClient{table: "health_events"}