Skip to content

Instantly share code, notes, and snippets.

@alexesDev
Created January 30, 2020 10:28
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save alexesDev/78a3482eb75791159cccc90e136eaea2 to your computer and use it in GitHub Desktop.
Save alexesDev/78a3482eb75791159cccc90e136eaea2 to your computer and use it in GitHub Desktop.
package main
import (
"encoding/json"
"fmt"
"log"
"net/http"
"regexp"
"time"
_ "go.uber.org/automaxprocs"
"github.com/caarlos0/env"
"github.com/go-pg/pg"
"github.com/go-telegram-bot-api/telegram-bot-api"
"golang.org/x/net/proxy"
)
type context struct {
TgToken string `env:"TG_TOKEN,required"`
PgConnectionString string `env:"PG_CONNECTION_STRING,required"`
Socks5Proxy string `env:"SOCKS5_PROXY"`
Bot *tgbotapi.BotAPI
DB *pg.DB
}
type raceRow struct {
tableName string `sql:"races"`
ID int64
ChatID int64
UserID int
InitialMessageID int
CreatedAt string
Duration string
Status string
Comment string
}
var goReplyKeyboard = tgbotapi.NewInlineKeyboardMarkup(
tgbotapi.NewInlineKeyboardRow(
tgbotapi.NewInlineKeyboardButtonData("Show remaining time", "left"),
),
)
var replyKeyboard = tgbotapi.NewReplyKeyboard(
tgbotapi.NewKeyboardButtonRow(
tgbotapi.NewKeyboardButton("/go"),
tgbotapi.NewKeyboardButton("/left"),
tgbotapi.NewKeyboardButton("/stats"),
tgbotapi.NewKeyboardButton("/cancel"),
),
tgbotapi.NewKeyboardButtonRow(
tgbotapi.NewKeyboardButton("/setcomment"),
tgbotapi.NewKeyboardButton("/getcomment"),
tgbotapi.NewKeyboardButton("/report"),
),
)
func (ctx *context) getBotAPI() (*tgbotapi.BotAPI, error) {
if ctx.Socks5Proxy != "" {
dialer, err := proxy.SOCKS5("tcp", ctx.Socks5Proxy, nil, proxy.Direct)
if err != nil {
log.Fatalf("can't connect to the proxy: %s", err)
}
httpTransport := &http.Transport{}
httpClient := &http.Client{Transport: httpTransport}
httpTransport.Dial = dialer.Dial
return tgbotapi.NewBotAPIWithClient(ctx.TgToken, httpClient)
}
return tgbotapi.NewBotAPI(ctx.TgToken)
}
func (ctx *context) getRaceComment(userID int) string {
var row struct {
Count int64
}
_, err := ctx.DB.QueryOne(&row, `
select count(*)
from races
where user_id = ?
and status = 'finished'
and created_at::date = now()::date
`, userID)
if err != nil {
log.Println(err)
return ""
}
if row.Count == 2 {
return "Кажется вы сели серьёзно поработать. План на сегодня: 14 помидор."
}
if row.Count == 7 {
return "Вы собрали половину запланированных помидор."
}
if row.Count == 12 {
return "Ещё два помидора и план готов."
}
if row.Count == 14 {
return "14 помидор - цель достигнута и можно отдохнуть."
}
return ""
}
func (ctx *context) checkFinished() {
var rows []raceRow
for {
_, err := ctx.DB.Query(&rows, `
update races set status = 'finished'
where status = 'pending' and created_at + duration < now()
returning chat_id, user_id, initial_message_id
`)
if err != nil {
log.Println(err)
time.Sleep(1000 * time.Millisecond)
continue
}
for _, row := range rows {
msg := tgbotapi.NewMessage(row.ChatID, "")
msg.Text = "Race finished."
msg.ReplyToMessageID = row.InitialMessageID
comment := ctx.getRaceComment(row.UserID)
if comment != "" {
msg.Text += " " + comment
}
if _, err := ctx.Bot.Send(msg); err != nil {
log.Println("failed to send: ", err)
}
}
time.Sleep(1000 * time.Millisecond)
}
}
type forgetRows struct {
ChatID int64
}
func (ctx *context) forgetRestart() {
var rows []forgetRows
for {
_, err := ctx.DB.Query(&rows, `
select f.chat_id
from races f
left join races p on p.user_id = f.user_id and p.status = 'pending' and p.created_at > f.created_at
where p.id is null
and f.status = 'finished'
and f.created_at + f.duration > now() - interval '11 minutes'
`)
if err != nil {
log.Println(err)
time.Sleep(5 * time.Minute)
continue
}
for _, row := range rows {
msg := tgbotapi.NewMessage(row.ChatID, "")
msg.Text = "Прошло 5 минут. Если вы работаете, то забыли сказать /go"
if _, err := ctx.Bot.Send(msg); err != nil {
log.Println("failed to send: ", err)
}
}
time.Sleep(5 * time.Minute)
}
}
func statusReply(msg *tgbotapi.MessageConfig, db *pg.DB, from int) {
var row struct {
MinutesLeft int
}
_, err := db.QueryOne(&row, `
select ceil(extract('epoch' from duration - (now() - created_at)) / 60) as minutes_left
from races
where status = 'pending'
and user_id = ?
`, from)
if err != nil {
log.Println(err)
msg.Text = "An active race not found."
} else {
msg.Text = fmt.Sprintf("%d minutes left.", row.MinutesLeft)
}
}
type dataRow struct {
ID int `json:"id"`
ChatID int `json:"chatId"`
CreatedAt string `json:"createdAt"`
Status string `json:"status"`
Duration string `json:"duration"`
UserName string `json:"userName"`
Index float32 `json:"index"`
}
func logRequest(handler http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
handler.ServeHTTP(w, r)
})
}
func runServer(db *pg.DB) {
fs := http.FileServer(http.Dir("dist"))
http.Handle("/", fs)
http.HandleFunc("/data", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Access-Control-Allow-Origin", "http://localhost:1234")
if r.Method == "OPTIONS" {
return
}
var rows []dataRow
_, err := db.Query(&rows, `
select r.id, r.chat_id, r.created_at, r.status, r.duration
, coalesce(u.user_name, u.first_name || ' ' || u.last_name) as user_name
, extract(epoch from created_at - now()) / 3600 + 24 as index
from races r
join users u on u.id = r.user_id
where created_at > now() - interval '1 day'
order by r.created_at
`)
if err != nil {
log.Println(err)
http.Error(w, http.StatusText(500), 500)
return
}
if err := json.NewEncoder(w).Encode(rows); err != nil {
log.Println(err)
http.Error(w, http.StatusText(500), 500)
}
})
log.Fatal(http.ListenAndServe(":3333", logRequest(http.DefaultServeMux)))
}
func getLastComment(db *pg.DB, userID int, chatID int64) string {
var row struct {
Comment string
}
_, err := db.QueryOne(&row, `
select comment
from races
where user_id = ?
and chat_id = ?
order by created_at desc
limit 1
`, userID, chatID)
if err != nil && err != pg.ErrNoRows {
log.Println(err)
}
return row.Comment
}
func main() {
ctx := context{}
err := env.Parse(&ctx)
if err != nil {
log.Fatal(err)
}
dbOptions, err := pg.ParseURL(ctx.PgConnectionString)
if err != nil {
panic(err)
}
db := pg.Connect(dbOptions)
defer func() {
if err := db.Close(); err != nil {
panic(err)
}
}()
_, err = db.Exec("select 1")
if err != nil {
panic(err)
}
ctx.DB = db
ctx.Bot, err = ctx.getBotAPI()
if err != nil {
log.Fatal(err)
}
bot := ctx.Bot
log.Printf("Authorized on account %s", bot.Self.UserName)
u := tgbotapi.NewUpdate(0)
u.Timeout = 60
updates, err := bot.GetUpdatesChan(u)
go ctx.checkFinished()
go ctx.forgetRestart()
go runServer(db)
nicknameR := regexp.MustCompile("@([a-zA-Z0-9_]+)")
for update := range updates {
if update.CallbackQuery != nil {
if update.CallbackQuery.Data == "left" {
msg := tgbotapi.NewMessage(update.CallbackQuery.Message.Chat.ID, "")
statusReply(&msg, db, update.CallbackQuery.From.ID)
if _, err := ctx.Bot.Send(msg); err != nil {
log.Println("failed to send: ", err)
}
}
}
if update.Message == nil { // ignore any non-Message Updates
continue
}
if update.Message.IsCommand() {
// Is a group
if update.Message.Chat.ID < 0 {
continue
}
msg := tgbotapi.NewMessage(update.Message.Chat.ID, "")
msg.ReplyToMessageID = update.Message.MessageID
msg.ReplyMarkup = replyKeyboard
switch update.Message.Command() {
case "go", "25":
comment := update.Message.CommandArguments()
if comment == "" {
comment = getLastComment(db, update.Message.From.ID, update.Message.Chat.ID)
}
err = db.Insert(&raceRow{
ChatID: update.Message.Chat.ID,
UserID: update.Message.From.ID,
InitialMessageID: update.Message.MessageID,
Comment: comment,
})
if err != nil {
if err.Error() == `ERROR #23505 duplicate key value violates unique constraint "race_pending_idx"` {
msg.Text = "You have an active race"
} else {
log.Println("go: ", err)
msg.Text = err.Error()
}
} else {
msg.Text = "Race started."
msg.ReplyMarkup = goReplyKeyboard
if comment != "" {
msg.Text += " Comment:\n" + comment
} else {
msg.Text += "You can put a comment like this /go issue #553 or update /set-comment new comment text"
}
// update user data
_, err := db.Exec(`
insert into users (id, first_name, last_name, user_name, last_chat_id)
values (?, ?, ?, ?, ?)
on conflict (id)
do update set
first_name = excluded.first_name,
last_name = excluded.last_name,
last_chat_id = excluded.last_chat_id,
user_name = excluded.user_name;
`, update.Message.From.ID, update.Message.From.FirstName, update.Message.From.LastName, update.Message.From.UserName, update.Message.Chat.ID)
if err != nil {
log.Println("failed to update user: ", err)
}
}
case "getcomment":
msg.Text = getLastComment(db, update.Message.From.ID, update.Message.Chat.ID)
case "setcomment":
comment := update.Message.CommandArguments()
if comment != "" {
res, err := db.Exec(`
update races
set comment = ?
where user_id = ?
and chat_id = ?
and status = 'pending'
`, update.Message.CommandArguments(), update.Message.From.ID, update.Message.Chat.ID)
if err != nil {
log.Println("failed to update user: ", err)
msg.Text = err.Error()
} else if res.RowsAffected() == 0 {
msg.Text = "No affected rows."
} else {
msg.Text = "Comment updated."
}
} else {
msg.Text = "Argument not found, use /setcomment text..."
}
case "report":
var rows []struct {
Comment string
Count int
}
_, err := db.Query(&rows, `
select comment, count(*)
from races
where user_id = ?
and chat_id = ?
and status = 'finished'
and created_at > now() - interval '1 day'
group by 1
`, update.Message.From.ID, update.Message.Chat.ID)
if err == nil {
if len(rows) > 0 {
for _, row := range rows {
msg.Text += fmt.Sprintf("\n(%d) %s", row.Count, row.Comment)
}
} else {
msg.Text = "Races not found."
}
} else {
msg.Text = err.Error()
log.Println("failed to get report", err)
}
case "cancel":
res, err := db.Exec(`update races set status = 'canceled' where status = 'pending' and user_id = ? and chat_id = ?`, update.Message.From.ID, update.Message.Chat.ID)
if err != nil {
log.Println("cancel: ", err)
msg.Text = err.Error()
} else if res.RowsAffected() > 0 {
msg.Text = "Ok, no problem."
} else {
msg.Text = "You have not active races."
}
case "status", "left":
statusReply(&msg, db, update.Message.From.ID)
case "stats":
var rows []struct {
Day string
Count int
}
_, err := db.Query(&rows, `
select created_at::date as day, count(*)
from races
where status = 'finished'
and created_at > now() - interval '7 days'
and user_id = ?
group by 1
order by 1
`, update.Message.From.ID)
if err != nil {
log.Println("cancel: ", err)
msg.Text = err.Error()
} else if len(rows) > 0 {
for index, row := range rows {
msg.Text += fmt.Sprintf("%s %d", row.Day, row.Count)
if index < len(rows)-1 {
msg.Text += "\n"
}
}
} else {
msg.Text = fmt.Sprintf("You have no completed pomodoros today. But don't be upset! There is still time to start one.")
}
default:
msg.Text = "I don't know that command. Available commands: /go, /cancel, /status, /stats"
}
if _, err := ctx.Bot.Send(msg); err != nil {
log.Println("failed to send: ", err)
}
} else { // if update.Message.Chat.ID < 0 {
// listen group
log.Println(update.Message.Text)
nickname := nicknameR.FindString(update.Message.Text)
if nickname == "" {
continue
}
var user struct {
ID int
LastChatID int64
}
_, err := ctx.DB.QueryOne(&user, `
select id, last_chat_id
from users
where user_name = ?
`, nickname[1:])
if err != nil {
log.Println("user not found", err, nickname[1:])
continue
}
var raceData struct {
PendingCount int64
}
_, err = ctx.DB.QueryOne(&raceData, `
select count(*) as pending_count
from races
where user_id = ?
and status = 'pending'
`, user.ID)
if err != nil {
log.Println("find races", err)
continue
}
if raceData.PendingCount == 0 {
msg := tgbotapi.NewMessage(user.LastChatID, update.Message.Text)
msg.ReplyToMessageID = update.Message.MessageID
if _, err := ctx.Bot.Send(msg); err != nil {
log.Println("failed to send: ", err)
}
}
}
}
}
create type race_status as enum (
'pending',
'finished',
'canceled'
);
create table races (
id serial primary key,
chat_id int not null,
user_id int not null,
initial_message_id int not null,
created_at timestamptz not null default statement_timestamp(),
duration interval not null default '25 minutes',
status race_status not null default 'pending'
);
create unique index race_pending_idx on races (user_id) where (status = 'pending');
create table users (
id bigint primary key,
first_name text,
last_name text,
user_name text
);
alter table races add column comment text;
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment