From 6a81a2afbf3aefc62a96d5334a1d13062f272238 Mon Sep 17 00:00:00 2001
From: Justin Chan <juchan@uber.com>
Date: Tue, 20 Nov 2018 13:50:05 -0500
Subject: [PATCH] Respect numShards in create

---
 src/query/api/v1/handler/database/create.go   |   5 +
 .../api/v1/handler/database/create_test.go    | 100 ++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/src/query/api/v1/handler/database/create.go b/src/query/api/v1/handler/database/create.go
index 99ec580f44..f41ed4bb05 100644
--- a/src/query/api/v1/handler/database/create.go
+++ b/src/query/api/v1/handler/database/create.go
@@ -381,6 +381,11 @@ func defaultedPlacementInitRequest(
 		return nil, errInvalidDBType
 	}
 
+	numShardsInput := r.GetNumShards()
+	if numShardsInput > 0 {
+		numShards = numShardsInput
+	}
+
 	return &admin.PlacementInitRequest{
 		NumShards:         numShards,
 		ReplicationFactor: replicationFactor,
diff --git a/src/query/api/v1/handler/database/create_test.go b/src/query/api/v1/handler/database/create_test.go
index 70359b9978..fa3fcd5ccc 100644
--- a/src/query/api/v1/handler/database/create_test.go
+++ b/src/query/api/v1/handler/database/create_test.go
@@ -175,6 +175,106 @@ func TestLocalType(t *testing.T) {
 		xtest.Diff(mustPrettyJSON(t, expectedResponse), mustPrettyJSON(t, string(body))))
 }
 
+func TestLocalTypeWithNumShards(t *testing.T) {
+	mockClient, mockKV, mockPlacementService := SetupDatabaseTest(t)
+	createHandler := NewCreateHandler(mockClient, config.Configuration{}, testDBCfg)
+	w := httptest.NewRecorder()
+
+	jsonInput := `
+		{
+			"namespaceName": "testNamespace",
+			"type": "local",
+			"numShards": 51
+		}
+	`
+
+	req := httptest.NewRequest("POST", "/database/create", strings.NewReader(jsonInput))
+	require.NotNil(t, req)
+
+	mockKV.EXPECT().Get(namespace.M3DBNodeNamespacesKey).Return(nil, kv.ErrNotFound)
+	mockKV.EXPECT().CheckAndSet(namespace.M3DBNodeNamespacesKey, gomock.Any(), gomock.Not(nil)).Return(1, nil)
+
+	placementProto := &placementpb.Placement{
+		Instances: map[string]*placementpb.Instance{
+			"localhost": &placementpb.Instance{
+				Id:             "m3db_local",
+				IsolationGroup: "local",
+				Zone:           "embedded",
+				Weight:         1,
+				Endpoint:       "http://localhost:9000",
+				Hostname:       "localhost",
+				Port:           9000,
+			},
+		},
+	}
+	newPlacement, err := placement.NewPlacementFromProto(placementProto)
+	require.NoError(t, err)
+	mockPlacementService.EXPECT().BuildInitialPlacement(gomock.Any(), 51, 1).Return(newPlacement, nil)
+
+	createHandler.ServeHTTP(w, req)
+
+	resp := w.Result()
+	body, err := ioutil.ReadAll(resp.Body)
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+
+	expectedResponse := `
+	{
+		"namespace": {
+			"registry": {
+				"namespaces": {
+					"testNamespace": {
+						"bootstrapEnabled": true,
+						"flushEnabled": true,
+						"writesToCommitLog": true,
+						"cleanupEnabled": true,
+						"repairEnabled": false,
+						"retentionOptions": {
+							"retentionPeriodNanos": "86400000000000",
+							"blockSizeNanos": "3600000000000",
+							"bufferFutureNanos": "120000000000",
+							"bufferPastNanos": "600000000000",
+							"blockDataExpiry": true,
+							"blockDataExpiryAfterNotAccessPeriodNanos": "300000000000"
+						},
+						"snapshotEnabled": false,
+						"indexOptions": {
+							"enabled": true,
+							"blockSizeNanos": "3600000000000"
+						}
+					}
+				}
+			}
+		},
+		"placement": {
+			"placement": {
+				"instances": {
+					"m3db_local": {
+						"id": "m3db_local",
+						"isolationGroup": "local",
+						"zone": "embedded",
+						"weight": 1,
+						"endpoint": "http://localhost:9000",
+						"shards": [],
+						"shardSetId": 0,
+						"hostname": "localhost",
+						"port": 9000
+					}
+				},
+				"replicaFactor": 0,
+				"numShards": 0,
+				"isSharded": false,
+				"cutoverTime": "0",
+				"isMirrored": false,
+				"maxShardSetId": 0
+			},
+			"version": 0
+		}
+	}
+	`
+	assert.Equal(t, stripAllWhitespace(expectedResponse), string(body),
+		xtest.Diff(mustPrettyJSON(t, expectedResponse), mustPrettyJSON(t, string(body))))
+}
 func TestLocalWithBlockSizeNanos(t *testing.T) {
 	mockClient, mockKV, mockPlacementService := SetupDatabaseTest(t)
 	createHandler := NewCreateHandler(mockClient, config.Configuration{}, testDBCfg)