Fix query creation via explorer; Extract header names to constants

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,

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,

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"]
+	// Handle query plan requests
-	if query != nil {
+	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,

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")
+		authHeader := c.Request.Header.Get(headers.Authorization)
-		date := c.Request.Header.Get("x-ms-date")
+		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
 	}

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,

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(headers.ETag, "\"420\"")
-		c.Header("lsn", "420")
+		c.Header(headers.LSN, "420")
-		c.Header("x-ms-cosmos-llsn", "420")
+		c.Header(headers.CosmosLsn, "420")
-		c.Header("x-ms-global-committed-lsn", "420")
+		c.Header(headers.GlobalCommittedLsn, "420")
-		c.Header("x-ms-item-count", fmt.Sprintf("%d", len(partitionKeyRanges)))
+		c.Header(headers.ItemCount, fmt.Sprintf("%d", len(partitionKeyRanges)))
 
 		collectionRid := collectionId
 		collection, _ := h.dataStore.GetCollection(databaseId, collectionId)

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
 	}

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
 	}

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
 	}

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
+)

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(headers.XDate, date)
-		req.Header.Add("authorization", "sig="+url.QueryEscape(signature))
+		req.Header.Add(headers.Authorization, "sig="+url.QueryEscape(signature))
 		res, err := httpClient.Do(req)
 
 		assert.Nil(t, err)