-
Notifications
You must be signed in to change notification settings - Fork 17
/
mysql.go
161 lines (129 loc) · 4.27 KB
/
mysql.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
package mysql
import (
"database/sql"
"fmt"
"github.com/GrappigPanda/notorious/config"
"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"
)
// OpenConnection does as its name dictates and opens a connection to the
// MysqlHost listed in the config
func OpenConnection() (db *gorm.DB, err error) {
c := config.LoadConfig()
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 {
return nil, fmt.Errorf("Failed to open connection to MySQL: %v", err)
}
return db, nil
}
// InitDB initializes database tables.
func InitDB(dbConn *gorm.DB) {
dbConn = assertOpenConnection(dbConn)
dbConn.CreateTable(&schemas.WhiteTorrent{})
dbConn.CreateTable(&schemas.Torrent{})
dbConn.CreateTable(&schemas.TrackerStats{})
dbConn.CreateTable(&schemas.PeerStats{})
}
// GetTorrent retrieves a torrent by its infoHash from the generic torrent
// table in the database. Note: there's also a whitelisted torrent table
// (`WhiteTorrent`).
func GetTorrent(dbConn *gorm.DB, infoHash string) (t *schemas.Torrent, err error) {
dbConn = assertOpenConnection(dbConn)
t = &schemas.Torrent{}
dbConn.Where("info_hash = ?", infoHash).Find(&t)
return
}
// GetWhitelistedTorrent Retrieves a single whitelisted torrent by its infoHash
func GetWhitelistedTorrent(dbConn *gorm.DB, infoHash string) (t *schemas.WhiteTorrent, err error) {
dbConn = assertOpenConnection(dbConn)
t = &schemas.WhiteTorrent{}
x := dbConn.Where("info_hash = ?", infoHash).First(&t)
if x.Error != nil {
err = x.Error
}
return
}
// UpdateStats Handles updating statistics relevant to our tracker.
func UpdateStats(dbConn *gorm.DB, uploaded uint64, downloaded uint64) {
dbConn = assertOpenConnection(dbConn)
ts := &schemas.TrackerStats{}
dbConn.First(&ts)
dbConn.Model(&ts).Updates(
schemas.TrackerStats{
Uploaded: ts.Uploaded + int64(uploaded),
Downloaded: ts.Downloaded + int64(downloaded),
})
return
}
// UpdateTorrentStats Handles updating statistics relevant to our tracker.
func UpdateTorrentStats(dbConn *gorm.DB, seederDelta int64, leecherDelta int64) {
dbConn = assertOpenConnection(dbConn)
t := &schemas.Torrent{}
dbConn.First(&t)
dbConn.Model(&t).Updates(
schemas.TrackerStats{
Uploaded: t.Seeders + seederDelta,
Downloaded: t.Leechers + leecherDelta,
})
return
}
// UpdatePeerStats handles updating peer info like hits per ip, downloaded
// amount, uploaded amounts.
func UpdatePeerStats(dbConn *gorm.DB, uploaded uint64, downloaded uint64, ip string) {
dbConn = assertOpenConnection(dbConn)
ps := &schemas.PeerStats{Ip: ip}
dbConn.First(&ps)
fmt.Printf("%v", *ps)
dbConn.Model(&ps).UpdateColumn(map[string]interface{}{
"Uploaded": ps.Uploaded + int64(uploaded),
"Downloaded": ps.Downloaded + int64(downloaded),
})
return
}
// GetWhitelistedTorrents allows us to retrieve all of the white listed
// torrents. Mostly used for populating the Redis KV storage with all of our
// whitelisted torrents.
func GetWhitelistedTorrents(dbConn *gorm.DB) (x *sql.Rows, err error) {
dbConn = assertOpenConnection(dbConn)
x, err = dbConn.Table("white_torrents").Rows()
if err != nil {
return
}
return
}
// ScrapeTorrent supports the Scrape convention
func ScrapeTorrent(dbConn *gorm.DB, infoHash string) (torrent *schemas.Torrent) {
dbConn = assertOpenConnection(dbConn)
dbConn.Where("info_hash = ?", infoHash).First(&torrent)
return
}
// formatConnectStrings concatenates the data from the config file into a
// usable MySQL connection string.
func formatConnectString(c config.ConfigStruct) string {
return fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?parseTime=true",
c.DBUser,
c.DBPass,
c.DBHost,
c.DBPort,
c.DBName,
)
}
// assertOpenConnection handles asserting a connection passed into a sql
// function is open, not nil. If nil, we'll create a new connection.
func assertOpenConnection(db *gorm.DB) *gorm.DB {
var err error
if db == nil {
db, err = OpenConnection()
if err != nil {
err = err
}
}
return db
}