diff --git a/go/store/prolly/tree/json_chunker.go b/go/store/prolly/tree/json_chunker.go index 10c8707b5f3..6c9544a838b 100644 --- a/go/store/prolly/tree/json_chunker.go +++ b/go/store/prolly/tree/json_chunker.go @@ -130,7 +130,7 @@ func (j *JsonChunker) Done(ctx context.Context) (Node, error) { // When inserting into the beginning of an object or array, we need to add an extra comma. // We could track then in the chunker, but it's easier to just check the next part of JSON to determine // whether we need the comma. - if j.jScanner.currentPath.getScannerState() == endOfValue && jsonBytes[0] != '}' && jsonBytes[0] != ']' && jsonBytes[0] != ',' { + if j.jScanner.currentPath.getScannerState() == endOfValue && len(jsonBytes) > 0 && jsonBytes[0] != '}' && jsonBytes[0] != ']' && jsonBytes[0] != ',' { j.appendJsonToBuffer([]byte(",")) } // Append the rest of the JsonCursor, then continue until we either exhaust the cursor, or we coincide with a boundary from the original tree. diff --git a/go/store/prolly/tree/json_indexed_document.go b/go/store/prolly/tree/json_indexed_document.go index d41b4683676..0c8c3eb38c2 100644 --- a/go/store/prolly/tree/json_indexed_document.go +++ b/go/store/prolly/tree/json_indexed_document.go @@ -144,7 +144,6 @@ func (i IndexedJsonDocument) Lookup(ctx context.Context, pathString string) (res } func (i IndexedJsonDocument) tryLookup(ctx context.Context, pathString string) (sql.JSONWrapper, error) { - path, err := jsonPathElementsFromMySQLJsonPath([]byte(pathString)) if err != nil { return nil, err @@ -386,22 +385,130 @@ func (i IndexedJsonDocument) tryRemove(ctx context.Context, path string) (types. return NewIndexedJsonDocument(ctx, newRoot, i.m.NodeStore), true, nil } -// Set is not yet implemented, so we call it on a types.JSONDocument instead. -func (i IndexedJsonDocument) Set(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) { - v, err := i.ToInterface() +// Set implements types.MutableJSON +func (i IndexedJsonDocument) Set(ctx context.Context, path string, val sql.JSONWrapper) (result types.MutableJSON, changed bool, err error) { + err = tryWithFallback( + ctx, + i, + func() error { + result, changed, err = i.trySet(ctx, path, val) + return err + }, + func(jsonDocument types.JSONDocument) error { + result, changed, err = jsonDocument.Set(ctx, path, val) + return err + }) + return result, changed, err +} + +func (i IndexedJsonDocument) trySet(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) { + keyPath, err := jsonPathElementsFromMySQLJsonPath([]byte(path)) if err != nil { return nil, false, err } - return types.JSONDocument{Val: v}.Set(ctx, path, val) + + jsonCursor, found, err := newJsonCursor(ctx, i.m.NodeStore, i.m.Root, keyPath, false) + if err != nil { + return nil, false, err + } + + // The supplied path may be 0-indexing into a scalar, which is the same as referencing the scalar. Remove + // the index and try again. + for !found && keyPath.size() > jsonCursor.jsonScanner.currentPath.size() { + lastKeyPathElement := keyPath.getLastPathElement() + if !lastKeyPathElement.isArrayIndex || lastKeyPathElement.getArrayIndex() != 0 { + // The key does not exist in the document. + break + } + + keyPath.pop() + found = compareJsonLocations(keyPath, jsonCursor.jsonScanner.currentPath) == 0 + } + + if found { + return i.replaceIntoCursor(ctx, keyPath, jsonCursor, val) + } else { + return i.insertIntoCursor(ctx, keyPath, jsonCursor, val) + } } -// Replace is not yet implemented, so we call it on a types.JSONDocument instead. -func (i IndexedJsonDocument) Replace(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) { - v, err := i.ToInterface() +// Replace implements types.MutableJSON +func (i IndexedJsonDocument) Replace(ctx context.Context, path string, val sql.JSONWrapper) (result types.MutableJSON, changed bool, err error) { + err = tryWithFallback( + ctx, + i, + func() error { + result, changed, err = i.tryReplace(ctx, path, val) + return err + }, + func(jsonDocument types.JSONDocument) error { + result, changed, err = jsonDocument.Replace(ctx, path, val) + return err + }) + return result, changed, err +} + +func (i IndexedJsonDocument) tryReplace(ctx context.Context, path string, val sql.JSONWrapper) (types.MutableJSON, bool, error) { + keyPath, err := jsonPathElementsFromMySQLJsonPath([]byte(path)) + if err != nil { + return nil, false, err + } + + jsonCursor, found, err := newJsonCursor(ctx, i.m.NodeStore, i.m.Root, keyPath, false) + if err != nil { + return nil, false, err + } + + // The supplied path may be 0-indexing into a scalar, which is the same as referencing the scalar. Remove + // the index and try again. + for !found && keyPath.size() > jsonCursor.jsonScanner.currentPath.size() { + lastKeyPathElement := keyPath.getLastPathElement() + if !lastKeyPathElement.isArrayIndex || lastKeyPathElement.getArrayIndex() != 0 { + // The key does not exist in the document. + return i, false, nil + } + + keyPath.pop() + found = compareJsonLocations(keyPath, jsonCursor.jsonScanner.currentPath) == 0 + } + + if !found { + // The key does not exist in the document. + return i, false, nil + } + + return i.replaceIntoCursor(ctx, keyPath, jsonCursor, val) +} + +func (i IndexedJsonDocument) replaceIntoCursor(ctx context.Context, keyPath jsonLocation, jsonCursor *JsonCursor, val sql.JSONWrapper) (types.MutableJSON, bool, error) { + + // The cursor is now pointing to the start of the value being replaced. + jsonChunker, err := newJsonChunker(ctx, jsonCursor, i.m.NodeStore) if err != nil { return nil, false, err } - return types.JSONDocument{Val: v}.Replace(ctx, path, val) + + // Advance the cursor to the end of the value being removed. + keyPath.setScannerState(endOfValue) + _, err = jsonCursor.AdvanceToLocation(ctx, keyPath, false) + if err != nil { + return nil, false, err + } + + insertedValueBytes, err := types.MarshallJson(val) + if err != nil { + return nil, false, err + } + + jsonChunker.appendJsonToBuffer(insertedValueBytes) + jsonChunker.processBuffer(ctx) + + newRoot, err := jsonChunker.Done(ctx) + if err != nil { + return nil, false, err + } + + return NewIndexedJsonDocument(ctx, newRoot, i.m.NodeStore), true, nil } // ArrayInsert is not yet implemented, so we call it on a types.JSONDocument instead. diff --git a/go/store/prolly/tree/json_indexed_document_test.go b/go/store/prolly/tree/json_indexed_document_test.go index 143857bc731..5a9b7e0e0a0 100644 --- a/go/store/prolly/tree/json_indexed_document_test.go +++ b/go/store/prolly/tree/json_indexed_document_test.go @@ -259,6 +259,28 @@ func TestIndexedJsonDocument_Extract(t *testing.T) { jsontests.RunJsonTests(t, testCases) } +func TestIndexedJsonDocument_Replace(t *testing.T) { + ctx := context.Background() + ns := NewTestNodeStore() + convertToIndexedJsonDocument := func(t *testing.T, s interface{}) interface{} { + return newIndexedJsonDocumentFromValue(t, ctx, ns, s) + } + + testCases := jsontests.JsonReplaceTestCases(t, convertToIndexedJsonDocument) + jsontests.RunJsonTests(t, testCases) +} + +func TestIndexedJsonDocument_Set(t *testing.T) { + ctx := context.Background() + ns := NewTestNodeStore() + convertToIndexedJsonDocument := func(t *testing.T, s interface{}) interface{} { + return newIndexedJsonDocumentFromValue(t, ctx, ns, s) + } + + testCases := jsontests.JsonSetTestCases(t, convertToIndexedJsonDocument) + jsontests.RunJsonTests(t, testCases) +} + func TestIndexedJsonDocument_Value(t *testing.T) { ctx := context.Background() ns := NewTestNodeStore()