Merge pull request #183 from GrappigPanda/issue-181

Fixes #181

database/database.go

@@ -1,11 +1,9 @@
 package db
 
 import (
-	"database/sql"
 	"github.com/GrappigPanda/notorious/config"
 	"github.com/GrappigPanda/notorious/database/mysql"
 	"github.com/GrappigPanda/notorious/database/postgres"
-	"github.com/GrappigPanda/notorious/database/schemas"
 	"github.com/jinzhu/gorm"
 	// We use a blank import here because I'm afraid of breaking anything
 	_ "github.com/jinzhu/gorm/dialects/mysql"
@@ -29,39 +27,14 @@ func InitDB(config *config.ConfigStruct) {
 	}
 }
 
-// PeerDeltaEvent allows us to set an event to handle how we're going to update
-// teh database.
-type PeerDeltaEvent int
+func OpenDBChoiceConnection() (*gorm.DB, error) {
+	cfg := config.LoadConfig()
 
-const (
-	// PEERUPDATE represents a change to a peer, so we'll update a tracker
-	// user's ratio.
-	PEERUPDATE PeerDeltaEvent = iota
-	// TRACKERUPDATE handles updating total tracker stats
-	TRACKERUPDATE
-	// TORRENTUPDATE represents the changes to a specific torrent where we
-	// update total upload/download for the torrent itself.
-	TORRENTUPDATE
-)
-
-// PeerTrackerDelta handles holding data to be updated by the `UpdateConsumer`.
-type PeerTrackerDelta struct {
-	Uploaded   uint64
-	Downloaded uint64
-	IP         string
-	Event      PeerDeltaEvent
-}
-
-// SQLStore is the base implementation for a database which will be used to
-// store stats and retrieve whitelisted torrents.
-type SQLStore interface {
-	OpenConnection() (*gorm.DB, error)
-	GetTorrent(string) (*schemas.Torrent, error)
-	GetWhitelistedTorrent(string) (*schemas.WhiteTorrent, error)
-	UpdateStats(uint64, uint64)
-	UpdateTorrentStats(int64, int64)
-	ScrapeTorrent(string) *schemas.Torrent
-	GetWhitelistedTorrents() (*sql.Rows, error)
-	UpdatePeerStats(uint64, uint64, string)
-	HandlePeerUpdates() chan PeerTrackerDelta
+	if cfg.DBChoice == "mysql" {
+		return mysql.OpenConnection()
+	} else if cfg.DBChoice == "postgres" {
+		return postgres.OpenConnection()
+	} else {
+		panic("Invalid database choice found for `OpenDBChoiceConnection`.")
+	}
 }

database/impl/mysqlStore.go

@@ -2,7 +2,6 @@ package sqlStoreImpl
 
 import (
 	"database/sql"
-	"github.com/GrappigPanda/notorious/database"
 	"github.com/GrappigPanda/notorious/database/mysql"
 	"github.com/GrappigPanda/notorious/database/schemas"
 	"github.com/jinzhu/gorm"
@@ -13,7 +12,7 @@ import (
 // MySQLStore represents the mysql implementation of `SQLStore`
 type MySQLStore struct {
 	dbPool         *gorm.DB
-	UpdateConsumer chan db.PeerTrackerDelta
+	UpdateConsumer chan PeerTrackerDelta
 }
 
 // InitMySQLStore Creates a `MySQLStore` object and initiates all necessary
@@ -42,18 +41,18 @@ func (m *MySQLStore) OpenConnection() (*gorm.DB, error) {
 // HandlePeerUpdates handles listening and aggregating peer updates. THis
 // allows block/asynchronous consumption of peer updates, rather than updating
 // the remote database at the end of every request.
-func (m *MySQLStore) HandlePeerUpdates() chan db.PeerTrackerDelta {
-	peerUpdatesChan := make(chan db.PeerTrackerDelta)
+func (m *MySQLStore) HandlePeerUpdates() chan PeerTrackerDelta {
+	peerUpdatesChan := make(chan PeerTrackerDelta)
 
 	go func() {
 		for {
 			update := <-peerUpdatesChan
 			switch update.Event {
-			case db.PEERUPDATE:
+			case PEERUPDATE:
 				m.UpdatePeerStats(update.Uploaded, update.Downloaded, update.IP)
-			case db.TRACKERUPDATE:
+			case TRACKERUPDATE:
 				m.UpdateStats(update.Uploaded, update.Downloaded)
-			case db.TORRENTUPDATE:
+			case TORRENTUPDATE:
 				m.UpdateTorrentStats(int64(update.Uploaded), int64(update.Downloaded))
 			}
 		}
@@ -89,7 +88,7 @@ func (m *MySQLStore) UpdatePeerStats(uploaded uint64, downloaded uint64, ip stri
 
 // UpdateStats wraps `mysql.UpdateStats`.
 func (m *MySQLStore) UpdateStats(uploaded uint64, downloaded uint64) {
-	m.UpdateConsumer <- db.PeerTrackerDelta{
+	m.UpdateConsumer <- PeerTrackerDelta{
 		Uploaded:   uploaded,
 		Downloaded: downloaded,
 	}

database/impl/postgresStore.go

@@ -2,7 +2,6 @@ package sqlStoreImpl
 
 import (
 	"database/sql"
-	"github.com/GrappigPanda/notorious/database"
 	"github.com/GrappigPanda/notorious/database/postgres"
 	"github.com/GrappigPanda/notorious/database/schemas"
 	"github.com/jinzhu/gorm"
@@ -41,18 +40,18 @@ func (m *PostgresStore) OpenConnection() (*gorm.DB, error) {
 // HandlePeerUpdates handles listening and aggregating peer updates. THis
 // allows block/asynchronous consumption of peer updates, rather than updating
 // the remote database at the end of every request.
-func (m *PostgresStore) HandlePeerUpdates() chan db.PeerTrackerDelta {
-	peerUpdatesChan := make(chan db.PeerTrackerDelta)
+func (m *PostgresStore) HandlePeerUpdates() chan PeerTrackerDelta {
+	peerUpdatesChan := make(chan PeerTrackerDelta)
 
 	go func() {
 		for {
 			update := <-peerUpdatesChan
 			switch update.Event {
-			case db.PEERUPDATE:
+			case PEERUPDATE:
 				m.UpdatePeerStats(update.Uploaded, update.Downloaded, update.IP)
-			case db.TRACKERUPDATE:
+			case TRACKERUPDATE:
 				m.UpdateStats(update.Uploaded, update.Downloaded)
-			case db.TORRENTUPDATE:
+			case TORRENTUPDATE:
 				m.UpdateTorrentStats(int64(update.Uploaded), int64(update.Downloaded))
 			}
 		}

database/impl/sqlstore.go

@@ -0,0 +1,56 @@
+package sqlStoreImpl
+
+import (
+	"database/sql"
+	"github.com/GrappigPanda/notorious/config"
+	"github.com/GrappigPanda/notorious/database/schemas"
+	"github.com/jinzhu/gorm"
+)
+
+// SQLStore is the base implementation for a database which will be used to
+// store stats and retrieve whitelisted torrents.
+type SQLStore interface {
+	OpenConnection() (*gorm.DB, error)
+	GetTorrent(string) (*schemas.Torrent, error)
+	GetWhitelistedTorrent(string) (*schemas.WhiteTorrent, error)
+	UpdateStats(uint64, uint64)
+	UpdateTorrentStats(int64, int64)
+	ScrapeTorrent(string) *schemas.Torrent
+	GetWhitelistedTorrents() (*sql.Rows, error)
+	UpdatePeerStats(uint64, uint64, string)
+	HandlePeerUpdates() chan PeerTrackerDelta
+}
+
+func InitSQLStoreByDBChoice() SQLStore {
+	cfg := config.LoadConfig()
+	if cfg.DBChoice == "mysql" {
+		return new(MySQLStore)
+	} else if cfg.DBChoice == "postgres" {
+		return new(PostgresStore)
+	} else {
+		panic("Invalid database choice found for `InitSQLStoreByDBChoice`.")
+	}
+}
+
+// PeerTrackerDelta handles holding data to be updated by the `UpdateConsumer`.
+type PeerTrackerDelta struct {
+	Uploaded   uint64
+	Downloaded uint64
+	IP         string
+	Event      PeerDeltaEvent
+}
+
+// PeerDeltaEvent allows us to set an event to handle how we're going to update
+// teh database.
+type PeerDeltaEvent int
+
+const (
+	// PEERUPDATE represents a change to a peer, so we'll update a tracker
+	// user's ratio.
+	PEERUPDATE PeerDeltaEvent = iota
+	// TRACKERUPDATE handles updating total tracker stats
+	TRACKERUPDATE
+	// TORRENTUPDATE represents the changes to a specific torrent where we
+	// update total upload/download for the torrent itself.
+	TORRENTUPDATE
+)

reaper/reaper.go

@@ -2,9 +2,7 @@ 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/database/impl"
 	"github.com/GrappigPanda/notorious/peerStore/redis"
 	"gopkg.in/redis.v3"
 	"strconv"
@@ -85,26 +83,13 @@ func StartReapingScheduler(waitTime time.Duration) {
 			addedBy := new(string)
 			dateAdded := new(int64)
 
-			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)
-					}
+			sqlStore := sqlStoreImpl.InitSQLStoreByDBChoice()
+			whitelistedTorrents, err := sqlStore.GetWhitelistedTorrents()
+			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

@@ -4,7 +4,6 @@ import (
 	"fmt"
 	a "github.com/GrappigPanda/notorious/announce"
 	"github.com/GrappigPanda/notorious/config"
-	"github.com/GrappigPanda/notorious/database"
 	"github.com/GrappigPanda/notorious/database/impl"
 	"github.com/GrappigPanda/notorious/peerStore"
 	"github.com/GrappigPanda/notorious/peerStore/impl"
@@ -17,7 +16,7 @@ type applicationContext struct {
 	config          config.ConfigStruct
 	trackerLevel    int
 	peerStoreClient peerStore.PeerStore
-	sqlObj          db.SQLStore
+	sqlObj          sqlStoreImpl.SQLStore
 }
 
 type scrapeData struct {
@@ -157,19 +156,10 @@ 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.")
-	}
+	sqlObj := sqlStoreImpl.InitSQLStoreByDBChoice()
 
 	app := applicationContext{
-		config:          cfg,
+		config:          config.LoadConfig(),
 		trackerLevel:    a.RATIOLESS,
 		peerStoreClient: new(redisPeerStoreImpl.RedisStore),
 		sqlObj:          sqlObj,