feat: update documents by function

Author: Ja7ad

index_document.go

@@ -223,6 +223,27 @@ func (i *index) UpdateDocumentsNdjsonInBatchesWithContext(ctx context.Context, d
 	return i.updateDocumentsNdjsonFromReaderInBatches(ctx, bytes.NewReader(documents), batchSize, primaryKey...)
 }
 
+func (i *index) UpdateDocumentsByFunction(req *UpdateDocumentByFunctionRequest) (*TaskInfo, error) {
+	return i.UpdateDocumentsByFunctionWithContext(context.Background(), req)
+}
+
+func (i *index) UpdateDocumentsByFunctionWithContext(ctx context.Context, req *UpdateDocumentByFunctionRequest) (*TaskInfo, error) {
+	resp := new(TaskInfo)
+	r := &internalRequest{
+		endpoint:            "/indexes/" + i.uid + "/documents/edit",
+		method:              http.MethodPost,
+		withRequest:         req,
+		withResponse:        resp,
+		contentType:         contentTypeJSON,
+		acceptedStatusCodes: []int{http.StatusAccepted},
+		functionName:        "UpdateDocumentsByFunction",
+	}
+	if err := i.client.executeRequest(ctx, r); err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
 func (i *index) GetDocument(identifier string, request *DocumentQuery, documentPtr interface{}) error {
 	return i.GetDocumentWithContext(context.Background(), identifier, request, documentPtr)
 }

index_document_test.go

@@ -1658,3 +1658,35 @@ func TestIndex_DeleteDocumentsByFilter(t *testing.T) {
 		})
 	}
 }
+
+func TestIndex_UpdateDocumentsByFunction(t *testing.T) {
+	c := setup(t, "")
+
+	exp := c.ExperimentalFeatures()
+	exp.SetEditDocumentsByFunction(true)
+	res, err := exp.Update()
+	require.NoError(t, err)
+	require.True(t, res.EditDocumentsByFunction)
+
+	idx := setupMovieIndex(t, c)
+
+	t.Run("Test Upper Case and Add Sparkles around Movie Titles", func(t *testing.T) {
+		task, err := idx.UpdateDocumentsByFunction(&UpdateDocumentByFunctionRequest{
+			Filter:   "id > 3000",
+			Function: "doc.title = `✨ ${doc.title.to_upper()} ✨`",
+		})
+		require.NoError(t, err)
+		testWaitForTask(t, idx, task)
+	})
+
+	t.Run("Test User-defined Context", func(t *testing.T) {
+		task, err := idx.UpdateDocumentsByFunction(&UpdateDocumentByFunctionRequest{
+			Context: map[string]interface{}{
+				"idmax": 50,
+			},
+			Function: "if doc.id >= context.idmax {\n\t\t    doc = ()\n\t\t  } else {\n\t\t\t  doc.title = `✨ ${doc.title} ✨`\n\t\t\t}",
+		})
+		require.NoError(t, err)
+		testWaitForTask(t, idx, task)
+	})
+}

index_interface.go

@@ -153,6 +153,12 @@ type DocumentManager interface {
 	// UpdateDocumentsNdjsonInBatchesWithContext updates documents in the index from a NDJSON byte array in batches of specified size using the provided context for cancellation.
 	UpdateDocumentsNdjsonInBatchesWithContext(ctx context.Context, documents []byte, batchsize int, primaryKey ...string) ([]TaskInfo, error)
 
+	// UpdateDocumentsByFunction update documents by using function
+	UpdateDocumentsByFunction(req *UpdateDocumentByFunctionRequest) (*TaskInfo, error)
+
+	// UpdateDocumentsByFunctionWithContext update documents by using function then provided context for cancellation.
+	UpdateDocumentsByFunctionWithContext(ctx context.Context, req *UpdateDocumentByFunctionRequest) (*TaskInfo, error)
+
 	// DeleteDocument deletes a single document from the index by identifier.
 	DeleteDocument(identifier string) (*TaskInfo, error)
 

main_test.go

@@ -210,6 +210,30 @@ func setUpBasicIndex(sv ServiceManager, indexUID string) {
 	}
 }
 
+func setupMovieIndex(t *testing.T, client ServiceManager) IndexManager {
+	t.Helper()
+
+	idx := client.Index("indexUID")
+
+	testdata, err := os.Open("./testdata/movies.json")
+	require.NoError(t, err)
+	defer testdata.Close()
+
+	tests := make([]map[string]interface{}, 0)
+
+	require.NoError(t, json.NewDecoder(testdata).Decode(&tests))
+
+	task, err := idx.AddDocuments(tests)
+	require.NoError(t, err)
+	testWaitForTask(t, idx, task)
+
+	task, err = idx.UpdateFilterableAttributes(&[]string{"id"})
+	require.NoError(t, err)
+	testWaitForTask(t, idx, task)
+
+	return idx
+}
+
 func setUpIndexForFaceting(client ServiceManager) {
 	idx := client.Index("indexUID")