1package storage 2 3import ( 4 "context" 5 "fmt" 6 "github.com/automuteus/utils/pkg/premium" 7 "github.com/georgysavva/scany/pgxscan" 8 "github.com/jackc/pgx/v4/pgxpool" 9 "io/ioutil" 10 "log" 11 "os" 12 "strconv" 13 "time" 14) 15 16type PsqlInterface struct { 17 Pool *pgxpool.Pool 18 19 // TODO does this require a lock? How should stuff be written/read from psql in an async way? Is this even a concern? 20 //https://brandur.org/postgres-connections 21} 22 23func ConstructPsqlConnectURL(addr, username, password string) string { 24 return fmt.Sprintf("postgres://%s?user=%s&password=%s", addr, username, password) 25} 26 27type PsqlParameters struct { 28 Addr string 29 Username string 30 Password string 31} 32 33func (psqlInterface *PsqlInterface) Init(addr string) error { 34 dbpool, err := pgxpool.Connect(context.Background(), addr) 35 if err != nil { 36 return err 37 } 38 psqlInterface.Pool = dbpool 39 return nil 40} 41 42func (psqlInterface *PsqlInterface) LoadAndExecFromFile(filepath string) error { 43 f, err := os.Open(filepath) 44 if err != nil { 45 return err 46 } 47 defer f.Close() 48 49 bytes, err := ioutil.ReadAll(f) 50 if err != nil { 51 return err 52 } 53 tag, err := psqlInterface.Pool.Exec(context.Background(), string(bytes)) 54 if err != nil { 55 return err 56 } 57 log.Println(tag.String()) 58 return nil 59} 60 61func (psqlInterface *PsqlInterface) insertGuild(guildID uint64, guildName string) error { 62 _, err := psqlInterface.Pool.Exec(context.Background(), "INSERT INTO guilds VALUES ($1, $2, 0);", guildID, guildName) 63 return err 64} 65 66func (psqlInterface *PsqlInterface) GetGuild(guildID uint64) (*PostgresGuild, error) { 67 guilds := []*PostgresGuild{} 68 err := pgxscan.Select(context.Background(), psqlInterface.Pool, &guilds, "SELECT * FROM guilds WHERE guild_id=$1", guildID) 69 if err != nil { 70 return nil, err 71 } 72 73 if len(guilds) > 0 { 74 return guilds[0], nil 75 } 76 return nil, nil 77} 78 79func (psqlInterface *PsqlInterface) insertUser(userID uint64) error { 80 _, err := psqlInterface.Pool.Exec(context.Background(), "INSERT INTO users VALUES ($1, true);", userID) 81 return err 82} 83 84func (psqlInterface *PsqlInterface) OptUserByString(userID string, opt bool) (bool, error) { 85 uid, err := strconv.ParseUint(userID, 10, 64) 86 if err != nil { 87 return false, err 88 } 89 user, err := psqlInterface.EnsureUserExists(uid) 90 if err != nil { 91 return false, err 92 } 93 if user.Opt == opt { 94 return false, nil 95 } 96 _, err = psqlInterface.Pool.Exec(context.Background(), "UPDATE users SET opt = $1 WHERE user_id = $2;", opt, uid) 97 if err != nil { 98 return false, err 99 } 100 if !opt { 101 _, err = psqlInterface.Pool.Exec(context.Background(), "UPDATE game_events SET user_id = NULL WHERE user_id = $1;", uid) 102 if err != nil { 103 log.Println(err) 104 } 105 106 _, err = psqlInterface.Pool.Exec(context.Background(), "DELETE FROM users_games WHERE user_id = $1;", uid) 107 if err != nil { 108 log.Println(err) 109 } 110 } 111 112 return true, nil 113} 114 115func (psqlInterface *PsqlInterface) GetUserByString(userID string) (*PostgresUser, error) { 116 uid, err := strconv.ParseUint(userID, 10, 64) 117 if err != nil { 118 return nil, err 119 } 120 return psqlInterface.GetUser(uid) 121} 122 123func (psqlInterface *PsqlInterface) GetUser(userID uint64) (*PostgresUser, error) { 124 users := []*PostgresUser{} 125 err := pgxscan.Select(context.Background(), psqlInterface.Pool, &users, "SELECT * FROM users WHERE user_id = $1", userID) 126 if err != nil { 127 return nil, err 128 } 129 130 if len(users) > 0 { 131 return users[0], nil 132 } 133 return nil, nil 134} 135 136func (psqlInterface *PsqlInterface) GetGame(guildID, connectCode, matchID string) (*PostgresGame, error) { 137 games := []*PostgresGame{} 138 err := pgxscan.Select(context.Background(), psqlInterface.Pool, &games, "SELECT * FROM games WHERE guild_id = $1 AND game_id = $2 AND connect_code = $3;", guildID, matchID, connectCode) 139 if err != nil { 140 return nil, err 141 } 142 if len(games) > 0 { 143 return games[0], nil 144 } 145 return nil, nil 146} 147 148func (psqlInterface *PsqlInterface) GetGameEvents(matchID string) ([]*PostgresGameEvent, error) { 149 events := []*PostgresGameEvent{} 150 err := pgxscan.Select(context.Background(), psqlInterface.Pool, &events, "SELECT * FROM game_events WHERE game_id = $1 ORDER BY event_id ASC;", matchID) 151 if err != nil { 152 return nil, err 153 } 154 return events, nil 155} 156 157func (psqlInterface *PsqlInterface) insertGame(game *PostgresGame) (uint64, error) { 158 t, err := psqlInterface.Pool.Query(context.Background(), "INSERT INTO games VALUES (DEFAULT, $1, $2, $3, $4, $5) RETURNING game_id;", game.GuildID, game.ConnectCode, game.StartTime, game.WinType, game.EndTime) 159 if t != nil { 160 for t.Next() { 161 g := uint64(0) 162 err := t.Scan(&g) 163 164 if err != nil { 165 log.Println(err) 166 t.Close() 167 return 0, err 168 } 169 t.Close() 170 return g, nil 171 } 172 } 173 return 0, err 174} 175 176func (psqlInterface *PsqlInterface) updateGame(gameID int64, winType int16, endTime int64) error { 177 _, err := psqlInterface.Pool.Exec(context.Background(), "UPDATE games SET (win_type, end_time) = ($1, $2) WHERE game_id = $3;", winType, endTime, gameID) 178 return err 179} 180 181func (psqlInterface *PsqlInterface) insertPlayer(player *PostgresUserGame) error { 182 _, err := psqlInterface.Pool.Exec(context.Background(), "INSERT INTO users_games VALUES ($1, $2, $3, $4, $5, $6, $7);", player.UserID, player.GuildID, player.GameID, player.PlayerName, player.PlayerColor, player.PlayerRole, player.PlayerWon) 183 return err 184} 185 186const SecsInADay = 86400 187 188func (psqlInterface *PsqlInterface) GetGuildPremiumStatus(guildID string) (premium.Tier, int) { 189 // self-hosting; only return the true guild status if this variable is set 190 if os.Getenv("AUTOMUTEUS_OFFICIAL") == "" { 191 return premium.SelfHostTier, premium.NoExpiryCode 192 } 193 194 gid, err := strconv.ParseUint(guildID, 10, 64) 195 if err != nil { 196 log.Println(err) 197 return premium.FreeTier, 0 198 } 199 200 guild, err := psqlInterface.GetGuild(gid) 201 if err != nil { 202 return premium.FreeTier, 0 203 } 204 205 daysRem := premium.NoExpiryCode 206 207 if guild.TxTimeUnix != nil { 208 diff := time.Now().Unix() - int64(*guild.TxTimeUnix) 209 // 31 - days elapsed 210 daysRem = int(premium.SubDays - (diff / SecsInADay)) 211 } 212 213 return premium.Tier(guild.Premium), daysRem 214} 215 216func (psqlInterface *PsqlInterface) EnsureGuildExists(guildID uint64, guildName string) (*PostgresGuild, error) { 217 guild, err := psqlInterface.GetGuild(guildID) 218 219 if guild == nil { 220 err := psqlInterface.insertGuild(guildID, guildName) 221 if err != nil { 222 return nil, err 223 } 224 return psqlInterface.GetGuild(guildID) 225 } 226 return guild, err 227} 228 229func (psqlInterface *PsqlInterface) EnsureUserExists(userID uint64) (*PostgresUser, error) { 230 user, err := psqlInterface.GetUser(userID) 231 232 if user == nil { 233 err := psqlInterface.insertUser(userID) 234 if err != nil { 235 log.Println(err) 236 } 237 return psqlInterface.GetUser(userID) 238 } 239 return user, err 240} 241 242func (psqlInterface *PsqlInterface) AddInitialGame(game *PostgresGame) (uint64, error) { 243 return psqlInterface.insertGame(game) 244} 245 246func (psqlInterface *PsqlInterface) AddEvent(event *PostgresGameEvent) error { 247 if event.UserID == nil { 248 _, err := psqlInterface.Pool.Exec(context.Background(), "INSERT INTO game_events VALUES (DEFAULT, NULL, $1, $2, $3, $4);", event.GameID, event.EventTime, event.EventType, event.Payload) 249 return err 250 } 251 _, err := psqlInterface.Pool.Exec(context.Background(), "INSERT INTO game_events VALUES (DEFAULT, $1, $2, $3, $4, $5);", event.UserID, event.GameID, event.EventTime, event.EventType, event.Payload) 252 return err 253} 254 255// make sure to call the relevant "ensure" methods before this one... 256func (psqlInterface *PsqlInterface) UpdateGameAndPlayers(gameID int64, winType int16, endTime int64, players []*PostgresUserGame) error { 257 err := psqlInterface.updateGame(gameID, winType, endTime) 258 if err != nil { 259 return err 260 } 261 262 for _, player := range players { 263 err := psqlInterface.insertPlayer(player) 264 if err != nil { 265 log.Println(err) 266 } 267 } 268 269 return nil 270} 271 272func (psqlInterface *PsqlInterface) Close() { 273 psqlInterface.Pool.Close() 274} 275