IMDb SQLite: add count and has_plots methods
[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 int add_plot (string title, string plot, string? author) {
147                 int rowid;
148
149                 if (!movie_exists (title, out rowid))
150                         return 1;
151
152                 string sql;
153                 if (author != null) {
154                         sql = "INSERT INTO Plots(rowid, Plot, Author) VALUES (%d, \"%s\", \"%s\");".printf (rowid, plot.replace ("\"", "&quot;"), author.replace ("\"", "&quot;"));
155                 } else {
156                         sql = "INSERT INTO Plots(rowid, Plot) VALUES (%d, \"%s\");".printf (rowid, plot.replace ("\"", "&quot;"));
157                 }
158                 int rc;
159                 rc = db.exec (sql, callback, null);
160                 if (rc != Sqlite.OK) {
161                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
162                         stderr.printf ("offending SQL: %s\n", sql);
163                         return 1;
164                 }
165
166                 return 0;
167         }
168
169         public bool movie_exists (string title, out int rowid = null) {
170                 string sql = "SELECT rowid FROM Movies WHERE Title=\"%s\"".printf (title);
171                 Statement stmt;
172                 int rc;
173
174                 rc = db.prepare_v2 (sql, -1, out stmt);
175                 if (rc != Sqlite.OK) {
176                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
177                         return false;
178                 }
179
180                 do {
181                         rc = stmt.step ();
182                         if (rc == Sqlite.ROW) {
183                                 if (&rowid != null) {
184                                         rowid = stmt.column_int (0);
185                                 }
186                                 return true;
187                         }
188                 } while (rc == Sqlite.ROW);
189
190                 return false;
191         }
192
193         public int count (string table) {
194                 string sql = "SELECT count(*) FROM %s".printf (table);
195                 Statement stmt;
196                 int rc;
197                 int count = 0;
198
199                 rc = db.prepare_v2 (sql, -1, out stmt);
200                 if (rc != Sqlite.OK) {
201                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
202                         db.progress_handler (0, null);
203                         return 0;
204                 }
205
206                 do {
207                         rc = stmt.step ();
208                         if (rc == Sqlite.ROW) {
209                                 count = stmt.column_int (0);
210                         }
211                 } while (rc == Sqlite.ROW);
212
213                 return count;
214         }
215
216         public bool has_plots () {
217                 return (count ("Plots") > 0);
218         }
219
220         public int clear () {
221                 int rc;
222
223                 rc = db.exec (
224                         "DROP TABLE IF EXISTS Movies;" +
225                         "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);" +
226                         "DROP TABLE IF EXISTS Genres;" +
227                         "CREATE TABLE Genres (Bit INTEGER PRIMARY KEY, Genre TEXT NOT NULL);" +
228                         "DROP TABLE IF EXISTS Akas;" +
229                         "CREATE TABLE Akas (Aka TEXT NOT NULL COLLATE NOCASE, TitleID INTEGER NOT NULL);" +
230                         "DROP TABLE IF EXISTS Plots;" +
231                         "CREATE TABLE Plots (Plot TEXT NOT NULL, Author TEXT)",
232                         callback, null);
233                 if (rc != Sqlite.OK) {
234                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
235                         return 1;
236                 }
237
238                 return 0;
239         }
240
241         public int create_aka_index () {
242                 int rc;
243
244                 rc = db.exec ("CREATE INDEX AkasAka ON Akas(Aka);", callback, null);
245                 if (rc != Sqlite.OK) {
246                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
247                         return 1;
248                 }
249
250                 return 0;
251         }
252
253         public int create_votes_index () {
254                 int rc;
255
256                 rc = db.exec ("CREATE INDEX MovieVotes ON Movies(Votes);", callback, null);
257                 if (rc != Sqlite.OK) {
258                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
259                         return 1;
260                 }
261
262                 return 0;
263         }
264
265         private Cancellable? _cancellable;
266         public async int query (MovieFilter filter, ReceiveMovieFunction receive_movie, int limit, Cancellable? cancellable) {
267                 var sql = "SELECT Title, Year, Rating, Genres FROM Movies";
268                 var sep = " WHERE ";
269                 string match = null;
270                 Statement stmt;
271                 int rc;
272
273                 // FIXME - how many opcodes until main loop iteration for best responsivity?
274                 _cancellable = cancellable;
275                 db.progress_handler (1000, progress_handler);
276
277                 if (filter.title != null && filter.title != "") {
278                         if ("*" in filter.title) {
279                                 match = "GLOB \"%s (*)\"".printf (filter.title);
280                         } else {
281                                 match = "LIKE \"%s%%\"".printf (filter.title);
282                         }
283                         sql += sep + "(Title %s OR rowid IN (SELECT TitleID FROM Akas WHERE Aka %s))".printf (match, match);
284                         sep = " AND ";
285                 }
286                 if (filter.year_min > 0) {
287                         sql += sep + "Year >= %d".printf (filter.year_min);
288                         sep = " AND ";
289                 }
290                 if (filter.year_max > 0) {
291                         sql += sep + "Year <= %d".printf (filter.year_max);
292                         sep = " AND ";
293                 }
294                 if (filter.rating_min > 0) {
295                         sql += sep + "Rating >= %d".printf (filter.rating_min);
296                         sep = " AND ";
297                 }
298                 if (filter.genres.field != 0) {
299                         sql += sep + "Genres&%d = %d".printf (filter.genres.field, filter.genres.field);
300                 }
301                 sql += " ORDER BY Votes DESC LIMIT %d;".printf (limit + 1);
302
303                 stdout.printf("SQL: \"%s\"\n", sql);
304
305                 rc = db.prepare_v2 (sql, -1, out stmt);
306                 if (rc != Sqlite.OK) {
307                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
308                         db.progress_handler (0, null);
309                         return 0;
310                 }
311
312                 int n = 0;
313                 do {
314                         if (++n > limit)
315                                 break;
316                         Idle.add (query.callback);
317                         yield;
318                         rc = stmt.step ();
319                         if (rc == Sqlite.ROW) {
320                                 int year = stmt.column_int (1);
321                                 string title = stmt.column_text (0);
322                                 int rating = stmt.column_int (2);
323                                 int genres = stmt.column_int (3);
324                                 string aka = null;
325                                 if (match != null && !(filter.matches_title (strip_year (title, year)))) {
326                                         aka = movie_aka (title, match);
327                                         if (aka != null)
328                                                 aka = strip_year (aka, year);
329                                 }
330                                 receive_movie (strip_year (title, year), aka, year, rating, genres);
331                         }
332                 } while (rc == Sqlite.ROW);
333
334                 db.progress_handler (0, null);
335                 return n;
336         }
337
338         private string movie_aka (string title, string match) {
339                 string sql = "SELECT Aka FROM Akas WHERE (TitleID = (SELECT rowid FROM Movies WHERE Title = \"%s\") AND Aka %s) LIMIT 1;".printf (title, match);
340                 Statement stmt;
341                 int rc;
342                 string aka = null;
343
344                 rc = db.prepare_v2 (sql, -1, out stmt);
345                 if (rc != Sqlite.OK) {
346                         stderr.printf ("SQL error: %d, %s\n", rc, db.errmsg ());
347                         return null;
348                 }
349
350                 do {
351                         rc = stmt.step ();
352                         if (rc == Sqlite.ROW) {
353                                 aka = stmt.column_text (0);
354                         }
355                 } while (rc == Sqlite.ROW);
356
357                 return aka;
358         }
359
360         private int progress_handler () {
361                 ((MainContext) null).iteration (false);
362                 return (int) _cancellable.is_cancelled ();
363         }
364
365         private string strip_year (string title, int year) {
366                 string year_suffix = " (%d)".printf (year);
367                 if (title.has_suffix (year_suffix))
368                         return title.substring (0, title.length - year_suffix.length);
369                 year_suffix = " (%d/I)".printf (year);
370                 if (title.has_suffix (year_suffix))
371                         return title.substring (0, title.length - year_suffix.length);
372                 year_suffix = " (%d/II)".printf (year);
373                 if (title.has_suffix (year_suffix))
374                         return title.substring (0, title.length - year_suffix.length);
375                 return title.dup ();
376         }
377 }