From 4a02ad7525768ce3f0fc8a2822b98c8a646e425d Mon Sep 17 00:00:00 2001
From: Tim Raymond <tim@timraymond.com>
Date: Fri, 4 Nov 2016 16:52:15 -0400
Subject: [PATCH] Enforce single default source on update

This ensures that there will only be on default source upon setting a
new default by unsetting the default flag on all other sources. This
only happens when the source to be updated has Default set to true to
avoid a performance hit when updating other attributes.
---
 bolt/sources.go      | 19 +++++++++++++++++++
 bolt/sources_test.go | 26 +++++++++++++++++++++-----
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/bolt/sources.go b/bolt/sources.go
index 5efc990163..13711a6ec4 100644
--- a/bolt/sources.go
+++ b/bolt/sources.go
@@ -103,6 +103,25 @@ func (s *SourcesStore) Update(ctx context.Context, src chronograf.Source) error
 			return chronograf.ErrSourceNotFound
 		}
 
+		//Unset any existing defaults if this is a new default
+		if src.Default {
+			srcs, err := s.All(ctx)
+			if err != nil {
+				return err
+			}
+
+			for _, other := range srcs {
+				if other.Default {
+					other.Default = false
+					if v, err := internal.MarshalSource(other); err != nil {
+						return err
+					} else if err := b.Put(itob(other.ID), v); err != nil {
+						return err
+					}
+				}
+			}
+		}
+
 		if v, err := internal.MarshalSource(src); err != nil {
 			return err
 		} else if err := b.Put(itob(src.ID), v); err != nil {
diff --git a/bolt/sources_test.go b/bolt/sources_test.go
index 60dbc9346d..30b7620421 100644
--- a/bolt/sources_test.go
+++ b/bolt/sources_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/influxdata/chronograf"
+	"github.com/influxdata/chronograf/bolt"
 )
 
 // Ensure an SourceStore can store, retrieve, update, and delete sources.
@@ -54,11 +55,8 @@ func TestSourceStore(t *testing.T) {
 	// Update source.
 	srcs[0].Username = "calvinklein"
 	srcs[1].Name = "Enchantment Under the Sea Dance"
-	if err := s.Update(nil, srcs[0]); err != nil {
-		t.Fatal(err)
-	} else if err := s.Update(nil, srcs[1]); err != nil {
-		t.Fatal(err)
-	}
+	mustUpdateSource(t, s, srcs[0])
+	mustUpdateSource(t, s, srcs[1])
 
 	// Confirm sources have updated.
 	if src, err := s.Get(nil, srcs[0].ID); err != nil {
@@ -72,6 +70,18 @@ func TestSourceStore(t *testing.T) {
 		t.Fatalf("source 1 update error: got %v, expected %v", src.Name, "Enchantment Under the Sea Dance")
 	}
 
+	// Attempt to make two default sources
+	srcs[0].Default = true
+	srcs[1].Default = true
+	mustUpdateSource(t, s, srcs[0])
+	mustUpdateSource(t, s, srcs[1])
+
+	if actual, err := s.Get(nil, srcs[0].ID); err != nil {
+		t.Fatal(err)
+	} else if actual.Default == true {
+		t.Fatal("Able to set two default sources when only one should be permitted")
+	}
+
 	// Delete an source.
 	if err := s.Delete(nil, srcs[0]); err != nil {
 		t.Fatal(err)
@@ -90,3 +100,9 @@ func TestSourceStore(t *testing.T) {
 		t.Fatalf("After delete All returned incorrect source; got %v, expected %v", bsrcs[0], srcs[1])
 	}
 }
+
+func mustUpdateSource(t *testing.T, s *bolt.SourcesStore, src chronograf.Source) {
+	if err := s.Update(nil, src); err != nil {
+		t.Fatal(err)
+	}
+}