diff --git a/api/handlers/collections.go b/api/handlers/collections.go index 58774a9..6cd96b1 100644 --- a/api/handlers/collections.go +++ b/api/handlers/collections.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" ) @@ -16,7 +17,7 @@ func (h *Handlers) GetAllCollections(c *gin.Context) { if status == datastore.StatusOk { database, _ := h.dataStore.GetDatabase(databaseId) - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(collections))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(collections))) c.IndentedJSON(http.StatusOK, gin.H{ "_rid": database.ResourceID, "DocumentCollections": collections, diff --git a/api/handlers/databases.go b/api/handlers/databases.go index 2f48d77..41d6c6a 100644 --- a/api/handlers/databases.go +++ b/api/handlers/databases.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" ) @@ -12,7 +13,7 @@ import ( func (h *Handlers) GetAllDatabases(c *gin.Context) { databases, status := h.dataStore.GetAllDatabases() if status == datastore.StatusOk { - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(databases))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(databases))) c.IndentedJSON(http.StatusOK, gin.H{ "_rid": "", "Databases": databases, diff --git a/api/handlers/documents.go b/api/handlers/documents.go index 71957ea..91c9021 100644 --- a/api/handlers/documents.go +++ b/api/handlers/documents.go @@ -9,6 +9,7 @@ import ( jsonpatch "github.com/cosmiumdev/json-patch/v5" "github.com/gin-gonic/gin" apimodels "github.com/pikami/cosmium/api/api_models" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/converters" "github.com/pikami/cosmium/internal/datastore" @@ -26,7 +27,7 @@ func (h *Handlers) GetAllDocuments(c *gin.Context) { if status == datastore.StatusOk { collection, _ := h.dataStore.GetCollection(databaseId, collectionId) - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(documents))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(documents))) c.IndentedJSON(http.StatusOK, gin.H{ "_rid": collection.ID, "Documents": documents, @@ -189,7 +190,7 @@ func (h *Handlers) DocumentsPost(c *gin.Context) { collectionId := c.Param("collId") // Handle batch requests - isBatchRequest, _ := strconv.ParseBool(c.GetHeader("x-ms-cosmos-is-batch-request")) + isBatchRequest, _ := strconv.ParseBool(c.GetHeader(headers.IsBatchRequest)) if isBatchRequest { h.handleBatchRequest(c) return @@ -201,8 +202,17 @@ func (h *Handlers) DocumentsPost(c *gin.Context) { return } - query := requestBody["query"] - if query != nil { + // Handle query plan requests + isQueryPlanRequest, _ := strconv.ParseBool(c.GetHeader(headers.IsQueryPlanRequest)) + if isQueryPlanRequest { + c.IndentedJSON(http.StatusOK, constants.QueryPlanResponse) + return + } + + // Handle query requests + isQueryRequest, _ := strconv.ParseBool(c.GetHeader(headers.IsQuery)) + isQueryRequestAltHeader, _ := strconv.ParseBool(c.GetHeader(headers.Query)) + if isQueryRequest || isQueryRequestAltHeader { h.handleDocumentQuery(c, requestBody) return } @@ -212,7 +222,7 @@ func (h *Handlers) DocumentsPost(c *gin.Context) { return } - isUpsert, _ := strconv.ParseBool(c.GetHeader("x-ms-documentdb-is-upsert")) + isUpsert, _ := strconv.ParseBool(c.GetHeader(headers.IsUpsert)) if isUpsert { h.dataStore.DeleteDocument(databaseId, collectionId, requestBody["id"].(string)) } @@ -247,11 +257,6 @@ func (h *Handlers) handleDocumentQuery(c *gin.Context, requestBody map[string]in databaseId := c.Param("databaseId") collectionId := c.Param("collId") - if c.GetHeader("x-ms-cosmos-is-query-plan-request") != "" { - c.IndentedJSON(http.StatusOK, constants.QueryPlanResponse) - return - } - var queryParameters map[string]interface{} if paramsArray, ok := requestBody["parameters"].([]interface{}); ok { queryParameters = parametersToMap(paramsArray) @@ -266,7 +271,7 @@ func (h *Handlers) handleDocumentQuery(c *gin.Context, requestBody map[string]in } collection, _ := h.dataStore.GetCollection(databaseId, collectionId) - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(docs))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(docs))) c.IndentedJSON(http.StatusOK, gin.H{ "_rid": collection.ResourceID, "Documents": docs, diff --git a/api/handlers/middleware/authentication.go b/api/handlers/middleware/authentication.go index 8aaa79c..9de888d 100644 --- a/api/handlers/middleware/authentication.go +++ b/api/handlers/middleware/authentication.go @@ -6,6 +6,7 @@ import ( "github.com/gin-gonic/gin" "github.com/pikami/cosmium/api/config" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/authentication" "github.com/pikami/cosmium/internal/logger" ) @@ -22,8 +23,8 @@ func Authentication(config *config.ServerConfig) gin.HandlerFunc { resourceType := urlToResourceType(requestUrl) resourceId := requestToResourceId(c) - authHeader := c.Request.Header.Get("authorization") - date := c.Request.Header.Get("x-ms-date") + authHeader := c.Request.Header.Get(headers.Authorization) + date := c.Request.Header.Get(headers.XDate) expectedSignature := authentication.GenerateSignature( c.Request.Method, resourceType, resourceId, date, config.AccountKey) @@ -85,7 +86,7 @@ func requestToResourceId(c *gin.Context) string { resourceId += "/udfs/" + udfId } - isFeed := c.Request.Header.Get("A-Im") == "Incremental Feed" + isFeed := c.Request.Header.Get(headers.AIM) == "Incremental Feed" if resourceType == "pkranges" && isFeed { resourceId = collId } diff --git a/api/handlers/offers.go b/api/handlers/offers.go index c03cb05..fe031ed 100644 --- a/api/handlers/offers.go +++ b/api/handlers/offers.go @@ -4,10 +4,11 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" ) func GetOffers(c *gin.Context) { - c.Header("x-ms-item-count", "0") + c.Header(headers.ItemCount, "0") c.IndentedJSON(http.StatusOK, gin.H{ "_rid": "", "_count": 0, diff --git a/api/handlers/partition_key_ranges.go b/api/handlers/partition_key_ranges.go index d22b697..299f4d1 100644 --- a/api/handlers/partition_key_ranges.go +++ b/api/handlers/partition_key_ranges.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" "github.com/pikami/cosmium/internal/resourceid" @@ -14,18 +15,18 @@ func (h *Handlers) GetPartitionKeyRanges(c *gin.Context) { databaseId := c.Param("databaseId") collectionId := c.Param("collId") - if c.Request.Header.Get("if-none-match") != "" { + if c.Request.Header.Get(headers.IfNoneMatch) != "" { c.AbortWithStatus(http.StatusNotModified) return } partitionKeyRanges, status := h.dataStore.GetPartitionKeyRanges(databaseId, collectionId) if status == datastore.StatusOk { - c.Header("etag", "\"420\"") - c.Header("lsn", "420") - c.Header("x-ms-cosmos-llsn", "420") - c.Header("x-ms-global-committed-lsn", "420") - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(partitionKeyRanges))) + c.Header(headers.ETag, "\"420\"") + c.Header(headers.LSN, "420") + c.Header(headers.CosmosLsn, "420") + c.Header(headers.GlobalCommittedLsn, "420") + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(partitionKeyRanges))) collectionRid := collectionId collection, _ := h.dataStore.GetCollection(databaseId, collectionId) diff --git a/api/handlers/stored_procedures.go b/api/handlers/stored_procedures.go index a500f13..a72b1af 100644 --- a/api/handlers/stored_procedures.go +++ b/api/handlers/stored_procedures.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" ) @@ -16,7 +17,7 @@ func (h *Handlers) GetAllStoredProcedures(c *gin.Context) { sps, status := h.dataStore.GetAllStoredProcedures(databaseId, collectionId) if status == datastore.StatusOk { - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(sps))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(sps))) c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "StoredProcedures": sps, "_count": len(sps)}) return } diff --git a/api/handlers/triggers.go b/api/handlers/triggers.go index 0dfd30d..c2e2112 100644 --- a/api/handlers/triggers.go +++ b/api/handlers/triggers.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" ) @@ -16,7 +17,7 @@ func (h *Handlers) GetAllTriggers(c *gin.Context) { triggers, status := h.dataStore.GetAllTriggers(databaseId, collectionId) if status == datastore.StatusOk { - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(triggers))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(triggers))) c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "Triggers": triggers, "_count": len(triggers)}) return } diff --git a/api/handlers/user_defined_functions.go b/api/handlers/user_defined_functions.go index 4cb77f0..3b190ab 100644 --- a/api/handlers/user_defined_functions.go +++ b/api/handlers/user_defined_functions.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/gin-gonic/gin" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/constants" "github.com/pikami/cosmium/internal/datastore" ) @@ -16,7 +17,7 @@ func (h *Handlers) GetAllUserDefinedFunctions(c *gin.Context) { udfs, status := h.dataStore.GetAllUserDefinedFunctions(databaseId, collectionId) if status == datastore.StatusOk { - c.Header("x-ms-item-count", fmt.Sprintf("%d", len(udfs))) + c.Header(headers.ItemCount, fmt.Sprintf("%d", len(udfs))) c.IndentedJSON(http.StatusOK, gin.H{"_rid": "", "UserDefinedFunctions": udfs, "_count": len(udfs)}) return } diff --git a/api/headers/headers.go b/api/headers/headers.go new file mode 100644 index 0000000..836d59e --- /dev/null +++ b/api/headers/headers.go @@ -0,0 +1,20 @@ +package headers + +const ( + AIM = "A-Im" + Authorization = "authorization" + CosmosLsn = "x-ms-cosmos-llsn" + ETag = "etag" + GlobalCommittedLsn = "x-ms-global-committed-lsn" + IfNoneMatch = "if-none-match" + IsBatchRequest = "x-ms-cosmos-is-batch-request" + IsQueryPlanRequest = "x-ms-cosmos-is-query-plan-request" + IsUpsert = "x-ms-documentdb-is-upsert" + ItemCount = "x-ms-item-count" + LSN = "lsn" + XDate = "x-ms-date" + + // Kinda retarded, but what can I do ¯\_(ツ)_/¯ + IsQuery = "x-ms-documentdb-isquery" // Sent from python sdk and web explorer + Query = "x-ms-documentdb-query" // Sent from Go sdk +) diff --git a/api/tests/documents_trailingslash_test.go b/api/tests/documents_trailingslash_test.go index 7e7fa26..0818975 100644 --- a/api/tests/documents_trailingslash_test.go +++ b/api/tests/documents_trailingslash_test.go @@ -8,6 +8,7 @@ import ( "time" "github.com/pikami/cosmium/api/config" + "github.com/pikami/cosmium/api/headers" "github.com/pikami/cosmium/internal/authentication" "github.com/stretchr/testify/assert" ) @@ -26,8 +27,8 @@ func Test_Documents_Read_Trailing_Slash(t *testing.T) { signature := authentication.GenerateSignature("GET", "docs", path, date, config.DefaultAccountKey) httpClient := &http.Client{} req, _ := http.NewRequest("GET", testUrl, nil) - req.Header.Add("x-ms-date", date) - req.Header.Add("authorization", "sig="+url.QueryEscape(signature)) + req.Header.Add(headers.XDate, date) + req.Header.Add(headers.Authorization, "sig="+url.QueryEscape(signature)) res, err := httpClient.Do(req) assert.Nil(t, err)