Completed
Push — master ( 95a2a0...997931 )
by Viktor
02:11
created

bot.*DBWorker.AddRadioStation   A

Complexity

Conditions 1

Size

Total Lines 4
Code Lines 4

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
eloc 4
nop 3
dl 0
loc 4
rs 10
c 0
b 0
f 0
1
package bot
2
3
import (
4
	"fmt"
5
	"github.com/bwmarrin/discordgo"
6
	"github.com/globalsign/mgo"
7
	"github.com/globalsign/mgo/bson"
8
	"os"
9
	"time"
10
)
11
12
// DBWorker MongoDB instance
13
type DBWorker struct {
14
	DBSession *mgo.Session
15
	DBName    string
16
}
17
18
type dbLog struct {
19
	Date   time.Time
20
	Text   string
21
	Module string
22
	Guild  string
23
}
24
25
// GuildData contains data about guild settings
26
type GuildData struct {
27
	ID          string
28
	WeatherCity string
29
	NewsCounty  string
30
	Language    string
31
	Timezone    int
32
	EmbedColor  int
33
	VoiceVolume float32
34
	Greeting    string
35
}
36
37
// GuildsMap contains guilds settings
38
type GuildsMap struct {
39
	Guilds map[string]*GuildData
40
}
41
42
// RadioStation contains info about radio station
43
type RadioStation struct {
44
	Name string
45
	URL  string
46
	Key  string
47
}
48
49
// NewDBSession creates new MongoDB instance
50
func NewDBSession(dbname string) *DBWorker {
51
	session, err := mgo.Dial(os.Getenv("MONGO_CONN"))
52
	if err != nil {
53
		fmt.Printf("Mongo connection error: %v", err)
54
	}
55
	count, err := session.DB("dtbot").C("logs").Count()
56
	if err != nil {
57
		fmt.Println("DB_ERR: ", err)
58
	}
59
	fmt.Printf("Mongo connected\nLogs in base: %v\n", count)
60
	return &DBWorker{DBSession: session, DBName: dbname}
61
}
62
63
// InitGuilds initialize guilds in database
64
func (db *DBWorker) InitGuilds(sess *discordgo.Session, conf *Config) *GuildsMap {
65
	var data = &GuildsMap{Guilds: make(map[string]*GuildData)}
66
	var loaded, initialized = 0, 0
67
	for _, guild := range sess.State.Guilds {
68
		count, err := db.DBSession.DB(db.DBName).C("guilds").Find(bson.M{"id": guild.ID}).Count()
69
		if err != nil {
70
			fmt.Println("Mongo: ", err)
71
		}
72
		if count == 0 {
73
			newData := &GuildData{
74
				ID:          guild.ID,
75
				WeatherCity: conf.Weather.City,
76
				NewsCounty:  conf.News.Country,
77
				Language:    conf.General.Language,
78
				Timezone:    conf.General.Timezone,
79
				EmbedColor:  conf.General.EmbedColor,
80
				VoiceVolume: conf.Voice.Volume,
81
				Greeting:    "",
82
			}
83
			_ = db.DBSession.DB(db.DBName).C("guilds").Insert(newData)
84
			data.Guilds[guild.ID] = newData
85
			initialized++
86
		} else {
87
			var newData = &GuildData{}
88
			_ = db.DBSession.DB(db.DBName).C("guilds").Find(bson.M{"id": guild.ID}).One(newData)
89
			if err != nil {
90
				fmt.Println("Mongo: ", err)
91
				continue
92
			}
93
			data.Guilds[guild.ID] = newData
94
			loaded++
95
		}
96
	}
97
	fmt.Printf("Guilds loaded [%v], initialized [%v]\n", loaded, initialized)
98
	return data
99
}
100
101
// InitNewGuild creates new guild in database
102
func (db *DBWorker) InitNewGuild(guildID string, conf *Config, data *GuildsMap) {
103
	newData := &GuildData{
104
		ID:          guildID,
105
		WeatherCity: conf.Weather.City,
106
		NewsCounty:  conf.News.Country,
107
		Language:    conf.General.Language,
108
		Timezone:    conf.General.Timezone,
109
		EmbedColor:  conf.General.EmbedColor,
110
		VoiceVolume: conf.Voice.Volume,
111
		Greeting:    "",
112
	}
113
	_ = db.DBSession.DB(db.DBName).C("guilds").Insert(newData)
114
	data.Guilds[guildID] = newData
115
}
116
117
// Log saves log in database
118
func (db *DBWorker) Log(module, guildID, text string) {
119
	_ = db.DBSession.DB(db.DBName).C("logs").Insert(dbLog{Date: time.Now(), Text: text, Module: module, Guild: guildID})
120
}
121
122
// LogGet returns last N log rows
123
func (db *DBWorker) LogGet(count int) []dbLog {
0 ignored issues
show
introduced by
exported method LogGet returns unexported type []bot.dbLog, which can be annoying to use
Loading history...
124
	var log = make([]dbLog, count)
125
	_ = db.DBSession.DB(db.DBName).C("logs").Find(nil).Sort("-$natural").Limit(count).All(&log)
126
	return log
127
}
128
129
// Guilds returns guilds collection from mongodb
130
func (db *DBWorker) Guilds() *mgo.Collection {
131
	return db.DBSession.DB(db.DBName).C("guilds")
132
}
133
134
// GetTwitchStreams returns twitch streams from mongodb
135
func (db *DBWorker) GetTwitchStreams(guildID string) map[string]*TwitchStream {
136
	streams := []TwitchStream{}
137
	err := db.DBSession.DB(db.DBName).C("streams").Find(bson.M{"guild": guildID}).All(&streams)
138
	if err != nil {
139
		fmt.Println("Mongo: ", err)
140
	}
141
	var newMap = make(map[string]*TwitchStream)
142
	for i, s := range streams {
143
		newMap[s.Login] = &streams[i]
144
	}
145
	return newMap
146
}
147
148
// UpdateStream updates stream in mongodb
149
func (db *DBWorker) UpdateStream(stream *TwitchStream) {
150
	err := db.DBSession.DB(db.DBName).C("streams").
151
		Update(
152
			bson.M{"guild": stream.Guild, "login": stream.Login},
153
			bson.M{"$set": bson.M{"isonline": stream.IsOnline}})
154
	if err != nil {
155
		fmt.Println(err.Error())
156
	}
157
}
158
159
// AddStream adds stream in mongodb
160
func (db *DBWorker) AddStream(stream *TwitchStream) {
161
	_ = db.DBSession.DB(db.DBName).C("streams").Insert(stream)
162
}
163
164
// RemoveStream removes stream from mongodb
165
func (db *DBWorker) RemoveStream(stream *TwitchStream) {
166
	_ = db.DBSession.DB(db.DBName).C("streams").Remove(bson.M{"login": stream.Login, "guild": stream.Guild})
167
}
168
169
// GetRadioStations gets stations from database and returns slice of them
170
func (db *DBWorker) GetRadioStations() []RadioStation {
171
	stations := []RadioStation{}
172
	err := db.DBSession.DB(db.DBName).C("stations").Find(nil).All(&stations)
173
	if err != nil {
174
		fmt.Println("Mongo: ", err)
175
	}
176
	return stations
177
}
178
179
// GetRadioStationByKey returns one station by key
180
func (db *DBWorker) GetRadioStationByKey(key string) (*RadioStation, error) {
181
	station := RadioStation{}
182
	err := db.DBSession.DB(db.DBName).C("stations").Find(bson.M{"key":key}).One(&station)
183
	if err != nil {
184
		fmt.Println("Mongo: ", err)
185
		return nil, fmt.Errorf("station not found")
186
	}
187
	return &station, nil
188
}
189
190
// RemoveRadioStation removes radio station by key
191
func (db *DBWorker) RemoveRadioStation(key string) error {
192
	err := db.DBSession.DB(db.DBName).C("stations").Remove(bson.M{"key": key})
193
	return err
194
}
195
196
// AddRadioStation adds new radio station
197
func (db *DBWorker) AddRadioStation(name, url, key string) error {
198
	station := RadioStation{Name:name, URL: url, Key: key}
199
	err := db.DBSession.DB(db.DBName).C("streams").Insert(&station)
200
	return err
201
}