Add ETag optimistic concurrency for document replace

Co-authored-by: Pijus Kamandulis <pikami@users.noreply.github.com>

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

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"

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)
 

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.",
+}