kEzNwtUF

1package sqlite3
2
3import (
4	"fmt"
5	"log"
6
7	"database/sql"
8
9	"net/url"
10	"sync"
11
12	_ "github.com/mattn/go-sqlite3"
13
14)
15
16// Sqlite3Storage implements a SQLite3 storage backend for Colly
17type Sqlite3Storage struct {
18	// Filename indicates the name of the sqlite file to use
19	Filename string
20	// handle to the db
21	dbh *sql.DB
22	mu  sync.RWMutex // Only used for cookie methods.
23}
24
25// Init initializes the sqlite3 storage
26func (s *Sqlite3Storage) Init() error {
27
28	if s.dbh == nil {
29		db, err := sql.Open("sqlite3", s.Filename)
30		if err != nil {
31			return fmt.Errorf("unable to open db file: %s", err.Error())
32		}
33
34		err = db.Ping()
35		if err != nil {
36			return fmt.Errorf("db init failure: %s", err.Error())
37		}
38		s.dbh = db
39	}
40	// create the data structures if necessary
41	statement, _ := s.dbh.Prepare("CREATE TABLE IF NOT EXISTS visited (id INTEGER AUTO_INCREMENT PRIMARY KEY, requestID INTEGER, visited INT)")
42	_, err := statement.Exec()
43	if err != nil {
44		return err
45	}
46	statement, _ = s.dbh.Prepare("CREATE INDEX IF NOT EXISTS idx_visited ON visited (requestID)")
47	_, err = statement.Exec()
48	if err != nil {
49		return err
50	}
51	statement, _ = s.dbh.Prepare("CREATE TABLE IF NOT EXISTS cookies (id INTEGER AUTO_INCREMENT PRIMARY KEY, host TEXT, cookies TEXT)")
52	_, err = statement.Exec()
53	if err != nil {
54		return err
55	}
56	statement, err = s.dbh.Prepare("CREATE INDEX IF NOT EXISTS idx_cookies ON cookies (host)")
57	_, err = statement.Exec()
58	if err != nil {
59		return err
60	}
61	return nil
62}
63
64// Clear removes all entries from the storage
65func (s *Sqlite3Storage) Clear() error {
66	s.mu.Lock()
67	defer s.mu.Unlock()
68
69	statement, err := s.dbh.Prepare("DROP TABLE visited")
70	if err != nil {
71		return err
72	}
73	_, err = statement.Exec()
74	if err != nil {
75		return err
76	}
77	statement, err = s.dbh.Prepare("DROP TABLE cookies")
78	if err != nil {
79		return err
80	}
81	_, err = statement.Exec()
82	if err != nil {
83		return err
84	}
85	return nil
86}
87
88//Close the db
89func (s *Sqlite3Storage) Close() error {
90	err := s.dbh.Close()
91	return err
92}
93
94// Visited implements colly/storage.Visited()
95func (s *Sqlite3Storage) Visited(requestID uint64) error {
96	statement, err := s.dbh.Prepare("INSERT INTO visited (requestID, visited) VALUES (?, 1)")
97	if err != nil {
98		return err
99	}
100	_, err = statement.Exec(requestID)
101	if err != nil {
102		return err
103	}
104	return nil
105}
106
107// IsVisited implements colly/storage.IsVisited()
108func (s *Sqlite3Storage) IsVisited(requestID uint64) (bool, error) {
109	var count int
110	statement, err := s.dbh.Prepare("SELECT COUNT(*) FROM visited where requestId = ?")
111	if err != nil {
112		return false, err
113	}
114	row := statement.QueryRow(requestID)
115	err = row.Scan(&count)
116	if err != nil {
117		return false, err
118	}
119	if count == 1 {
120		return true, nil
121	}
122	return false, nil
123}
124
125// SetCookies implements colly/storage..SetCookies()
126func (s *Sqlite3Storage) SetCookies(u *url.URL, cookies string) {
127	// TODO Cookie methods currently have no way to return an error.
128
129	// We need to use a write lock to prevent a race in the db:
130	// if two callers set cookies in a very small window of time,
131	// it is possible to drop the new cookies from one caller
132	// ('last update wins' == best avoided).
133	s.mu.Lock()
134	defer s.mu.Unlock()
135
136	statement, err := s.dbh.Prepare("INSERT INTO cookies (host, cookies) VALUES (?,?)")
137	if err != nil {
138		log.Printf("SetCookies() .Set error %s", err)
139	}
140	_, err = statement.Exec(u.Host, cookies)
141	if err != nil {
142		log.Printf("SetCookies() .Set error %s", err)
143	}
144
145}
146
147// Cookies implements colly/storage.Cookies()
148func (s *Sqlite3Storage) Cookies(u *url.URL) string {
149	// TODO Cookie methods currently have no way to return an error.
150	var cookies string
151	s.mu.RLock()
152
153	//cookiesStr, err := s.Client.Get(s.getCookieID(u.Host)).Result()
154	statement, err := s.dbh.Prepare("SELECT cookies FROM cookies where host = ?")
155	if err != nil {
156		log.Printf("Cookies() .Get error %s", err)
157	}
158	row := statement.QueryRow(u.Host)
159	err = row.Scan(&cookies)
160
161	s.mu.RUnlock()
162
163	if err != nil {
164		cookies = ""
165		log.Printf("Cookies() .Get error %s", err)
166	}
167
168	return cookies
169}