Add v2 test for getting collection shards

informalict · informalict · commit 77de3ae6cf3d · 2021-09-10T06:20:30.000+02:00
diff --git a/test/collection_test.go b/test/collection_test.go
@@ -966,8 +966,8 @@ func TestCollectionWriteConcernSetPropInvalid(t *testing.T) {
 		minRepl, prop.WriteConcern)
 }
 
-// TestCollectionShards creates a collection and gets the shards' information.
-func TestCollectionShards(t *testing.T) {
+// Test_CollectionShards creates a collection and gets the shards' information.
+func Test_CollectionShards(t *testing.T) {
 	if getTestMode() != testModeCluster {
 		t.Skipf("Not a cluster mode")
 	}
@@ -990,5 +990,5 @@ func TestCollectionShards(t *testing.T) {
 		require.Lenf(t, dbServers, 2, "expected 2 DB servers for the shard")
 		leaders = append(leaders, dbServers[0])
 	}
-	assert.NotEqualf(t, leaders[0], leaders[1], "the leader can not be on the same server")
+	assert.NotEqualf(t, leaders[0], leaders[1], "the leader shard can not be on the same server")
 }
diff --git a/v2/arangodb/collection.go b/v2/arangodb/collection.go
@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2020 ArangoDB GmbH, Cologne, Germany
+// Copyright 2020-2021 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 // Author Adam Janikowski
+// Author Tomasz Mielech
 //
 
 package arangodb
@@ -28,6 +29,9 @@ type Collection interface {
 	Name() string
 	Database() Database
 
+	// Shards fetches shards information of the collection.
+	Shards(ctx context.Context, details bool) (CollectionShards, error)
+
 	// Remove removes the entire collection.
 	// If the collection does not exist, a NotFoundError is returned.
 	Remove(ctx context.Context) error
diff --git a/v2/arangodb/collection_impl.go b/v2/arangodb/collection_impl.go
@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2020 ArangoDB GmbH, Cologne, Germany
+// Copyright 2020-2021 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 // Author Adam Janikowski
+// Author Tomasz Mielech
 //
 
 package arangodb
@@ -26,8 +27,9 @@ import (
 	"context"
 	"net/http"
 
-	"github.com/arangodb/go-driver/v2/arangodb/shared"
+	"github.com/pkg/errors"
 
+	"github.com/arangodb/go-driver/v2/arangodb/shared"
 	"github.com/arangodb/go-driver/v2/connection"
 )
 
@@ -104,3 +106,24 @@ func (c collection) connection() connection.Connection {
 func (c collection) url(api string, parts ...string) string {
 	return c.db.url(append([]string{"_api", api, c.name}, parts...)...)
 }
+
+// Shards fetches shards information of the collection.
+func (c *collection) Shards(ctx context.Context, details bool) (CollectionShards, error) {
+	var body struct {
+		shared.ResponseStruct `json:",inline"`
+		CollectionShards      `json:",inline"`
+	}
+
+	resp, err := connection.CallGet(ctx, c.connection(), c.url("collection", "shards"), &body,
+		connection.WithQuery("details", "true"))
+	if err != nil {
+		return CollectionShards{}, errors.WithStack(err)
+	}
+
+	switch code := resp.Code(); code {
+	case http.StatusOK:
+		return body.CollectionShards, nil
+	default:
+		return CollectionShards{}, body.AsArangoErrorWithCode(code)
+	}
+}
diff --git a/v2/arangodb/collection_opts.go b/v2/arangodb/collection_opts.go
@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2017 ArangoDB GmbH, Cologne, Germany
+// Copyright 2017-2021 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 // Author Ewout Prangsma
+// Author Tomasz Mielech
 //
 
 package arangodb
@@ -26,7 +27,7 @@ import (
 	"time"
 )
 
-// CollectionInfo contains information about a collection
+// CollectionInfo contains basic information about a collection.
 type CollectionInfo struct {
 	// The identifier of the collection.
 	ID string `json:"id,omitempty"`
@@ -42,17 +43,10 @@ type CollectionInfo struct {
 	GloballyUniqueId string `json:"globallyUniqueId,omitempty"`
 }
 
-// CollectionProperties contains extended information about a collection.
-type CollectionProperties struct {
+// CollectionExtendedInfo contains extended information about a collection.
+type CollectionExtendedInfo struct {
 	CollectionInfo
-
-	// WaitForSync; If true then creating, changing or removing documents will wait until the data has been synchronized to disk.
-	WaitForSync bool `json:"waitForSync,omitempty"`
-	// DoCompact specifies whether or not the collection will be compacted.
-	DoCompact bool `json:"doCompact,omitempty"`
-	// JournalSize is the maximal size setting for journals / datafiles in bytes.
-	JournalSize int64 `json:"journalSize,omitempty"`
-	// CacheEnabled set cacheEnabled option in collection properties
+	// CacheEnabled set cacheEnabled option in collection properties.
 	CacheEnabled bool `json:"cacheEnabled,omitempty"`
 	KeyOptions   struct {
 		// Type specifies the type of the key generator. The currently available generators are traditional and autoincrement.
@@ -62,29 +56,41 @@ type CollectionProperties struct {
 		// the _key attribute of documents is considered an error.
 		AllowUserKeys bool `json:"allowUserKeys,omitempty"`
 	} `json:"keyOptions,omitempty"`
+	// Deprecated: use 'WriteConcern' instead.
+	MinReplicationFactor int `json:"minReplicationFactor,omitempty"`
 	// NumberOfShards is the number of shards of the collection.
 	// Only available in cluster setup.
 	NumberOfShards int `json:"numberOfShards,omitempty"`
+	// This attribute specifies the name of the sharding strategy to use for the collection.
+	// Can not be changed after creation.
+	ShardingStrategy ShardingStrategy `json:"shardingStrategy,omitempty"`
 	// ShardKeys contains the names of document attributes that are used to determine the target shard for documents.
 	// Only available in cluster setup.
 	ShardKeys []string `json:"shardKeys,omitempty"`
 	// ReplicationFactor contains how many copies of each shard are kept on different DBServers.
 	// Only available in cluster setup.
 	ReplicationFactor int `json:"replicationFactor,omitempty"`
-	// Deprecated: use 'WriteConcern' instead
-	MinReplicationFactor int `json:"minReplicationFactor,omitempty"`
+	// WaitForSync; If true then creating, changing or removing documents will wait
+	// until the data has been synchronized to disk.
+	WaitForSync bool `json:"waitForSync,omitempty"`
 	// WriteConcern contains how many copies must be available before a collection can be written.
 	// It is required that 1 <= WriteConcern <= ReplicationFactor.
 	// Default is 1. Not available for satellite collections.
 	// Available from 3.6 arangod version.
 	WriteConcern int `json:"writeConcern,omitempty"`
+}
+
+// CollectionProperties contains extended information about a collection.
+type CollectionProperties struct {
+	CollectionExtendedInfo
+	// DoCompact specifies whether or not the collection will be compacted.
+	DoCompact bool `json:"doCompact,omitempty"`
+	// JournalSize is the maximal size setting for journals / datafiles in bytes.
+	JournalSize int64 `json:"journalSize,omitempty"`
 	// SmartJoinAttribute
 	// See documentation for smart joins.
 	// This requires ArangoDB Enterprise Edition.
 	SmartJoinAttribute string `json:"smartJoinAttribute,omitempty"`
-	// This attribute specifies the name of the sharding strategy to use for the collection.
-	// Can not be changed after creation.
-	ShardingStrategy ShardingStrategy `json:"shardingStrategy,omitempty"`
 	// This attribute specifies that the sharding of a collection follows that of another
 	// one.
 	DistributeShardsLike string `json:"distributeShardsLike,omitempty"`
@@ -213,3 +219,18 @@ type CollectionStatistics struct {
 		} `json:"revisions"`
 	} `json:"figures"`
 }
+
+// CollectionShards contains shards information about a collection.
+type CollectionShards struct {
+	CollectionExtendedInfo
+	// Set to create a smart edge or vertex collection.
+	// This requires ArangoDB Enterprise Edition.
+	IsSmart bool `json:"isSmart,omitempty"`
+
+	// Shards is a list of shards that belong to the collection.
+	// Each shard contains a list of DB servers where the first one is the leader and the rest are followers.
+	Shards map[ShardID][]ServerID `json:"shards,omitempty"`
+
+	// StatusString represents status as a string.
+	StatusString string `json:"statusString,omitempty"`
+}
diff --git a/v2/arangodb/meta.go b/v2/arangodb/meta.go
@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2020 ArangoDB GmbH, Cologne, Germany
+// Copyright 2020-2021 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 // Author Adam Janikowski
+// Author Tomasz Mielech
 //
 
 package arangodb
@@ -27,6 +28,12 @@ import (
 	"github.com/pkg/errors"
 )
 
+// ShardID is an internal identifier of a specific shard.
+type ShardID string
+
+// ServerID identifies an ArangoDB server in a cluster.
+type ServerID string
+
 type DocumentID string
 
 // DocumentMeta contains all meta data used to identifier a document.
diff --git a/v2/tests/database_collection_operations_test.go b/v2/tests/database_collection_operations_test.go
@@ -1,7 +1,7 @@
 //
 // DISCLAIMER
 //
-// Copyright 2020 ArangoDB GmbH, Cologne, Germany
+// Copyright 2020-2021 ArangoDB GmbH, Cologne, Germany
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 // Author Adam Janikowski
+// Author Tomasz Mielech
 //
 
 package tests
@@ -28,12 +29,42 @@ import (
 	"testing"
 	"time"
 
-	"github.com/arangodb/go-driver/v2/arangodb"
-	"github.com/arangodb/go-driver/v2/arangodb/shared"
 	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/arangodb/go-driver/v2/arangodb"
+	"github.com/arangodb/go-driver/v2/arangodb/shared"
 )
 
+// Test_CollectionShards creates a collection and gets the shards' information.
+func Test_CollectionShards(t *testing.T) {
+	requireClusterMode(t)
+
+	options := arangodb.CreateCollectionOptions{
+		ReplicationFactor: 2,
+		NumberOfShards:    2,
+	}
+
+	Wrap(t, func(t *testing.T, client arangodb.Client) {
+		WithDatabase(t, client, nil, func(db arangodb.Database) {
+			WithCollection(t, db, &options, func(col arangodb.Collection) {
+				shards, err := col.Shards(context.Background(), true)
+				require.NoError(t, err)
+
+				require.Len(t, shards.Shards, 2, "expected 2 shards")
+				var leaders []arangodb.ServerID
+
+				for _, dbServers := range shards.Shards {
+					require.Lenf(t, dbServers, 2, "expected 2 DB servers for the shard")
+					leaders = append(leaders, dbServers[0])
+				}
+				assert.NotEqualf(t, leaders[0], leaders[1], "the leader shard can not be on the same server")
+			})
+		})
+	})
+}
+
 func Test_DatabaseCollectionOperations(t *testing.T) {
 
 	Wrap(t, func(t *testing.T, client arangodb.Client) {
diff --git a/v2/tests/util_test.go b/v2/tests/util_test.go
@@ -0,0 +1,59 @@
+//
+// DISCLAIMER
+//
+// Copyright 2021 ArangoDB GmbH, Cologne, Germany
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Copyright holder is ArangoDB GmbH, Cologne, Germany
+//
+// Author Tomasz Mielech
+//
+
+package tests
+
+import (
+	"os"
+	"strings"
+	"testing"
+)
+
+func getTestMode() string {
+	return strings.TrimSpace(os.Getenv("TEST_MODE"))
+}
+
+type mode string
+
+const (
+	testModeCluster         mode = "cluster"
+	testModeResilientSingle mode = "resilientsingle"
+	testModeSingle          mode = "single"
+)
+
+func requireMode(t *testing.T, mode mode) {
+	if getTestMode() != string(mode) {
+		t.Skipf("the test requires %s mode", mode)
+	}
+}
+
+func requireClusterMode(t *testing.T) {
+	requireMode(t, testModeCluster)
+}
+
+func requireSingleMode(t *testing.T) {
+	requireMode(t, testModeSingle)
+}
+
+func requireResilientSingleMode(t *testing.T) {
+	requireMode(t, testModeResilientSingle)
+}

Original file line number	Diff line number	Diff line change
`@@ -966,8 +966,8 @@ func TestCollectionWriteConcernSetPropInvalid(t *testing.T) {`
`966`	`966`	`minRepl, prop.WriteConcern)`
`967`	`967`	`}`
`968`	`968`
`969`		`-// TestCollectionShards creates a collection and gets the shards' information.`
`970`		`-func TestCollectionShards(t *testing.T) {`
	`969`	`+// Test_CollectionShards creates a collection and gets the shards' information.`
	`970`	`+func Test_CollectionShards(t *testing.T) {`
`971`	`971`	`if getTestMode() != testModeCluster {`
`972`	`972`	`t.Skipf("Not a cluster mode")`
`973`	`973`	`}`
`@@ -990,5 +990,5 @@ func TestCollectionShards(t *testing.T) {`
`990`	`990`	`require.Lenf(t, dbServers, 2, "expected 2 DB servers for the shard")`
`991`	`991`	`leaders = append(leaders, dbServers[0])`
`992`	`992`	`}`
`993`		`- assert.NotEqualf(t, leaders[0], leaders[1], "the leader can not be on the same server")`
	`993`	`+ assert.NotEqualf(t, leaders[0], leaders[1], "the leader shard can not be on the same server")`
`994`	`994`	`}`