From 97dc9fb465da4e2c16a8fe50ef03450be07cfd6b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 5 Jun 2026 20:12:55 +0000 Subject: [PATCH] Add ETag optimistic concurrency for document replace Co-authored-by: Pijus Kamandulis --- api/handlers/documents.go | 22 +++++++++- api/headers/headers.go | 1 + api/tests/documents_test.go | 75 +++++++++++++++++++++++++++++++++ internal/constants/responses.go | 4 ++ 4 files changed, 101 insertions(+), 1 deletion(-) diff --git a/api/handlers/documents.go b/api/handlers/documents.go index 5f3fa68..5667851 100644 --- a/api/handlers/documents.go +++ b/api/handlers/documents.go @@ -47,6 +47,9 @@ func (h *Handlers) GetDocument(c *gin.Context) { document, status := h.dataStore.GetDocument(databaseId, collectionId, documentId) if status == datastore.StatusOk { + if etag, ok := document["_etag"].(string); ok { + c.Header(headers.ETag, etag) + } c.IndentedJSON(http.StatusOK, document) return } @@ -90,7 +93,24 @@ func (h *Handlers) ReplaceDocument(c *gin.Context) { return } - status := h.dataStore.DeleteDocument(databaseId, collectionId, documentId) + existingDocument, status := h.dataStore.GetDocument(databaseId, collectionId, documentId) + if status == datastore.StatusNotFound { + c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse) + return + } + if status != datastore.StatusOk { + c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse) + return + } + + if ifMatch := c.GetHeader(headers.IfMatch); ifMatch != "" { + if existingDocument["_etag"] != ifMatch { + c.JSON(http.StatusPreconditionFailed, constants.PreconditionFailedResponse) + return + } + } + + status = h.dataStore.DeleteDocument(databaseId, collectionId, documentId) if status == datastore.StatusNotFound { c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse) return diff --git a/api/headers/headers.go b/api/headers/headers.go index 76a5999..d69f8ac 100644 --- a/api/headers/headers.go +++ b/api/headers/headers.go @@ -6,6 +6,7 @@ const ( CosmosLsn = "x-ms-cosmos-llsn" ETag = "etag" GlobalCommittedLsn = "x-ms-global-committed-lsn" + IfMatch = "if-match" IfNoneMatch = "if-none-match" IsBatchRequest = "x-ms-cosmos-is-batch-request" IsQueryPlanRequest = "x-ms-cosmos-is-query-plan-request" diff --git a/api/tests/documents_test.go b/api/tests/documents_test.go index 8c8cc58..132aed9 100644 --- a/api/tests/documents_test.go +++ b/api/tests/documents_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "io" "net/http" "reflect" "sync" @@ -379,6 +380,80 @@ func Test_Documents(t *testing.T) { }) }) + runTestsWithPresets(t, "Test_Documents_ETag_OptimisticConcurrency", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) { + collectionClient := documents_InitializeDb(t, ts) + + t.Run("Should fail replace with incorrect etag", func(t *testing.T) { + context := context.TODO() + + item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": true} + bytes, err := json.Marshal(item) + assert.Nil(t, err) + + wrongETag := azcore.ETag("\"incorrect-etag\"") + _, err = collectionClient.ReplaceItem( + context, + azcosmos.PartitionKey{}, + "12345", + bytes, + &azcosmos.ItemOptions{IfMatchEtag: &wrongETag}, + ) + assert.NotNil(t, err) + + var respErr *azcore.ResponseError + if errors.As(err, &respErr) { + assert.Equal(t, http.StatusPreconditionFailed, respErr.StatusCode) + assert.Equal(t, "PreconditionFailed", respErr.ErrorCode) + + responseBody, readErr := io.ReadAll(respErr.RawResponse.Body) + assert.Nil(t, readErr) + assert.JSONEq(t, + `{"code":"PreconditionFailed","message":"Operation cannot be performed because one of the specified precondition is not met."}`, + string(responseBody), + ) + } else { + panic(err) + } + + document, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345") + assert.Equal(t, datastore.StatusOk, status) + assert.Equal(t, false, document["isCool"]) + }) + + t.Run("Should replace with correct etag", func(t *testing.T) { + context := context.TODO() + + readResponse, err := collectionClient.ReadItem(context, azcosmos.PartitionKey{}, "12345", nil) + assert.Nil(t, err) + assert.NotEmpty(t, readResponse.ETag) + + var item map[string]interface{} + err = json.Unmarshal(readResponse.Value, &item) + assert.Nil(t, err) + assert.Equal(t, string(readResponse.ETag), item["_etag"]) + + item["pk"] = "999" + item["isCool"] = true + bytes, err := json.Marshal(item) + assert.Nil(t, err) + + etag := readResponse.ETag + _, err = collectionClient.ReplaceItem( + context, + azcosmos.PartitionKey{}, + "12345", + bytes, + &azcosmos.ItemOptions{IfMatchEtag: &etag}, + ) + assert.Nil(t, err) + + document, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345") + assert.Equal(t, datastore.StatusOk, status) + assert.Equal(t, "999", document["pk"]) + assert.Equal(t, true, document["isCool"]) + }) + }) + runTestsWithPresets(t, "Test_Documents_TransactionalBatch", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) { collectionClient := documents_InitializeDb(t, ts) diff --git a/internal/constants/responses.go b/internal/constants/responses.go index ddf1e38..b8b55eb 100644 --- a/internal/constants/responses.go +++ b/internal/constants/responses.go @@ -35,3 +35,7 @@ var UnknownErrorResponse = gin.H{"message": "Unknown error"} var NotFoundResponse = gin.H{"message": "NotFound"} var ConflictResponse = gin.H{"message": "Conflict"} var BadRequestResponse = gin.H{"message": "BadRequest"} +var PreconditionFailedResponse = gin.H{ + "code": "PreconditionFailed", + "message": "Operation cannot be performed because one of the specified precondition is not met.", +}