-
Notifications
You must be signed in to change notification settings - Fork 0
/
storage.go
123 lines (104 loc) · 2.64 KB
/
storage.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
package crawler
import (
"github.com/jinzhu/gorm"
"fmt"
"time"
_ "github.com/go-sql-driver/mysql"
)
type MysqlStorage struct {
db *gorm.DB
}
type MysqlConfig struct {
User string
Password string
Host string
Port string
DBName string
}
type GormTask struct {
*gorm.Model
Task
}
type GormRecord struct {
*gorm.Model
Record
}
func fromTask(t *Task) *GormTask {
return &GormTask{Task: *t}
}
func fromRecord(r *Record) *GormRecord {
return &GormRecord{Record: *r}
}
func NewMysqlStorage(conf *MysqlConfig) (*MysqlStorage, error) {
db, err := gorm.Open("mysql", fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?charset=utf8&parseTime=True&loc=Local", conf.User, conf.Password, conf.Host, conf.Port, conf.DBName))
if err != nil {
return nil, err
}
db.AutoMigrate(&GormTask{})
db.AutoMigrate(&GormRecord{}).AddUniqueIndex("unique_url", "url")
return &MysqlStorage{
db,
}, nil
}
func (s *MysqlStorage) GetTask(taskId int64) (*Task, error) {
var t GormTask
err := s.db.First(&t, "id = ?", taskId).Error
if err != nil {
return nil, err
}
return &t.Task, nil
}
func (s *MysqlStorage) AddTask(t *Task) (int64, error) {
gormT := fromTask(t)
err := s.db.Create(&gormT).Error
return int64(gormT.ID), err
}
func (s *MysqlStorage) RemoveTask(id int64) error {
return s.db.Delete("id = ?", id).Error
}
func (s *MysqlStorage) ListTasks(uid int64) ([]*Task, error) {
tasks := make([]*GormTask, 0)
err := s.db.Where("user_id = ?", uid).Find(&tasks).Error
result := make([]*Task, len(tasks))
for i, _ := range tasks {
result[i] = &tasks[i].Task
result[i].id = int64(tasks[i].ID)
}
return result, err
}
func (s *MysqlStorage) AddRecord(r *Record) error {
return s.db.Create(fromRecord(r)).Error
}
func (s *MysqlStorage) SetRecordChecked(url string) error {
return s.db.Model(&GormRecord{}).Where("url = ?", url).Update("checked", true).Error
}
func (s *MysqlStorage) ListUncheckedRecords(taskId int64) ([]*Record, error) {
tasks := make([]*GormRecord, 0)
err := s.db.Where("task_id = ?", taskId).Find(&tasks, "checked = ?", false).Error
results := make([]*Record, len(tasks))
for i, _ := range tasks {
results[i] = &tasks[i].Record
}
return results, err
}
func (s *MysqlStorage) GetRecord(url string, maxTimeStamp *time.Time) (*Record, error) {
var r GormRecord
if maxTimeStamp != nil {
err := s.db.First(&r, "createdAt >= ? AND url = ?", time.Time{}, url).Error
if err == gorm.ErrRecordNotFound {
return nil, nil
}
if err != nil {
return nil, err
}
return &r.Record, err
}
err := s.db.First(&r, "url = ? ", url).Error
if err == gorm.ErrRecordNotFound {
return nil, nil
}
if err != nil {
return nil, err
}
return &r.Record, err
}