Merge pull request #182 from GrappigPanda/issue-179

Fixes #179

config.yaml

@@ -7,7 +7,7 @@ UseMySQL: false
 DBHost: localhost
 DBPort: "3306"
 DBUser: testuser
-DBPass: testuser 
+DBPass: testuser
 DBName: testdb
 # Redis Options
 RedisServer: localhost

database/mysql/database_test.go

@@ -18,14 +18,13 @@ var CONFIG = config.ConfigStruct{
 	false,
 }
 
-var DBCONN, _ = OpenConnectionWithConfig(&CONFIG)
+var DBCONN, ERR = OpenConnectionWithConfig(&CONFIG)
 
-func TestOpenConn(t *testing.T) {
-	dbConn, err := OpenConnection()
-	if err != nil {
-		t.Fatalf("%v", err)
+func TestOpenConnPostgres(t *testing.T) {
+	if ERR != nil {
+		t.Fatalf("Unable to connect %v", ERR)
 	}
-	InitDB(dbConn)
+	InitDB(DBCONN)
 }
 
 func TestAddWhitelistedTorrent(t *testing.T) {

database/mysql/mysql.go

@@ -14,24 +14,18 @@ import (
 // MysqlHost listed in the config
 func OpenConnection() (db *gorm.DB, err error) {
 	c := config.LoadConfig()
-
-	db, err = gorm.Open("mysql", formatConnectString(c))
-	if err != nil {
-		err = fmt.Errorf("Failed to open connection to MySQL: %v", err)
-	}
-
-	return
+	return OpenConnectionWithConfig(&c)
 }
 
 // OpenConnectionWithConfig does as its name dictates and opens a connection to the
 // MysqlHost listed in the config
 func OpenConnectionWithConfig(cfg *config.ConfigStruct) (db *gorm.DB, err error) {
 	db, err = gorm.Open("mysql", formatConnectString(*cfg))
 	if err != nil {
-		err = fmt.Errorf("Failed to open connection to MySQL: %v", err)
+		return nil, fmt.Errorf("Failed to open connection to MySQL: %v", err)
 	}
 
-	return
+	return db, nil
 }
 
 // InitDB initializes database tables.

reaper/reaper.go

@@ -2,7 +2,9 @@ package reaper
 
 import (
 	"fmt"
+	"github.com/GrappigPanda/notorious/config"
 	"github.com/GrappigPanda/notorious/database/mysql"
+	"github.com/GrappigPanda/notorious/database/postgres"
 	"github.com/GrappigPanda/notorious/peerStore/redis"
 	"gopkg.in/redis.v3"
 	"strconv"
@@ -83,10 +85,26 @@ func StartReapingScheduler(waitTime time.Duration) {
 			addedBy := new(string)
 			dateAdded := new(int64)
 
-			x, err := mysql.GetWhitelistedTorrents(nil)
-			for x.Next() {
-				x.Scan(infoHash, name, addedBy, dateAdded)
-				redisPeerStore.CreateNewTorrentKey(nil, *infoHash)
+			cfg := config.LoadConfig()
+
+			if cfg.DBChoice == "mysql" {
+				whitelistedTorrents, err := mysql.GetWhitelistedTorrents(nil)
+				if err == nil {
+					for whitelistedTorrents.Next() {
+						whitelistedTorrents.Scan(infoHash, name, addedBy, dateAdded)
+						redisPeerStore.CreateNewTorrentKey(nil, *infoHash)
+					}
+				}
+			} else if cfg.DBChoice == "postgres" {
+				whitelistedTorrents, err := postgres.GetWhitelistedTorrents(nil)
+				if err == nil {
+					for whitelistedTorrents.Next() {
+						whitelistedTorrents.Scan(infoHash, name, addedBy, dateAdded)
+						redisPeerStore.CreateNewTorrentKey(nil, *infoHash)
+					}
+				}
+			} else {
+				panic("Invalid DBChoice encountered when starting the peer reaper.")
 			}
 
 			// Start the actual peer reaper.

server/server.go

@@ -156,11 +156,23 @@ func writeResponse(w http.ResponseWriter, values string) {
 
 // RunServer spins up the server and muxes the routes.
 func RunServer() {
+	// Load the config and initiate a `SQLStore` implementation.
+	var sqlObj db.SQLStore
+	cfg := config.LoadConfig()
+
+	if cfg.DBChoice == "mysql" {
+		sqlObj = new(sqlStoreImpl.MySQLStore)
+	} else if cfg.DBChoice == "postgres" {
+		sqlObj = new(sqlStoreImpl.PostgresStore)
+	} else {
+		panic("Invalid config value for DBChoice. Must be either postgres or MySQL.")
+	}
+
 	app := applicationContext{
-		config:          config.LoadConfig(),
+		config:          cfg,
 		trackerLevel:    a.RATIOLESS,
 		peerStoreClient: new(redisPeerStoreImpl.RedisStore),
-		sqlObj:          new(sqlStoreImpl.MySQLStore),
+		sqlObj:          sqlObj,
 	}
 
 	mux := http.NewServeMux()