From 8b0035b0f07a43cb2fc5bb4dd8fd9b79d118c723 Mon Sep 17 00:00:00 2001
From: Nicolas <zepouet@gmail.com>
Date: Thu, 9 Jul 2015 17:46:01 +0200
Subject: [PATCH] Add filtering options to select plugin at startup

---
 agent.go                 | 42 ++++++++++++++++++++++++++++++++--------
 agent_test.go            | 42 ++++++++++++++++++++++++++++++++++++++++
 cmd/telegraf/telegraf.go |  3 ++-
 3 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/agent.go b/agent.go
index 7ac924f63..ee2f753cc 100644
--- a/agent.go
+++ b/agent.go
@@ -8,6 +8,7 @@ import (
 	"sort"
 	"sync"
 	"time"
+	"strings"
 
 	"github.com/influxdb/influxdb/client"
 	"github.com/influxdb/telegraf/plugins"
@@ -81,24 +82,49 @@ func (agent *Agent) Connect() error {
 	return nil
 }
 
-func (a *Agent) LoadPlugins() ([]string, error) {
+func (a *Agent) LoadPlugins(pluginsFilter string) ([]string, error) {
 	var names []string
+	var filters []string
+
+	pluginsFilter = strings.TrimSpace(pluginsFilter)
+	if pluginsFilter != "" {
+		filters = strings.Split(":"+pluginsFilter+":", ":")
+	}
 
 	for _, name := range a.Config.PluginsDeclared() {
+
 		creator, ok := plugins.Plugins[name]
+
 		if !ok {
 			return nil, fmt.Errorf("Undefined but requested plugin: %s", name)
 		}
 
-		plugin := creator()
-
-		config, err := a.Config.ApplyPlugin(name, plugin)
-		if err != nil {
-			return nil, err
+		// to know if plugin is enabled or not through filter flag
+		isPluginEnabled := false
+		if len(filters)>0 {
+			for _, runeValue := range filters {
+				if runeValue != "" && strings.ToLower(runeValue) == strings.ToLower(name) {
+					fmt.Printf("plugin [%s] is enabled because present is filter options\n", name)
+					isPluginEnabled = true
+					break
+				}
+			}
+		} else {
+			// if no filter, we ALWAYS accept the plugin
+			isPluginEnabled = true
+		}
+
+		if isPluginEnabled {
+			plugin := creator()
+			config, err := a.Config.ApplyPlugin(name, plugin)
+			if err != nil {
+				return nil, err
+			}
+
+			a.plugins = append(a.plugins, &runningPlugin{name, plugin, config})
+			names = append(names, name)
 		}
 
-		a.plugins = append(a.plugins, &runningPlugin{name, plugin, config})
-		names = append(names, name)
 	}
 
 	sort.Strings(names)
diff --git a/agent_test.go b/agent_test.go
index 8811fa810..4d9f7a8ee 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -1,5 +1,47 @@
 package telegraf
 
+import (
+	"testing"
+	"github.com/stretchr/testify/assert"
+
+	// needing to load the plugins
+	_ "github.com/influxdb/telegraf/plugins/all"
+)
+
+func TestAgent_LoadPlugin(t *testing.T) {
+
+	// load a dedicated configuration file
+	config, _ := LoadConfig("./testdata/telegraf-agent.toml")
+	a, _ := NewAgent(config)
+
+	pluginsEnabled, _ := a.LoadPlugins("mysql")
+	assert.Equal(t, 1, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("foo")
+	assert.Equal(t, 0, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("mysql:foo")
+	assert.Equal(t, 1, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("mysql:redis")
+	assert.Equal(t, 2, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins(":mysql:foo:redis:bar")
+	assert.Equal(t, 2, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("")
+	assert.Equal(t, 14, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins(" ")
+	assert.Equal(t, 14, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("		")
+	assert.Equal(t, 14, len(pluginsEnabled))
+
+	pluginsEnabled, _ = a.LoadPlugins("\n\t")
+	assert.Equal(t, 14, len(pluginsEnabled))
+}
+
 /*
 func TestAgent_DrivesMetrics(t *testing.T) {
 	var (
diff --git a/cmd/telegraf/telegraf.go b/cmd/telegraf/telegraf.go
index f36693ebc..53fbfd44e 100644
--- a/cmd/telegraf/telegraf.go
+++ b/cmd/telegraf/telegraf.go
@@ -18,6 +18,7 @@ var fConfig = flag.String("config", "", "configuration file to load")
 var fVersion = flag.Bool("version", false, "display the version")
 var fSampleConfig = flag.Bool("sample-config", false, "print out full sample configuration")
 var fPidfile = flag.String("pidfile", "", "file to write our pid to")
+var fPLuginsFilter = flag.String("filter", "", "filter the plugin to enable")
 
 var Version = "unreleased"
 var Commit = ""
@@ -58,7 +59,7 @@ func main() {
 		ag.Debug = true
 	}
 
-	plugins, err := ag.LoadPlugins()
+	plugins, err := ag.LoadPlugins(*fPLuginsFilter)
 	if err != nil {
 		log.Fatal(err)
 	}