storage constructor

cadvisor.go

@@ -19,8 +19,6 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	"os"
-	"time"
 
 	"github.com/google/cadvisor/api"
 	"github.com/google/cadvisor/container/docker"
@@ -29,9 +27,6 @@ import (
 	"github.com/google/cadvisor/manager"
 	"github.com/google/cadvisor/pages"
 	"github.com/google/cadvisor/pages/static"
-	"github.com/google/cadvisor/storage"
-	"github.com/google/cadvisor/storage/influxdb"
-	"github.com/google/cadvisor/storage/memory"
 )
 
 var argPort = flag.Int("port", 8080, "port to listen")
@@ -39,45 +34,12 @@ var argSampleSize = flag.Int("samples", 1024, "number of samples we want to keep
 var argHistoryDuration = flag.Int("history_duration", 60, "number of seconds of container history to keep")
 var argAllowLmctfy = flag.Bool("allow_lmctfy", true, "whether to allow lmctfy as a container handler")
 
-// database related arguments
 var argDbDriver = flag.String("db", "memory", "database driver name")
-var argDbUsername = flag.String("db.user", "root", "database username")
-var argDbPassword = flag.String("db.password", "root", "database password")
-var argDbHost = flag.String("db.host", "localhost:8086", "database host:port")
-var argDbName = flag.String("db.name", "cadvisor", "database name")
-var argDbIsSecure = flag.Bool("db.secure", false, "use secure connection with database")
 
 func main() {
 	flag.Parse()
 
-	var storageDriver storage.StorageDriver
+	storageDriver, err := NewStorage(*argDbDriver)
-	var err error
-
-	switch *argDbDriver {
-	case "":
-		// empty string by default is the in memory store
-		fallthrough
-	case "memory":
-		storageDriver = memory.New(*argSampleSize, *argHistoryDuration)
-	case "influxdb":
-		var hostname string
-		hostname, err = os.Hostname()
-		if err == nil {
-			storageDriver, err = influxdb.New(
-				hostname,
-				"cadvisorTable",
-				*argDbName,
-				*argDbUsername,
-				*argDbPassword,
-				*argDbHost,
-				*argDbIsSecure,
-				// TODO(monnand): One hour? Or user-defined?
-				1*time.Hour,
-			)
-		}
-	default:
-		err = fmt.Errorf("Unknown database driver: %v", *argDbDriver)
-	}
 	if err != nil {
 		log.Fatalf("Failed to connect to database: %s", err)
 	}

storagedriver.go

@@ -0,0 +1,69 @@
+// Copyright 2014 Google Inc. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/google/cadvisor/storage"
+	"github.com/google/cadvisor/storage/influxdb"
+	"github.com/google/cadvisor/storage/memory"
+)
+
+var argDbUsername = flag.String("db.user", "root", "database username")
+var argDbPassword = flag.String("db.password", "root", "database password")
+var argDbHost = flag.String("db.host", "localhost:8086", "database host:port")
+var argDbName = flag.String("db.name", "cadvisor", "database name")
+var argDbIsSecure = flag.Bool("db.secure", false, "use secure connection with database")
+
+func NewStorage(driverName string) (storage.StorageDriver, error) {
+	var storageDriver storage.StorageDriver
+	var err error
+	switch driverName {
+	case "":
+		// empty string by default is the in memory store
+		fallthrough
+	case "memory":
+		storageDriver = memory.New(*argSampleSize, *argHistoryDuration)
+		return storageDriver, nil
+	case "influxdb":
+		var hostname string
+		hostname, err = os.Hostname()
+		if err != nil {
+			return nil, err
+		}
+
+		storageDriver, err = influxdb.New(
+			hostname,
+			"cadvisorTable",
+			*argDbName,
+			*argDbUsername,
+			*argDbPassword,
+			*argDbHost,
+			*argDbIsSecure,
+			// TODO(monnand): One hour? Or user-defined?
+			1*time.Hour,
+		)
+	default:
+		err = fmt.Errorf("Unknown database driver: %v", *argDbDriver)
+	}
+	if err != nil {
+		return nil, err
+	}
+	return storageDriver, nil
+}