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