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