git.sesse.net Git - nageru/blob - futatabi/db.cpp

   1 #include "db.h"
   2
   3 #include "frame.pb.h"
   4
   5 #include <string>
   6 #include <unordered_set>
   7
   8 using namespace std;
   9
  10 DB::DB(const string &filename)
  11 {
  12         int ret = sqlite3_open(filename.c_str(), &db);
  13         if (ret != SQLITE_OK) {
  14                 fprintf(stderr, "%s: %s\n", filename.c_str(), sqlite3_errmsg(db));
  15                 exit(1);
  16         }
  17
  18         sqlite3_exec(db, R"(
  19                 CREATE TABLE IF NOT EXISTS state (state BLOB);
  20         )", nullptr, nullptr, nullptr);  // Ignore errors.
  21
  22         sqlite3_exec(db, R"(
  23                 CREATE TABLE IF NOT EXISTS settings (settings BLOB);
  24         )", nullptr, nullptr, nullptr);  // Ignore errors.
  25
  26         sqlite3_exec(db, R"(
  27                 DROP TABLE file;
  28         )", nullptr, nullptr, nullptr);  // Ignore errors.
  29
  30         sqlite3_exec(db, R"(
  31                 DROP TABLE frame;
  32         )", nullptr, nullptr, nullptr);  // Ignore errors.
  33
  34         sqlite3_exec(db, R"(
  35                 CREATE TABLE IF NOT EXISTS filev2 (
  36                         file INTEGER NOT NULL PRIMARY KEY,
  37                         filename VARCHAR NOT NULL UNIQUE,
  38                         size BIGINT NOT NULL,
  39                         frames BLOB NOT NULL
  40                 );
  41         )", nullptr, nullptr, nullptr);  // Ignore errors.
  42
  43         sqlite3_exec(db, "PRAGMA journal_mode=WAL", nullptr, nullptr, nullptr);  // Ignore errors.
  44         sqlite3_exec(db, "PRAGMA synchronous=NORMAL", nullptr, nullptr, nullptr);  // Ignore errors.
  45 }
  46
  47 StateProto DB::get_state()
  48 {
  49         StateProto state;
  50
  51         sqlite3_stmt *stmt;
  52         int ret = sqlite3_prepare_v2(db, "SELECT state FROM state", -1, &stmt, 0);
  53         if (ret != SQLITE_OK) {
  54                 fprintf(stderr, "SELECT prepare: %s\n", sqlite3_errmsg(db));
  55                 exit(1);
  56         }
  57
  58         ret = sqlite3_step(stmt);
  59         if (ret == SQLITE_ROW) {
  60                 bool ok = state.ParseFromArray(sqlite3_column_blob(stmt, 0), sqlite3_column_bytes(stmt, 0));
  61                 if (!ok) {
  62                         fprintf(stderr, "State in database is corrupted!\n");
  63                         exit(1);
  64                 }
  65         } else if (ret != SQLITE_DONE) {
  66                 fprintf(stderr, "SELECT step: %s\n", sqlite3_errmsg(db));
  67                 exit(1);
  68         }
  69
  70         ret = sqlite3_finalize(stmt);
  71         if (ret != SQLITE_OK) {
  72                 fprintf(stderr, "SELECT finalize: %s\n", sqlite3_errmsg(db));
  73                 exit(1);
  74         }
  75
  76         return state;
  77 }
  78
  79 void DB::store_state(const StateProto &state)
  80 {
  81         string serialized;
  82         state.SerializeToString(&serialized);
  83
  84         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
  85         if (ret != SQLITE_OK) {
  86                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
  87                 exit(1);
  88         }
  89
  90         ret = sqlite3_exec(db, "DELETE FROM state", nullptr, nullptr, nullptr);
  91         if (ret != SQLITE_OK) {
  92                 fprintf(stderr, "DELETE: %s\n", sqlite3_errmsg(db));
  93                 exit(1);
  94         }
  95
  96         sqlite3_stmt *stmt;
  97         ret = sqlite3_prepare_v2(db, "INSERT INTO state VALUES (?)", -1, &stmt, 0);
  98         if (ret != SQLITE_OK) {
  99                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 100                 exit(1);
 101         }
 102
 103         sqlite3_bind_blob(stmt, 1, serialized.data(), serialized.size(), SQLITE_STATIC);
 104
 105         ret = sqlite3_step(stmt);
 106         if (ret == SQLITE_ROW) {
 107                 fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 108                 exit(1);
 109         }
 110
 111         ret = sqlite3_finalize(stmt);
 112         if (ret != SQLITE_OK) {
 113                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 114                 exit(1);
 115         }
 116
 117         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 118         if (ret != SQLITE_OK) {
 119                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 120                 exit(1);
 121         }
 122 }
 123
 124 SettingsProto DB::get_settings()
 125 {
 126         SettingsProto settings;
 127
 128         sqlite3_stmt *stmt;
 129         int ret = sqlite3_prepare_v2(db, "SELECT settings FROM settings", -1, &stmt, 0);
 130         if (ret != SQLITE_OK) {
 131                 fprintf(stderr, "SELECT prepare: %s\n", sqlite3_errmsg(db));
 132                 exit(1);
 133         }
 134
 135         ret = sqlite3_step(stmt);
 136         if (ret == SQLITE_ROW) {
 137                 bool ok = settings.ParseFromArray(sqlite3_column_blob(stmt, 0), sqlite3_column_bytes(stmt, 0));
 138                 if (!ok) {
 139                         fprintf(stderr, "State in database is corrupted!\n");
 140                         exit(1);
 141                 }
 142         } else if (ret != SQLITE_DONE) {
 143                 fprintf(stderr, "SELECT step: %s\n", sqlite3_errmsg(db));
 144                 exit(1);
 145         }
 146
 147         ret = sqlite3_finalize(stmt);
 148         if (ret != SQLITE_OK) {
 149                 fprintf(stderr, "SELECT finalize: %s\n", sqlite3_errmsg(db));
 150                 exit(1);
 151         }
 152
 153         return settings;
 154 }
 155
 156 void DB::store_settings(const SettingsProto &settings)
 157 {
 158         string serialized;
 159         settings.SerializeToString(&serialized);
 160
 161         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
 162         if (ret != SQLITE_OK) {
 163                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
 164                 exit(1);
 165         }
 166
 167         ret = sqlite3_exec(db, "DELETE FROM settings", nullptr, nullptr, nullptr);
 168         if (ret != SQLITE_OK) {
 169                 fprintf(stderr, "DELETE: %s\n", sqlite3_errmsg(db));
 170                 exit(1);
 171         }
 172
 173         sqlite3_stmt *stmt;
 174         ret = sqlite3_prepare_v2(db, "INSERT INTO settings VALUES (?)", -1, &stmt, 0);
 175         if (ret != SQLITE_OK) {
 176                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 177                 exit(1);
 178         }
 179
 180         sqlite3_bind_blob(stmt, 1, serialized.data(), serialized.size(), SQLITE_STATIC);
 181
 182         ret = sqlite3_step(stmt);
 183         if (ret == SQLITE_ROW) {
 184                 fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 185                 exit(1);
 186         }
 187
 188         ret = sqlite3_finalize(stmt);
 189         if (ret != SQLITE_OK) {
 190                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 191                 exit(1);
 192         }
 193
 194         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 195         if (ret != SQLITE_OK) {
 196                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 197                 exit(1);
 198         }
 199 }
 200
 201 vector<DB::FrameOnDiskAndStreamIdx> DB::load_frame_file(const string &filename, size_t size, unsigned filename_idx)
 202 {
 203         FileContentsProto file_contents;
 204
 205         sqlite3_stmt *stmt;
 206         int ret = sqlite3_prepare_v2(db, "SELECT frames FROM filev2 WHERE filename=? AND size=?", -1, &stmt, 0);
 207         if (ret != SQLITE_OK) {
 208                 fprintf(stderr, "SELECT prepare: %s\n", sqlite3_errmsg(db));
 209                 exit(1);
 210         }
 211
 212         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 213         sqlite3_bind_int64(stmt, 2, size);
 214
 215         ret = sqlite3_step(stmt);
 216         if (ret == SQLITE_ROW) {
 217                 bool ok = file_contents.ParseFromArray(sqlite3_column_blob(stmt, 0), sqlite3_column_bytes(stmt, 0));
 218                 if (!ok) {
 219                         fprintf(stderr, "Frame list in database is corrupted!\n");
 220                         exit(1);
 221                 }
 222         } else if (ret != SQLITE_DONE) {
 223                 fprintf(stderr, "SELECT step: %s\n", sqlite3_errmsg(db));
 224                 exit(1);
 225         }
 226
 227         ret = sqlite3_finalize(stmt);
 228         if (ret != SQLITE_OK) {
 229                 fprintf(stderr, "SELECT finalize: %s\n", sqlite3_errmsg(db));
 230                 exit(1);
 231         }
 232
 233         vector<FrameOnDiskAndStreamIdx> frames;
 234         for (const StreamContentsProto &stream : file_contents.stream()) {
 235                 FrameOnDiskAndStreamIdx frame;
 236                 frame.stream_idx = stream.stream_idx();
 237                 for (int i = 0; i < stream.pts_size(); ++i) {
 238                         frame.frame.filename_idx = filename_idx;
 239                         frame.frame.pts = stream.pts(i);
 240                         frame.frame.offset = stream.offset(i);
 241                         frame.frame.size = stream.file_size(i);
 242                         frames.push_back(frame);
 243                 }
 244         }
 245
 246         return frames;
 247 }
 248
 249 void DB::store_frame_file(const string &filename, size_t size, const vector<FrameOnDiskAndStreamIdx> &frames)
 250 {
 251         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
 252         if (ret != SQLITE_OK) {
 253                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
 254                 exit(1);
 255         }
 256
 257         // Delete any existing instances with this filename.
 258         sqlite3_stmt *stmt;
 259
 260         ret = sqlite3_prepare_v2(db, "DELETE FROM filev2 WHERE filename=?", -1, &stmt, 0);
 261         if (ret != SQLITE_OK) {
 262                 fprintf(stderr, "DELETE prepare: %s\n", sqlite3_errmsg(db));
 263                 exit(1);
 264         }
 265
 266         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 267
 268         ret = sqlite3_step(stmt);
 269         if (ret == SQLITE_ROW) {
 270                 fprintf(stderr, "DELETE step: %s\n", sqlite3_errmsg(db));
 271                 exit(1);
 272         }
 273
 274         ret = sqlite3_finalize(stmt);
 275         if (ret != SQLITE_OK) {
 276                 fprintf(stderr, "DELETE finalize: %s\n", sqlite3_errmsg(db));
 277                 exit(1);
 278         }
 279
 280         // Create the protobuf blob for the new row.
 281         FileContentsProto file_contents;
 282         unordered_set<unsigned> seen_stream_idx;  // Usually only one.
 283         for (const FrameOnDiskAndStreamIdx &frame : frames) {
 284                 seen_stream_idx.insert(frame.stream_idx);
 285         }
 286         for (unsigned stream_idx : seen_stream_idx) {
 287                 StreamContentsProto *stream = file_contents.add_stream();
 288                 stream->set_stream_idx(stream_idx);
 289                 stream->mutable_pts()->Reserve(frames.size());
 290                 stream->mutable_offset()->Reserve(frames.size());
 291                 stream->mutable_file_size()->Reserve(frames.size());
 292                 for (const FrameOnDiskAndStreamIdx &frame : frames) {
 293                         if (frame.stream_idx != stream_idx) {
 294                                 continue;
 295                         }
 296                         stream->add_pts(frame.frame.pts);
 297                         stream->add_offset(frame.frame.offset);
 298                         stream->add_file_size(frame.frame.size);
 299                 }
 300         }
 301         string serialized;
 302         file_contents.SerializeToString(&serialized);
 303
 304         // Insert the new row.
 305         ret = sqlite3_prepare_v2(db, "INSERT INTO filev2 (filename, size, frames) VALUES (?, ?, ?)", -1, &stmt, 0);
 306         if (ret != SQLITE_OK) {
 307                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 308                 exit(1);
 309         }
 310
 311         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 312         sqlite3_bind_int64(stmt, 2, size);
 313         sqlite3_bind_blob(stmt, 3, serialized.data(), serialized.size(), SQLITE_STATIC);
 314
 315         ret = sqlite3_step(stmt);
 316         if (ret == SQLITE_ROW) {
 317                 fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 318                 exit(1);
 319         }
 320
 321         ret = sqlite3_finalize(stmt);
 322         if (ret != SQLITE_OK) {
 323                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 324                 exit(1);
 325         }
 326
 327         // Commit.
 328         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 329         if (ret != SQLITE_OK) {
 330                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 331                 exit(1);
 332         }
 333 }
 334
 335 void DB::clean_unused_frame_files(const vector<string> &used_filenames)
 336 {
 337         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
 338         if (ret != SQLITE_OK) {
 339                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
 340                 exit(1);
 341         }
 342
 343         ret = sqlite3_exec(db, R"(
 344                 CREATE TEMPORARY TABLE used_filenames ( filename VARCHAR NOT NULL PRIMARY KEY )
 345         )", nullptr, nullptr, nullptr);
 346
 347         if (ret != SQLITE_OK) {
 348                 fprintf(stderr, "CREATE TEMPORARY TABLE: %s\n", sqlite3_errmsg(db));
 349                 exit(1);
 350         }
 351
 352         // Insert the new rows.
 353         sqlite3_stmt *stmt;
 354         ret = sqlite3_prepare_v2(db, "INSERT INTO used_filenames (filename) VALUES (?)", -1, &stmt, 0);
 355         if (ret != SQLITE_OK) {
 356                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 357                 exit(1);
 358         }
 359
 360         for (const string &filename : used_filenames) {
 361                 sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 362
 363                 ret = sqlite3_step(stmt);
 364                 if (ret == SQLITE_ROW) {
 365                         fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 366                         exit(1);
 367                 }
 368
 369                 ret = sqlite3_reset(stmt);
 370                 if (ret == SQLITE_ROW) {
 371                         fprintf(stderr, "INSERT reset: %s\n", sqlite3_errmsg(db));
 372                         exit(1);
 373                 }
 374         }
 375
 376         ret = sqlite3_finalize(stmt);
 377         if (ret != SQLITE_OK) {
 378                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 379                 exit(1);
 380         }
 381
 382         ret = sqlite3_exec(db, R"(
 383                 DELETE FROM filev2 WHERE filename NOT IN ( SELECT filename FROM used_filenames )
 384         )", nullptr, nullptr, nullptr);
 385
 386         if (ret != SQLITE_OK) {
 387                 fprintf(stderr, "DELETE: %s\n", sqlite3_errmsg(db));
 388                 exit(1);
 389         }
 390
 391         ret = sqlite3_exec(db, R"(
 392                 DROP TABLE used_filenames
 393         )", nullptr, nullptr, nullptr);
 394
 395         if (ret != SQLITE_OK) {
 396                 fprintf(stderr, "DROP TABLE: %s\n", sqlite3_errmsg(db));
 397                 exit(1);
 398         }
 399
 400         // Commit.
 401         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 402         if (ret != SQLITE_OK) {
 403                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 404                 exit(1);
 405         }
 406 }