[Feature] Query ForceOneShardAttributeValue (#327)

Author: ajanikow

query.go

@@ -28,16 +28,17 @@ import (
 )
 
 const (
-	keyQueryCount          = "arangodb-query-count"
-	keyQueryBatchSize      = "arangodb-query-batchSize"
-	keyQueryCache          = "arangodb-query-cache"
-	keyQueryMemoryLimit    = "arangodb-query-memoryLimit"
-	keyQueryTTL            = "arangodb-query-ttl"
-	keyQueryOptSatSyncWait = "arangodb-query-opt-satSyncWait"
-	keyQueryOptFullCount   = "arangodb-query-opt-fullCount"
-	keyQueryOptStream      = "arangodb-query-opt-stream"
-	keyQueryOptProfile     = "arangodb-query-opt-profile"
-	keyQueryOptMaxRuntime  = "arangodb-query-opt-maxRuntime"
+	keyQueryCount                       = "arangodb-query-count"
+	keyQueryBatchSize                   = "arangodb-query-batchSize"
+	keyQueryCache                       = "arangodb-query-cache"
+	keyQueryMemoryLimit                 = "arangodb-query-memoryLimit"
+	keyQueryForceOneShardAttributeValue = "arangodb-query-forceOneShardAttributeValue"
+	keyQueryTTL                         = "arangodb-query-ttl"
+	keyQueryOptSatSyncWait              = "arangodb-query-opt-satSyncWait"
+	keyQueryOptFullCount                = "arangodb-query-opt-fullCount"
+	keyQueryOptStream                   = "arangodb-query-opt-stream"
+	keyQueryOptProfile                  = "arangodb-query-opt-profile"
+	keyQueryOptMaxRuntime               = "arangodb-query-opt-maxRuntime"
 )
 
 // WithQueryCount is used to configure a context that will set the Count of a query request,
@@ -70,6 +71,11 @@ func WithQueryMemoryLimit(parent context.Context, value int64) context.Context {
 	return context.WithValue(contextOrBackground(parent), keyQueryMemoryLimit, value)
 }
 
+// WithQueryForceOneShardAttributeValue is used to configure a context that will set the ForceOneShardAttributeValue of a query request,
+func WithQueryForceOneShardAttributeValue(parent context.Context, value string) context.Context {
+	return context.WithValue(contextOrBackground(parent), keyQueryForceOneShardAttributeValue, value)
+}
+
 // WithQueryTTL is used to configure a context that will set the TTL of a query request,
 func WithQueryTTL(parent context.Context, value time.Duration) context.Context {
 	return context.WithValue(contextOrBackground(parent), keyQueryTTL, value)
@@ -179,6 +185,9 @@ type queryRequest struct {
 		// MaxRuntime specify the timeout which can be used to kill a query on the server after the specified
 		// amount in time. The timeout value is specified in seconds. A value of 0 means no timeout will be enforced.
 		MaxRuntime float64 `json:"maxRuntime,omitempty"`
+		// ForceOneShardAttributeValue This query option can be used in complex queries in case the query optimizer cannot
+		// automatically detect that the query can be limited to only a single server (e.g. in a disjoint smart graph case).
+		ForceOneShardAttributeValue *string `json:"forceOneShardAttributeValue,omitempty"`
 	} `json:"options,omitempty"`
 }
 
@@ -207,6 +216,11 @@ func (q *queryRequest) applyContextSettings(ctx context.Context) {
 			q.MemoryLimit = value
 		}
 	}
+	if rawValue := ctx.Value(keyQueryForceOneShardAttributeValue); rawValue != nil {
+		if value, ok := rawValue.(string); ok {
+			q.Options.ForceOneShardAttributeValue = &value
+		}
+	}
 	if rawValue := ctx.Value(keyQueryTTL); rawValue != nil {
 		if value, ok := rawValue.(time.Duration); ok {
 			q.TTL = value.Seconds()

test/query_test.go

@@ -229,3 +229,41 @@ func TestProfileQuery(t *testing.T) {
 		}
 	})
 }
+
+// TestForceOneShardAttributeValue test ForceOneShardAttributeValue query attribute.
+func TestForceOneShardAttributeValue(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c := createClientFromEnv(t, true)
+
+	EnsureVersion(t, ctx, c).CheckVersion(MinimumVersion("3.9.0")).Cluster().Enterprise()
+
+	db := ensureDatabase(ctx, c, "force_one_shard_attribute_value", nil, t)
+
+	db, clean := prepareQueryDatabase(t, ctx, c, "force_one_shard_attribute_value")
+	defer clean(t)
+
+	// Setup tests
+	tests := []profileQueryTest{
+		{
+			Query: "FOR d IN books SORT d.Title RETURN d",
+		},
+		{
+			Query: "FOR d IN books FILTER d.Title==@title SORT d.Title RETURN d",
+			BindVars: map[string]interface{}{
+				"title": "Book 16",
+			},
+		},
+	}
+
+	t.Run("With ForceOneShardAttributeValue", func(t *testing.T) {
+		for i, test := range tests {
+			t.Run(fmt.Sprintf("Run %d", i), func(t *testing.T) {
+				nCtx := driver.WithQueryForceOneShardAttributeValue(ctx, "value")
+				_, err := db.Query(nCtx, test.Query, test.BindVars)
+				require.NoError(t, err)
+			})
+		}
+	})
+}