IMDb plugin + downloader: parse IMDb alternative titles and use to match
[cinaest] / src / imdb / imdb-sqlite.vala
1 /* This file is part of Cinaest.
2  *
3  * Copyright (C) 2009 Philipp Zabel
4  *
5  * Cinaest is free software: you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation, either version 3 of the License, or
8  * (at your option) any later version.
9  *
10  * Cinaest is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with Cinaest. If not, see <http://www.gnu.org/licenses/>.
17  */
18
19 using Sqlite;
20
21 class IMDbSqlite : Object {
22         Database db;
23         List<string> genres;
24
25         public delegate void ReceiveMovieFunction (string title, string? aka, int year, int rating, int genres);
26
27         public IMDbSqlite (string filename) {
28                 int rc;
29
30                 genres = new List<string> ();
31
32                 rc = Database.open (filename, out db);
33                 if (rc != Sqlite.OK) {
34                         stderr.printf ("Can't open database: %d, %s\n", rc, db.errmsg ());
35                         return;
36                 }
37
38                 rc = db.exec ("PRAGMA journal_mode = OFF;", callback, null);
39                 if (rc != Sqlite.OK) {
40                         stderr.printf ("Can't turn off journal mode: %d, %s\n", rc, db.errmsg ());
41                         return;
42                 }
43
44                 rc = db.exec ("PRAGMA locking_mode = EXCLUSIVE;", callback, null);
45                 if (rc != Sqlite.OK) {
46                         stderr.printf ("Can't get exclusive lock: %d, %s\n", rc, db.errmsg ());
47                         return;
48                 }
49
50                 rc = db.exec ("PRAGMA synchronous = OFF;", callback, null);
51                 if (rc != Sqlite.OK)
52                         stderr.printf ("Can't turn off synchronous access: %d, %s\n", rc, db.errmsg ());
53
54         }
55
56         public static int callback (int n_columns, string[] values,
57                                     string[] column_names) {
58                 for (int i = 0; i < n_columns; i++) {
59                         stdout.printf ("%s = %s\n", column_names[i], values[i]);
60                 }
61                 stdout.printf ("\n");
62
63                 return 0;
64         }
65
66         public int add_movie (string title, int year) {
67                 string sql = "INSERT INTO Movies(Title, Year) VALUES (\"%s\", %d);".printf (title, year);
68                 int rc;
69
70                 rc = db.exec (sql, callback, null);
71                 if (rc != Sqlite.OK) {
72                         stderr.printf ("Failed to insert movie \"%s\" (%d): %d, %s\n", title, year, rc, db.errmsg ());
73                         return 1;
74                 }
75
76                 return 0;
77         }
78
79         public int movie_set_rating (string title, int rating, int votes) {
80                 var sql = "UPDATE Movies SET Rating=%d, Votes=%d WHERE Title=\"%s\";".printf (rating, votes, title);
81                 int rc;
82
83                 rc = db.exec (sql, callback, null);
84                 if (rc != Sqlite.OK) {
85                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
86                         return 1;
87                 }
88
89                 return 0;
90         }
91
92         public int movie_add_genre (string title, string genre) {
93                 string sql;
94                 int bit;
95                 int rc;
96
97                 bit = genre_bit (genre);
98                 if (bit == 0) {
99                         genres.append (genre);
100                         bit = genre_bit (genre);
101
102                         sql = "INSERT INTO Genres(Bit, Genre) VALUES (%d, \"%s\");".printf (bit, genre);
103
104                         rc = db.exec (sql, callback, null);
105                         if (rc != Sqlite.OK) {
106                                 stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
107                                 return 1;
108                         }
109                 }
110
111                 sql = "UPDATE Movies SET Genres=Genres|%d WHERE Title=\"%s\";".printf (bit, title);
112                 rc = db.exec (sql, callback, null);
113                 if (rc != Sqlite.OK) {
114                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
115                         return 1;
116                 }
117
118                 return 0;
119         }
120
121         int genre_bit (string genre) {
122                 for (int i = 0; i < genres.length (); i++) {
123                         if (genres.nth_data (i) == genre)
124                                 return 1 << i;
125                 }
126                 return 0;
127         }
128
129         public int add_aka (string title, string aka) {
130                 int rowid;
131
132                 if (!movie_exists (title, out rowid))
133                         return 1;
134
135                 string sql = "INSERT INTO Akas(Aka, TitleID) VALUES (\"%s\", %d);".printf (aka, rowid);
136                 int rc;
137                 rc = db.exec (sql, callback, null);
138                 if (rc != Sqlite.OK) {
139                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
140                         return 1;
141                 }
142
143                 return 0;
144         }
145
146         public bool movie_exists (string title, out int rowid = null) {
147                 string sql = "SELECT rowid FROM Movies WHERE Title=\"%s\"".printf (title);
148                 Statement stmt;
149                 int rc;
150                 int count = 0;
151
152                 rc = db.prepare_v2 (sql, -1, out stmt);
153                 if (rc != Sqlite.OK) {
154                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
155                         return false;
156                 }
157
158                 do {
159                         rc = stmt.step ();
160                         if (rc == Sqlite.ROW) {
161                                 if (&rowid != null) {
162                                         rowid = stmt.column_int (0);
163                                 }
164                                 return true;
165                         }
166                 } while (rc == Sqlite.ROW);
167
168                 return false;
169         }
170
171         public int clear () {
172                 int rc;
173
174                 rc = db.exec (
175                         "DROP TABLE IF EXISTS Movies;" +
176                         "CREATE TABLE Movies (Title TEXT PRIMARY KEY COLLATE NOCASE, Year INTEGER, Rating INTEGER, Votes INTEGER NOT NULL DEFAULT 0, Genres INTEGER NOT NULL DEFAULT 0);" +
177                         "DROP TABLE IF EXISTS Genres;" +
178                         "CREATE TABLE Genres (Bit INTEGER PRIMARY KEY, Genre TEXT NOT NULL);" +
179                         "DROP TABLE IF EXISTS Akas;" +
180                         "CREATE TABLE Akas (Aka TEXT NOT NULL COLLATE NOCASE, TitleID INTEGER NOT NULL);",
181                         callback, null);
182                 if (rc != Sqlite.OK) {
183                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
184                         return 1;
185                 }
186
187                 return 0;
188         }
189
190         public int create_aka_index () {
191                 int rc;
192
193                 rc = db.exec ("CREATE INDEX AkasAka ON Akas(Aka);", callback, null);
194                 if (rc != Sqlite.OK) {
195                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
196                         return 1;
197                 }
198
199                 return 0;
200         }
201
202         public int create_votes_index () {
203                 int rc;
204
205                 rc = db.exec ("CREATE INDEX MovieVotes ON Movies(Votes);", callback, null);
206                 if (rc != Sqlite.OK) {
207                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
208                         return 1;
209                 }
210
211                 return 0;
212         }
213
214         private Cancellable? _cancellable;
215         public async int query (MovieFilter filter, ReceiveMovieFunction receive_movie, int limit, Cancellable? cancellable) {
216                 var sql = "SELECT Title, Year, Rating, Genres FROM Movies";
217                 var sep = " WHERE ";
218                 string match = null;
219                 Statement stmt;
220                 int rc;
221
222                 // FIXME - how many opcodes until main loop iteration for best responsivity?
223                 _cancellable = cancellable;
224                 db.progress_handler (1000, progress_handler);
225
226                 if (filter.title != null && filter.title != "") {
227                         if ("*" in filter.title) {
228                                 match = "GLOB \"%s (*)\"".printf (filter.title);
229                         } else {
230                                 match = "LIKE \"%s%%\"".printf (filter.title);
231                         }
232                         sql += sep + "(Title %s OR rowid IN (SELECT TitleID FROM Akas WHERE Aka %s))".printf (match, match);
233                         sep = " AND ";
234                 }
235                 if (filter.year_min > 0) {
236                         sql += sep + "Year >= %d".printf (filter.year_min);
237                         sep = " AND ";
238                 }
239                 if (filter.year_max > 0) {
240                         sql += sep + "Year <= %d".printf (filter.year_max);
241                         sep = " AND ";
242                 }
243                 if (filter.rating_min > 0) {
244                         sql += sep + "Rating >= %d".printf (filter.rating_min);
245                         sep = " AND ";
246                 }
247                 if (filter.genres.field != 0) {
248                         sql += sep + "Genres&%d = %d".printf (filter.genres.field, filter.genres.field);
249                 }
250                 sql += " ORDER BY Votes DESC LIMIT %d;".printf (limit);
251
252                 stdout.printf("SQL: \"%s\"\n", sql);
253
254                 rc = db.prepare_v2 (sql, -1, out stmt);
255                 if (rc != Sqlite.OK) {
256                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
257                         db.progress_handler (0, null);
258                         return 1;
259                 }
260
261                 do {
262                         Idle.add (query.callback);
263                         yield;
264                         rc = stmt.step ();
265                         if (rc == Sqlite.ROW) {
266                                 int year = stmt.column_int (1);
267                                 string title = stmt.column_text (0);
268                                 int rating = stmt.column_int (2);
269                                 int genres = stmt.column_int (3);
270                                 string aka = null;
271                                 if (match != null && !(filter.matches_title (strip_year (title, year)))) {
272                                         aka = movie_aka (title, match);
273                                         if (aka != null)
274                                                 aka = strip_year (aka, year);
275                                 }
276                                 receive_movie (strip_year (title, year), aka, year, rating, genres);
277                         }
278                 } while (rc == Sqlite.ROW);
279
280                 db.progress_handler (0, null);
281                 return 0;
282         }
283
284         private string movie_aka (string title, string match) {
285                 string sql = "SELECT Aka FROM Akas WHERE (TitleID = (SELECT rowid FROM Movies WHERE Title = \"%s\") AND Aka %s) LIMIT 1;".printf (title, match);
286                 Statement stmt;
287                 int rc;
288                 string aka = null;
289
290                 rc = db.prepare_v2 (sql, -1, out stmt);
291                 if (rc != Sqlite.OK) {
292                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
293                         return null;
294                 }
295
296                 do {
297                         rc = stmt.step ();
298                         if (rc == Sqlite.ROW) {
299                                 aka = stmt.column_text (0);
300                         }
301                 } while (rc == Sqlite.ROW);
302
303                 return aka;
304         }
305
306         private int progress_handler () {
307                 ((MainContext) null).iteration (false);
308                 return (int) _cancellable.is_cancelled ();
309         }
310
311         private string strip_year (string title, int year) {
312                 string year_suffix = " (%d)".printf (year);
313                 if (title.has_suffix (year_suffix))
314                         return title.substring (0, title.length - year_suffix.length);
315                 year_suffix = " (%d/I)".printf (year);
316                 if (title.has_suffix (year_suffix))
317                         return title.substring (0, title.length - year_suffix.length);
318                 year_suffix = " (%d/II)".printf (year);
319                 if (title.has_suffix (year_suffix))
320                         return title.substring (0, title.length - year_suffix.length);
321                 return title.dup ();
322         }
323 }