Feature/get shards info (#331)

* Get shards information

* Add v2 test for getting collection shards

* Create more reliable collection shards test

Author: informalict

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Change Log
 
 ## [master](https://github.com/arangodb/go-driver/tree/master) (N/A)
+- Add support for fetching shards' info by the given collection name.
 
 ## [1.2.0](https://github.com/arangodb/go-driver/tree/1.2.0) (2021-08-04)
 - Add support for AQL, Pipeline, Stopwords, GeoJSON and GeoPoint Arango Search analyzers.

collection.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 driver
@@ -55,10 +56,13 @@ type Collection interface {
 	// SetProperties changes properties of the collection.
 	SetProperties(ctx context.Context, options SetCollectionPropertiesOptions) error
 
+	// Shards fetches shards information of the collection.
+	Shards(ctx context.Context, details bool) (CollectionShards, error)
+
 	// Load the collection into memory.
 	Load(ctx context.Context) error
 
-	// UnLoad the collection from memory.
+	// Unload unloads the collection from memory.
 	Unload(ctx context.Context) error
 
 	// Remove removes the entire collection.
@@ -262,3 +266,60 @@ type CollectionStatistics struct {
 		} `json:"revisions"`
 	} `json:"figures"`
 }
+
+// CollectionShards contains shards information about a collection.
+type CollectionShards struct {
+	CollectionInfo
+
+	// CacheEnabled set cacheEnabled option in collection properties
+	CacheEnabled bool `json:"cacheEnabled,omitempty"`
+
+	// Set to create a smart edge or vertex collection.
+	// This requires ArangoDB Enterprise Edition.
+	IsSmart bool `json:"isSmart,omitempty"`
+
+	KeyOptions struct {
+		// Type specifies the type of the key generator. The currently available generators are traditional and autoincrement.
+		Type KeyGeneratorType `json:"type,omitempty"`
+		// AllowUserKeys; if set to true, then it is allowed to supply own key values in the _key attribute of a document.
+		// If set to false, then the key generator is solely responsible for generating keys and supplying own key values in
+		// 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"`
+
+	// 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"`
+
+	// ReplicationFactor contains how many copies of each shard are kept on different DBServers.
+	// Only available in cluster setup.
+	ReplicationFactor int `json:"replicationFactor,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 ArangoDB version.
+	WriteConcern int `json:"writeConcern,omitempty"`
+}

collection_impl.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 driver
@@ -190,6 +191,34 @@ func (c *collection) SetProperties(ctx context.Context, options SetCollectionPro
 	return nil
 }
 
+// Shards fetches shards information of the collection.
+func (c *collection) Shards(ctx context.Context, details bool) (CollectionShards, error) {
+	req, err := c.conn.NewRequest("GET", path.Join(c.relPath("collection"), "shards"))
+	if err != nil {
+		return CollectionShards{}, WithStack(err)
+	}
+	if details {
+		req.SetQuery("details", "true")
+	}
+
+	resp, err := c.conn.Do(ctx, req)
+	if err != nil {
+		return CollectionShards{}, WithStack(err)
+	}
+
+	if err := resp.CheckStatus(200); err != nil {
+		return CollectionShards{}, WithStack(err)
+	}
+
+	shards := CollectionShards{}
+	if err := resp.ParseBody("", &shards); err != nil {
+		return CollectionShards{}, WithStack(err)
+	}
+
+	return shards, nil
+
+}
+
 // Load the collection into memory.
 func (c *collection) Load(ctx context.Context) error {
 	req, err := c.conn.NewRequest("PUT", path.Join(c.relPath("collection"), "load"))

edge_collection_impl.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 driver
@@ -126,6 +127,15 @@ func (c *edgeCollection) SetProperties(ctx context.Context, options SetCollectio
 	return nil
 }
 
+// Shards fetches shards information of the collection.
+func (c *edgeCollection) Shards(ctx context.Context, details bool) (CollectionShards, error) {
+	result, err := c.rawCollection().Shards(ctx, details)
+	if err != nil {
+		return result, WithStack(err)
+	}
+	return result, nil
+}
+
 // Load the collection into memory.
 func (c *edgeCollection) Load(ctx context.Context) error {
 	if err := c.rawCollection().Load(ctx); err != nil {

test/collection_test.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 test
@@ -964,3 +965,46 @@ func TestCollectionWriteConcernSetPropInvalid(t *testing.T) {
 	assert.Equalf(t, defaultWriteConcern, prop.WriteConcern, "MinReplicationFactor not updated, expected %d, found %d",
 		minRepl, prop.WriteConcern)
 }
+
+// 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")
+	}
+
+	databaseName := getCallerFunctionName()
+	c := createClientFromEnv(t, true)
+	db := ensureDatabase(nil, c, databaseName, nil, t)
+	name := "test_collection_set_properties"
+	col, err := db.CreateCollection(nil, name, &driver.CreateCollectionOptions{
+		ReplicationFactor: 2,
+		NumberOfShards:    2,
+	})
+	require.NoError(t, err)
+
+	shards, err := col.Shards(context.Background(), true)
+	require.NoError(t, err)
+	assert.NotEmpty(t, shards.ID)
+	assert.Equal(t, name, shards.Name)
+	assert.NotEmpty(t, shards.Status)
+	assert.Equal(t, driver.CollectionTypeDocument, shards.Type)
+	assert.Equal(t, false, shards.IsSystem)
+	assert.NotEmpty(t, shards.GloballyUniqueId)
+	assert.Equal(t, false, shards.CacheEnabled)
+	assert.Equal(t, false, shards.IsSmart)
+	assert.Equal(t, driver.KeyGeneratorTraditional, shards.KeyOptions.Type)
+	assert.Equal(t, true, shards.KeyOptions.AllowUserKeys)
+	assert.Equal(t, 2, shards.NumberOfShards)
+	assert.Equal(t, driver.ShardingStrategyHash, shards.ShardingStrategy)
+	assert.Equal(t, []string{"_key"}, shards.ShardKeys)
+	require.Len(t, shards.Shards, 2, "expected 2 shards")
+	var leaders []driver.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")
+	assert.Equal(t, 2, shards.ReplicationFactor)
+	assert.Equal(t, false, shards.WaitForSync)
+	assert.Equal(t, 1, shards.WriteConcern)
+}

test/util.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 test
@@ -28,13 +29,14 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
 	"time"
 
+	"github.com/dchest/uniuri"
 	"github.com/google/uuid"
-
 	"github.com/stretchr/testify/require"
 
 	driver "github.com/arangodb/go-driver"
@@ -212,3 +214,18 @@ func min(max int, ints ...int) int {
 
 	return z
 }
+
+// getCallerFunctionName returns the name of the function of the caller.
+func getCallerFunctionName() string {
+	programCounters := make([]uintptr, 10)
+	// skip this function and 'runtime.Callers' function
+	runtime.Callers(2, programCounters)
+	functionPackage := runtime.FuncForPC(programCounters[0])
+
+	function := strings.Split(functionPackage.Name(), ".")
+	if len(function) > 1 {
+		return function[len(function)-1] + "_" + uniuri.NewLen(6)
+	}
+
+	return function[0] + "_" + uniuri.NewLen(6)
+}

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

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