Initial RNTBD server implementation

api/config/config.go

@@ -23,6 +23,7 @@ const (
 func ParseFlags() ServerConfig {
 	host := flag.String("Host", "localhost", "Hostname")
 	port := flag.Int("Port", 8081, "Listen port")
+	rntbdPort := flag.Int("RntbdPort", 10000, "RNTBD listen port")
 	explorerPath := flag.String("ExplorerDir", "", "Path to cosmos-explorer files")
 	tlsCertificatePath := flag.String("Cert", "", "Hostname")
 	tlsCertificateKey := flag.String("CertKey", "", "Hostname")
@@ -35,6 +36,7 @@ func ParseFlags() ServerConfig {
 	flag.Var(logLevel, "LogLevel", fmt.Sprintf("Sets the logging level %s", logLevel.AllowedValuesList()))
 	dataStore := NewEnumValue("json", []string{DataStoreJson, DataStoreBadger})
 	flag.Var(dataStore, "DataStore", fmt.Sprintf("Sets the data store %s", dataStore.AllowedValuesList()))
+	enableRntbd := flag.Bool("ExperimentalEnableRntbd", false, "EXPERIMENTAL: Enable RNTBD (CosmosDB Direct Connection Mode)")
 
 	flag.Parse()
 	setFlagsFromEnvironment()
@@ -42,6 +44,7 @@ func ParseFlags() ServerConfig {
 	config := ServerConfig{}
 	config.Host = *host
 	config.Port = *port
+	config.RntbdPort = *rntbdPort
 	config.ExplorerPath = *explorerPath
 	config.TLS_CertificatePath = *tlsCertificatePath
 	config.TLS_CertificateKey = *tlsCertificateKey
@@ -52,6 +55,7 @@ func ParseFlags() ServerConfig {
 	config.AccountKey = *accountKey
 	config.LogLevel = logLevel.value
 	config.DataStore = dataStore.value
+	config.EnableRntbd = *enableRntbd
 
 	config.PopulateCalculatedFields()
 
@@ -62,6 +66,7 @@ func (c *ServerConfig) PopulateCalculatedFields() {
 	c.DatabaseAccount = c.Host
 	c.DatabaseDomain = c.Host
 	c.DatabaseEndpoint = fmt.Sprintf("https://%s:%d/", c.Host, c.Port)
+	c.RntbdEndpoint = fmt.Sprintf("rntbd://%s:%d/", c.Host, c.RntbdPort)
 	c.ExplorerBaseUrlLocation = ExplorerBaseUrlLocation
 
 	switch c.LogLevel {

api/config/models.go

@@ -4,10 +4,12 @@ type ServerConfig struct {
 	DatabaseAccount  string `json:"databaseAccount"`
 	DatabaseDomain   string `json:"databaseDomain"`
 	DatabaseEndpoint string `json:"databaseEndpoint"`
+	RntbdEndpoint    string `json:"rntbdEndpoint"`
 	AccountKey       string `json:"accountKey"`
 
 	ExplorerPath            string `json:"explorerPath"`
 	Port                    int    `json:"port"`
+	RntbdPort               int    `json:"rntbdPort"`
 	Host                    string `json:"host"`
 	TLS_CertificatePath     string `json:"tlsCertificatePath"`
 	TLS_CertificateKey      string `json:"tlsCertificateKey"`
@@ -17,6 +19,7 @@ type ServerConfig struct {
 	DisableTls              bool   `json:"disableTls"`
 	LogLevel                string `json:"logLevel"`
 	ExplorerBaseUrlLocation string `json:"explorerBaseUrlLocation"`
+	EnableRntbd             bool   `json:"enableRntbd"`
 
 	DataStore string `json:"dataStore"`
 }

api/handlers/server_info.go

@@ -3,6 +3,7 @@ package handlers
 import (
 	"fmt"
 	"net/http"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 )
@@ -41,3 +42,46 @@ func (h *Handlers) GetServerInfo(c *gin.Context) {
 		"queryEngineConfiguration": "{\"allowNewKeywords\":true,\"maxJoinsPerSqlQuery\":10,\"maxQueryRequestTimeoutFraction\":0.9,\"maxSqlQueryInputLength\":524288,\"maxUdfRefPerSqlQuery\":10,\"queryMaxInMemorySortDocumentCount\":-1000,\"spatialMaxGeometryPointCount\":256,\"sqlAllowNonFiniteNumbers\":false,\"sqlDisableOptimizationFlags\":0,\"enableSpatialIndexing\":true,\"maxInExpressionItemsCount\":2147483647,\"maxLogicalAndPerSqlQuery\":2147483647,\"maxLogicalOrPerSqlQuery\":2147483647,\"maxSpatialQueryCells\":2147483647,\"sqlAllowAggregateFunctions\":true,\"sqlAllowGroupByClause\":true,\"sqlAllowLike\":true,\"sqlAllowSubQuery\":true,\"sqlAllowScalarSubQuery\":true,\"sqlAllowTop\":true}",
 	})
 }
+
+type Address struct {
+	IsPrimary                     bool   `json:"isPrimary"`
+	PhyscialUri                   string `json:"physcialUri"`
+	IsAuxiliary                   bool   `json:"isAuxiliary"`
+	PartitionTargetReplicaSetSize int    `json:"partitionTargetReplicaSetSize"`
+	Protocol                      string `json:"protocol"`
+	PartitionKeyRangeId           string `json:"partitionKeyRangeId"`
+	PartitionIndex                string `json:"partitionIndex"`
+}
+
+func (h *Handlers) GetAddresses(c *gin.Context) {
+	addresses := []Address{}
+
+	if h.config.EnableRntbd {
+		addresses = append(addresses, Address{
+			IsPrimary:                     true,
+			PhyscialUri:                   h.config.RntbdEndpoint,
+			IsAuxiliary:                   false,
+			PartitionTargetReplicaSetSize: 1,
+			Protocol:                      "rntbd",
+			PartitionKeyRangeId:           "0",
+			PartitionIndex:                "0@0",
+		})
+	}
+
+	if !strings.Contains(c.Request.RequestURI, "protocol%20eq%20rntbd") {
+		addresses = append(addresses, Address{
+			IsPrimary:                     true,
+			PhyscialUri:                   h.config.DatabaseEndpoint,
+			IsAuxiliary:                   false,
+			PartitionTargetReplicaSetSize: 1,
+			Protocol:                      "https",
+			PartitionKeyRangeId:           "0",
+			PartitionIndex:                "0@0",
+		})
+	}
+
+	c.IndentedJSON(http.StatusOK, gin.H{
+		"Addresss": addresses,
+		"_count":   len(addresses),
+	})
+}

api/headers/headers.go

@@ -21,4 +21,10 @@ const (
 	// 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
+
+	// I kinda don't use these, but I've seen them in the wild xd
+	SupportedCapabilities     = "x-ms-cosmos-sdk-supportedcapabilities"
+	ClientRetryAttemptCount   = "x-ms-client-retry-attempt-count"
+	RemainingTimeInMsOnClient = "x-ms-remaining-time-in-ms-on-client"
+	ConsistencyLevel          = "x-ms-consistency-level"
 )

api/router.go

@@ -31,6 +31,7 @@ func (s *ApiServer) CreateRouter(dataStore datastore.DataStore) {
 
 	router := gin.Default(func(e *gin.Engine) {
 		e.RedirectTrailingSlash = false
+		e.RemoveExtraSlash = true
 	})
 
 	if s.config.LogLevel == "debug" {
@@ -79,6 +80,7 @@ func (s *ApiServer) CreateRouter(dataStore datastore.DataStore) {
 
 	router.GET("/offers", handlers.GetOffers)
 	router.GET("/", routeHandlers.GetServerInfo)
+	router.GET("//addresses", routeHandlers.GetAddresses)
 
 	router.GET("/cosmium/export", routeHandlers.CosmiumExport)