Fix unique-by-args handling of literal JSON keys

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed `UniqueOpts.ByArgs` skipping distinct jobs or failing inserts when JSON keys contain path syntax (like `user.id`), are empty, or come from unnamed tags like `json:",omitempty"`. Unaffected unique keys remain unchanged; affected jobs may be inserted again after upgrading or by old and new clients during a rolling upgrade. [PR #1387](https://github.com/riverqueue/river/pull/1387).
 - Fixed SQLite job list pagination skipping or repeating jobs by formatting cursor timestamps consistently with stored timestamps. [PR #1374](https://github.com/riverqueue/river/pull/1374).
 - Improved PostgreSQL job listing performance when filtering by one finalized state (`completed`, `cancelled`, or `discarded`) and sorting by finalized time, including in River UI. [PR #1374](https://github.com/riverqueue/river/pull/1374).
 - Fixed `JobRescuer` overwriting jobs that complete, leave the running state, or are claimed again by another worker after being fetched for rescue, preserving their state, errors, metadata, and timestamps across PostgreSQL and SQLite drivers. Fixes [#1302](https://github.com/riverqueue/river/issues/1302). [PR #1373](https://github.com/riverqueue/river/pull/1373).

internal/dbunique/db_unique.go

@@ -2,6 +2,8 @@ package dbunique
 
 import (
 	"crypto/sha256"
+	"encoding/json"
+	"errors"
 	"slices"
 	"strings"
 	"time"
@@ -12,7 +14,6 @@ import (
 	"github.com/riverqueue/river/rivershared/structtag"
 	"github.com/riverqueue/river/rivershared/uniquestates"
 	"github.com/riverqueue/river/rivershared/util/ptrutil"
-	"github.com/riverqueue/river/rivershared/util/sliceutil"
 	"github.com/riverqueue/river/rivertype"
 )
 
@@ -93,20 +94,10 @@ func buildUniqueKeyString(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpt
 			encodedArgsForUnique = sortedJSONWithOnlyUniqueValues
 		} else {
 			// Use all keys from EncodedArgs sorted alphabetically
-			keys := sliceutil.Map(gjson.GetBytes(params.EncodedArgs, "@keys").Array(), func(v gjson.Result) string { return v.String() })
-			slices.Sort(keys)
-
-			sortedJSON := make([]byte, 0, len(params.EncodedArgs))
-			sortedJSON = append(sortedJSON, "{}"...)
-			sjsonOpts := &sjson.Options{ReplaceInPlace: true}
-			for _, key := range keys {
-				sortedJSON, err = sjson.SetRawBytesOptions(sortedJSON, key, []byte(gjson.GetBytes(params.EncodedArgs, key).Raw), sjsonOpts)
-				if err != nil {
-					// Should not happen unless key was invalid
-					return "", err
-				}
+			encodedArgsForUnique, err = appendSortedObject(make([]byte, 0, len(params.EncodedArgs)), params.EncodedArgs)
+			if err != nil {
+				return "", err
 			}
-			encodedArgsForUnique = sortedJSON
 		}
 
 		sb.WriteString("&args=")
@@ -129,3 +120,73 @@ func buildUniqueKeyString(timeGen rivertype.TimeGenerator, uniqueOpts *UniqueOpt
 
 	return sb.String(), nil
 }
+
+// appendJSONKey appends key to buf as a JSON string, encoded the same way sjson
+// encodes object keys: verbatim between quotes, unless the key contains a byte
+// below 0x20 or above 0x7f, `"`, or `\`, in which case it's encoded with
+// encoding/json (which also escapes `<`, `>`, `&`, U+2028, and U+2029, and
+// replaces invalid UTF-8 with U+FFFD). Matching sjson keeps unique keys
+// identical to those of earlier versions, which built unique args with sjson.
+func appendJSONKey(buf []byte, key string) []byte {
+	for i := range len(key) {
+		if key[i] < ' ' || key[i] > 0x7f || key[i] == '"' || key[i] == '\\' {
+			encodedKey, _ := json.Marshal(key) //nolint:errchkjson // marshaling a string can't fail
+			return append(buf, encodedKey...)
+		}
+	}
+
+	buf = append(buf, '"')
+	buf = append(buf, key...)
+	return append(buf, '"')
+}
+
+// appendSortedObject appends to buf a compact JSON object containing each
+// top-level key of encodedObject along with its raw value, sorted by key. If
+// a key appears more than once, its first value is used. As in the previous
+// path-based implementation, empty input and an empty array produce `{}`;
+// other non-object input is rejected.
+//
+// Keys are walked directly rather than addressed as gjson/sjson paths so that
+// keys containing path syntax (like `.`, `@`, or a leading `:`) and empty
+// keys are included literally, and distinct values of such keys produce
+// distinct output. For all other keys, output is byte-identical to that
+// produced by setting each key onto `{}` with sjson.
+func appendSortedObject(buf, encodedObject []byte) ([]byte, error) {
+	type keyValue struct {
+		key      string
+		rawValue string
+	}
+
+	var (
+		keyValues []keyValue
+		keysSeen  = make(map[string]struct{})
+		err       error
+	)
+	gjson.ParseBytes(encodedObject).ForEach(func(key, value gjson.Result) bool {
+		if key.Type != gjson.String {
+			err = errors.New("unique args must encode a JSON object")
+			return false
+		}
+		if _, ok := keysSeen[key.Str]; !ok {
+			keysSeen[key.Str] = struct{}{}
+			keyValues = append(keyValues, keyValue{key: key.Str, rawValue: value.Raw})
+		}
+		return true
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	slices.SortFunc(keyValues, func(a, b keyValue) int { return strings.Compare(a.key, b.key) })
+
+	buf = append(buf, '{')
+	for i, keyValue := range keyValues {
+		if i > 0 {
+			buf = append(buf, ',')
+		}
+		buf = appendJSONKey(buf, keyValue.key)
+		buf = append(buf, ':')
+		buf = append(buf, keyValue.rawValue...)
+	}
+	return append(buf, '}'), nil
+}

internal/dbunique/db_unique_test.go

@@ -22,6 +22,60 @@ func (a JobArgsStaticKind) Kind() string {
 	return a.kind
 }
 
+func TestAppendSortedObject(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Encodes", func(t *testing.T) {
+		t.Parallel()
+
+		// Compare exact bytes: equivalent JSON with different key escaping or
+		// whitespace would change the hashes of jobs inserted before upgrading.
+		for _, tt := range []struct {
+			name          string
+			encodedObject string
+			expected      string
+		}{
+			{name: "ASCIIHTML", encodedObject: `{"a\u003cb\u0026c\u003e":1}`, expected: `{"a<b&c>":1}`},
+			{name: "Colon", encodedObject: `{"x":2,":x":1}`, expected: `{":x":1,"x":2}`},
+			{name: "ControlCharacters", encodedObject: `{"line\nbreak":1,"\u0000":2}`, expected: `{"\u0000":2,"line\nbreak":1}`},
+			{name: "DEL", encodedObject: `{"\u007f":1}`, expected: "{\"\x7f\":1}"},
+			{name: "DuplicateKeys", encodedObject: `{"a":1,"\u0061":2}`, expected: `{"a":1}`},
+			{name: "EmptyArray", encodedObject: `[]`, expected: `{}`},
+			{name: "EmptyInput", encodedObject: ``, expected: `{}`},
+			{name: "EmptyKey", encodedObject: `{"a":2,"":1}`, expected: `{"":1,"a":2}`},
+			{name: "EmptyObject", encodedObject: `{}`, expected: `{}`},
+			{name: "EscapedKey", encodedObject: `{"quote\"slash\\":1}`, expected: `{"quote\"slash\\":1}`},
+			{name: "LiteralAndNestedKeys", encodedObject: `{"file.name":"a","file":{"name":"b"}}`, expected: `{"file":{"name":"b"},"file.name":"a"}`},
+			{name: "NumericKeys", encodedObject: `{"10":2,"0":1}`, expected: `{"0":1,"10":2}`},
+			{name: "RawValues", encodedObject: `{"b":null,"a":{"y":1, "x":[1, 2]}}`, expected: `{"a":{"y":1, "x":[1, 2]},"b":null}`},
+			{name: "UnicodeHTML", encodedObject: `{"\u00e9\u003c":1}`, expected: `{"é\u003c":1}`},
+			{name: "UnicodeSeparators", encodedObject: `{"\u2029":2,"\u2028":1}`, expected: `{"\u2028":1,"\u2029":2}`},
+		} {
+			t.Run(tt.name, func(t *testing.T) {
+				t.Parallel()
+
+				encoded, err := appendSortedObject(nil, []byte(tt.encodedObject))
+				require.NoError(t, err)
+				require.Equal(t, tt.expected, string(encoded))
+			})
+		}
+	})
+
+	t.Run("RejectsNonObjectArgs", func(t *testing.T) {
+		t.Parallel()
+
+		for _, encodedArgs := range []string{`null`, `[1,2]`, `"str"`, `true`, `123`} {
+			args := JobArgsStaticKind{kind: "kind"}
+			_, err := UniqueKey(&riversharedtest.TimeStub{}, &UniqueOpts{ByArgs: true}, &rivertype.JobInsertParams{
+				Args:        args,
+				EncodedArgs: []byte(encodedArgs),
+				Kind:        args.Kind(),
+			})
+			require.EqualError(t, err, "unique args must encode a JSON object", "encoded args: %s", encodedArgs)
+		}
+	})
+}
+
 func TestUniqueKey(t *testing.T) {
 	t.Parallel()
 
@@ -101,6 +155,88 @@ func TestUniqueKey(t *testing.T) {
 			uniqueOpts:   UniqueOpts{ByArgs: true},
 			expectedJSON: `&kind=worker_1&args={"Recipient":"john@example.com","Subject":"Another Test Email"}`,
 		},
+		{
+			name: "ByArgsUniqueWithUnnamedJSONTagUsesFieldName",
+			argsFunc: func() rivertype.JobArgs {
+				//nolint:tagliatelle // non-snake keys are intentional
+				type EmailJobArgs struct {
+					JobArgsStaticKind
+
+					Recipient  string `json:",omitempty" river:"unique"`
+					Subject    string `json:"subject"    river:"unique"`
+					TemplateID int
+				}
+				return EmailJobArgs{
+					JobArgsStaticKind: JobArgsStaticKind{kind: "worker_1"},
+					Recipient:         "john@example.com",
+					Subject:           "Another Test Email",
+					TemplateID:        102,
+				}
+			},
+			uniqueOpts:   UniqueOpts{ByArgs: true},
+			expectedJSON: `&kind=worker_1&args={"Recipient":"john@example.com","subject":"Another Test Email"}`,
+		},
+		{
+			name: "ByArgsWithCompatibleJSONNames",
+			argsFunc: func() rivertype.JobArgs {
+				//nolint:tagliatelle // Exercise names whose escaped sort order differs.
+				type Nested struct {
+					Amount int    `json:"a&b" river:"unique"`
+					Upper  string `json:"Y"   river:"unique"`
+				}
+				//nolint:tagliatelle // Exercise valid names whose escaping must not change hashes.
+				type Args struct {
+					JobArgsStaticKind
+
+					Dollar  string `json:"$x"     river:"unique"`
+					Nested  Nested `json:"nested"`
+					Unicode string `json:"é<"     river:"unique"`
+					Upper   string `json:"Y"      river:"unique"`
+				}
+				return Args{
+					JobArgsStaticKind: JobArgsStaticKind{kind: "worker_1"},
+					Dollar:            "dollar",
+					Nested:            Nested{Amount: 1, Upper: "nested"},
+					Unicode:           "unicode",
+					Upper:             "upper",
+				}
+			},
+			uniqueOpts: UniqueOpts{ByArgs: true},
+			// Preserve the old byte ordering even though escaped `\$x` sorts after `Y`.
+			expectedJSON: `&kind=worker_1&args={"$x":"dollar","Y":"upper","nested":{"Y":"nested","a&b":1},"é\u003c":"unicode"}`,
+		},
+		{
+			name: "ByArgsUniqueWithPathSyntaxInJSONTags",
+			argsFunc: func() rivertype.JobArgs {
+				type Nested struct {
+					Value string `json:"inner@key" river:"unique"`
+				}
+				//nolint:tagliatelle // non-snake keys are intentional
+				type PathSyntaxJobArgs struct {
+					JobArgsStaticKind
+
+					Bang    string `json:"!bang"               river:"unique"`
+					Colon   string `json:":x"                  river:"unique"`
+					Email   string `json:"alice@example.com"   river:"unique"`
+					Literal string `json:"outer.key.inner@key" river:"unique"`
+					Nested  Nested `json:"outer.key"`
+					UserID  string `json:"user.id"             river:"unique"`
+					X       string `json:"x"                   river:"unique"`
+				}
+				return PathSyntaxJobArgs{
+					JobArgsStaticKind: JobArgsStaticKind{kind: "worker_1"},
+					Bang:              "bang",
+					Colon:             "colon",
+					Email:             "email",
+					Literal:           "literal",
+					Nested:            Nested{Value: "nested"},
+					UserID:            "u1",
+					X:                 "x",
+				}
+			},
+			uniqueOpts:   UniqueOpts{ByArgs: true},
+			expectedJSON: `&kind=worker_1&args={"!bang":"bang",":x":"colon","alice@example.com":"email","outer.key":{"inner@key":"nested"},"outer.key.inner@key":"literal","user.id":"u1","x":"x"}`,
+		},
 		{
 			name: "ByArgsWithPointerToStruct",
 			argsFunc: func() rivertype.JobArgs {
@@ -344,6 +480,15 @@ func TestUniqueKey(t *testing.T) {
 			// args JSON should be sorted alphabetically:
 			expectedJSON: `&kind=worker_3&args={"count":10,"description":"A generic job without unique fields."}`,
 		},
+		{
+			name:     "ByArgsWithNoUniqueFieldsAndLiteralKeys",
+			argsFunc: func() rivertype.JobArgs { return JobArgsStaticKind{kind: "worker_3"} },
+			modifyInsertParamsFunc: func(params *rivertype.JobInsertParams) {
+				params.EncodedArgs = []byte(`{"x":"x","file.name":"file","alice@example.com":"email",":x":"colon","[x":"bracket","{x":"brace","":"empty"}`)
+			},
+			uniqueOpts:   UniqueOpts{ByArgs: true},
+			expectedJSON: `&kind=worker_3&args={"":"empty",":x":"colon","[x":"bracket","alice@example.com":"email","file.name":"file","x":"x","{x":"brace"}`,
+		},
 		{
 			name: "ByArgsWithEmptyEncodedArgs",
 			argsFunc: func() rivertype.JobArgs {