diff --git a/database/database.go b/database/database.go index f63eea3..3fd769a 100644 --- a/database/database.go +++ b/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`.") + } } diff --git a/database/impl/mysqlStore.go b/database/impl/mysqlStore.go index 8c8470a..a000930 100644 --- a/database/impl/mysqlStore.go +++ b/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, } diff --git a/database/impl/postgresStore.go b/database/impl/postgresStore.go index f1a7c82..05d842b 100644 --- a/database/impl/postgresStore.go +++ b/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)) } } diff --git a/database/impl/sqlstore.go b/database/impl/sqlstore.go new file mode 100644 index 0000000..2034544 --- /dev/null +++ b/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 +) diff --git a/reaper/reaper.go b/reaper/reaper.go index 6e94d81..21b09ed 100644 --- a/reaper/reaper.go +++ b/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. diff --git a/server/server.go b/server/server.go index 8cb6160..63c1240 100644 --- a/server/server.go +++ b/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,