From 363b6157c444a3337a4b44c0e3de2fe589513caf Mon Sep 17 00:00:00 2001
From: Zachary Wasserman <zach@dactiv.llc>
Date: Tue, 13 Aug 2019 09:42:02 -0700
Subject: [PATCH] Allow scheduled query shard to be set to empty in frontend
 (#2093)

Fixes #1452
---
 Gopkg.lock                                     |  9 +++++++++
 Gopkg.toml                                     |  4 ++++
 server/kolide/scheduled_queries.go             | 18 ++++++++++--------
 server/service/service_scheduled_queries.go    |  7 ++++++-
 .../transport_scheduled_queries_test.go        |  5 +++--
 5 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/Gopkg.lock b/Gopkg.lock
index 3a43288f41..f3168c51cc 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -270,6 +270,14 @@
   revision = "ea4d1f681babbce9545c9c5f3d5194a789c89f5b"
   version = "v1.2.0"
 
+[[projects]]
+  digest = "1:a4b04b5314d49a67ad022fa7693453ee2b823ab99314f52fa3930ed31085dfef"
+  name = "github.com/guregu/null"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "80515d440932108546bcade467bb7d6968e812e2"
+  version = "v3.4.0"
+
 [[projects]]
   digest = "1:67474f760e9ac3799f740db2c489e6423a4cde45520673ec123ac831ad849cb8"
   name = "github.com/hashicorp/golang-lru"
@@ -858,6 +866,7 @@
     "github.com/gomodule/redigo/redis",
     "github.com/gorilla/mux",
     "github.com/gorilla/websocket",
+    "github.com/guregu/null",
     "github.com/igm/sockjs-go/sockjs",
     "github.com/jmoiron/sqlx",
     "github.com/jmoiron/sqlx/reflectx",
diff --git a/Gopkg.toml b/Gopkg.toml
index 0cb716dc55..907f9113c4 100644
--- a/Gopkg.toml
+++ b/Gopkg.toml
@@ -125,3 +125,7 @@
 [[constraint]]
   name = "cloud.google.com/go"
   version = "0.37.4"
+
+[[constraint]]
+  name = "github.com/guregu/null"
+  version = "3.4.0"
diff --git a/server/kolide/scheduled_queries.go b/server/kolide/scheduled_queries.go
index 1f91d27839..d49d4f2946 100644
--- a/server/kolide/scheduled_queries.go
+++ b/server/kolide/scheduled_queries.go
@@ -2,6 +2,8 @@ package kolide
 
 import (
 	"context"
+
+	"github.com/guregu/null"
 )
 
 type ScheduledQueryStore interface {
@@ -39,12 +41,12 @@ type ScheduledQuery struct {
 }
 
 type ScheduledQueryPayload struct {
-	PackID   *uint   `json:"pack_id"`
-	QueryID  *uint   `json:"query_id"`
-	Interval *uint   `json:"interval"`
-	Snapshot *bool   `json:"snapshot"`
-	Removed  *bool   `json:"removed"`
-	Platform *string `json:"platform"`
-	Version  *string `json:"version"`
-	Shard    *uint   `json:"shard"`
+	PackID   *uint     `json:"pack_id"`
+	QueryID  *uint     `json:"query_id"`
+	Interval *uint     `json:"interval"`
+	Snapshot *bool     `json:"snapshot"`
+	Removed  *bool     `json:"removed"`
+	Platform *string   `json:"platform"`
+	Version  *string   `json:"version"`
+	Shard    *null.Int `json:"shard"`
 }
diff --git a/server/service/service_scheduled_queries.go b/server/service/service_scheduled_queries.go
index 2272c29783..0a6ea37d97 100644
--- a/server/service/service_scheduled_queries.go
+++ b/server/service/service_scheduled_queries.go
@@ -64,7 +64,12 @@ func (svc service) ModifyScheduledQuery(ctx context.Context, id uint, p kolide.S
 	}
 
 	if p.Shard != nil {
-		sq.Shard = p.Shard
+		if p.Shard.Valid {
+			val := uint(p.Shard.Int64)
+			sq.Shard = &val
+		} else {
+			sq.Shard = nil
+		}
 	}
 
 	return svc.ds.SaveScheduledQuery(sq)
diff --git a/server/service/transport_scheduled_queries_test.go b/server/service/transport_scheduled_queries_test.go
index 0ad649e341..c3b87d4a6e 100644
--- a/server/service/transport_scheduled_queries_test.go
+++ b/server/service/transport_scheduled_queries_test.go
@@ -67,12 +67,13 @@ func TestDecodeModifyScheduledQueryRequest(t *testing.T) {
 		assert.Equal(t, uint(6), *params.payload.QueryID)
 		assert.Equal(t, true, *params.payload.Removed)
 		assert.Equal(t, uint(60), *params.payload.Interval)
-		assert.Equal(t, uint(1), *params.payload.Shard)
+		assert.Equal(t, true, params.payload.Shard.Valid)
+		assert.Equal(t, int64(1), params.payload.Shard.Int64)
 	}).Methods("PATCH")
 
 	var body bytes.Buffer
 	body.Write([]byte(`{
-        "pack_id": 5,
+	        "pack_id": 5,
 		"query_id": 6,
 		"removed": true,
 		"interval": 60,