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

   1 #include "db.h"
   2
   3 #include <string>
   4
   5 using namespace std;
   6
   7 DB::DB(const string &filename)
   8 {
   9         int ret = sqlite3_open(filename.c_str(), &db);
  10         if (ret != SQLITE_OK) {
  11                 fprintf(stderr, "%s: %s\n", filename.c_str(), sqlite3_errmsg(db));
  12                 exit(1);
  13         }
  14
  15         sqlite3_exec(db, R"(
  16                 CREATE TABLE IF NOT EXISTS state (state BLOB);
  17         )", nullptr, nullptr, nullptr);  // Ignore errors.
  18
  19         sqlite3_exec(db, R"(
  20                 CREATE TABLE IF NOT EXISTS file (
  21                         file INTEGER NOT NULL PRIMARY KEY,
  22                         filename VARCHAR NOT NULL UNIQUE,
  23                         size BIGINT NOT NULL
  24                 );
  25         )", nullptr, nullptr, nullptr);  // Ignore errors.
  26
  27         sqlite3_exec(db, R"(
  28                 CREATE TABLE IF NOT EXISTS frame (
  29                         file INTEGER NOT NULL REFERENCES file ON DELETE CASCADE,
  30                         stream_idx INTEGER NOT NULL,
  31                         pts BIGINT NOT NULL,
  32                         offset BIGINT NOT NULL,
  33                         size INTEGER NOT NULL
  34                 );
  35         )", nullptr, nullptr, nullptr);  // Ignore errors.
  36
  37         sqlite3_exec(db, "CREATE INDEX frame_file ON FRAME ( file );", nullptr, nullptr, nullptr);  // Ignore errors.
  38
  39         sqlite3_exec(db, "PRAGMA journal_mode=WAL", nullptr, nullptr, nullptr);  // Ignore errors.
  40         sqlite3_exec(db, "PRAGMA synchronous=NORMAL", nullptr, nullptr, nullptr);  // Ignore errors.
  41 }
  42
  43 StateProto DB::get_state()
  44 {
  45         StateProto state;
  46
  47         sqlite3_stmt *stmt;
  48         int ret = sqlite3_prepare_v2(db, "SELECT state FROM state", -1, &stmt, 0);
  49         if (ret != SQLITE_OK) {
  50                 fprintf(stderr, "SELECT prepare: %s\n", sqlite3_errmsg(db));
  51                 exit(1);
  52         }
  53
  54         ret = sqlite3_step(stmt);
  55         if (ret == SQLITE_ROW) {
  56                 bool ok = state.ParseFromArray(sqlite3_column_blob(stmt, 0), sqlite3_column_bytes(stmt, 0));
  57                 if (!ok) {
  58                         fprintf(stderr, "State in database is corrupted!\n");
  59                         exit(1);
  60                 }
  61         } else if (ret != SQLITE_DONE) {
  62                 fprintf(stderr, "SELECT step: %s\n", sqlite3_errmsg(db));
  63                 exit(1);
  64         }
  65
  66         ret = sqlite3_finalize(stmt);
  67         if (ret != SQLITE_OK) {
  68                 fprintf(stderr, "SELECT finalize: %s\n", sqlite3_errmsg(db));
  69                 exit(1);
  70         }
  71
  72         return state;
  73 }
  74
  75 void DB::store_state(const StateProto &state)
  76 {
  77         string serialized;
  78         state.SerializeToString(&serialized);
  79
  80         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
  81         if (ret != SQLITE_OK) {
  82                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
  83                 exit(1);
  84         }
  85
  86         ret = sqlite3_exec(db, "DELETE FROM state", nullptr, nullptr, nullptr);
  87         if (ret != SQLITE_OK) {
  88                 fprintf(stderr, "DELETE: %s\n", sqlite3_errmsg(db));
  89                 exit(1);
  90         }
  91
  92         sqlite3_stmt *stmt;
  93         ret = sqlite3_prepare_v2(db, "INSERT INTO state VALUES (?)", -1, &stmt, 0);
  94         if (ret != SQLITE_OK) {
  95                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
  96                 exit(1);
  97         }
  98
  99         sqlite3_bind_blob(stmt, 1, serialized.data(), serialized.size(), SQLITE_STATIC);
 100
 101         ret = sqlite3_step(stmt);
 102         if (ret == SQLITE_ROW) {
 103                 fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 104                 exit(1);
 105         }
 106
 107         ret = sqlite3_finalize(stmt);
 108         if (ret != SQLITE_OK) {
 109                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 110                 exit(1);
 111         }
 112
 113         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 114         if (ret != SQLITE_OK) {
 115                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 116                 exit(1);
 117         }
 118 }
 119
 120 vector<DB::FrameOnDiskAndStreamIdx> DB::load_frame_file(const string &filename, size_t size, unsigned filename_idx)
 121 {
 122         sqlite3_stmt *stmt;
 123         int ret = sqlite3_prepare_v2(db, "SELECT pts, offset, frame.size, stream_idx FROM file JOIN frame USING (file) WHERE filename=? AND file.size=?", -1, &stmt, 0);
 124         if (ret != SQLITE_OK) {
 125                 fprintf(stderr, "SELECT prepare: %s\n", sqlite3_errmsg(db));
 126                 exit(1);
 127         }
 128
 129         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 130         sqlite3_bind_int64(stmt, 2, size);
 131
 132         vector<FrameOnDiskAndStreamIdx> frames;
 133         do {
 134                 ret = sqlite3_step(stmt);
 135                 if (ret == SQLITE_ROW) {
 136                         FrameOnDiskAndStreamIdx frame;
 137                         frame.frame.filename_idx = filename_idx;
 138                         frame.frame.pts = sqlite3_column_int64(stmt, 0);
 139                         frame.frame.offset = sqlite3_column_int64(stmt, 1);
 140                         frame.frame.size = sqlite3_column_int(stmt, 2);
 141                         frame.stream_idx = sqlite3_column_int(stmt, 3);
 142                         frames.push_back(frame);
 143                 } else if (ret != SQLITE_DONE) {
 144                         fprintf(stderr, "SELECT step: %s\n", sqlite3_errmsg(db));
 145                         exit(1);
 146                 }
 147         } while (ret != SQLITE_DONE);
 148
 149         ret = sqlite3_finalize(stmt);
 150         if (ret != SQLITE_OK) {
 151                 fprintf(stderr, "SELECT finalize: %s\n", sqlite3_errmsg(db));
 152                 exit(1);
 153         }
 154
 155         return frames;
 156 }
 157
 158 void DB::store_frame_file(const string &filename, size_t size, const vector<FrameOnDiskAndStreamIdx> &frames)
 159 {
 160         int ret = sqlite3_exec(db, "BEGIN", nullptr, nullptr, nullptr);
 161         if (ret != SQLITE_OK) {
 162                 fprintf(stderr, "BEGIN: %s\n", sqlite3_errmsg(db));
 163                 exit(1);
 164         }
 165
 166         // Delete any existing instances with this filename. This also includes
 167         // deleting any associated frames, due to the ON CASCADE DELETE constraint.
 168         sqlite3_stmt *stmt;
 169         ret = sqlite3_prepare_v2(db, "DELETE FROM file WHERE filename=?", -1, &stmt, 0);
 170         if (ret != SQLITE_OK) {
 171                 fprintf(stderr, "DELETE prepare: %s\n", sqlite3_errmsg(db));
 172                 exit(1);
 173         }
 174
 175         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 176
 177         ret = sqlite3_step(stmt);
 178         if (ret == SQLITE_ROW) {
 179                 fprintf(stderr, "DELETE step: %s\n", sqlite3_errmsg(db));
 180                 exit(1);
 181         }
 182
 183         ret = sqlite3_finalize(stmt);
 184         if (ret != SQLITE_OK) {
 185                 fprintf(stderr, "DELETE finalize: %s\n", sqlite3_errmsg(db));
 186                 exit(1);
 187         }
 188
 189         // Insert the new row.
 190         ret = sqlite3_prepare_v2(db, "INSERT INTO file (filename, size) VALUES (?, ?)", -1, &stmt, 0);
 191         if (ret != SQLITE_OK) {
 192                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 193                 exit(1);
 194         }
 195
 196         sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), SQLITE_STATIC);
 197         sqlite3_bind_int64(stmt, 2, size);
 198
 199         ret = sqlite3_step(stmt);
 200         if (ret == SQLITE_ROW) {
 201                 fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 202                 exit(1);
 203         }
 204
 205         ret = sqlite3_finalize(stmt);
 206         if (ret != SQLITE_OK) {
 207                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 208                 exit(1);
 209         }
 210
 211         // Insert the actual frames.
 212         int64_t rowid = sqlite3_last_insert_rowid(db);
 213
 214         ret = sqlite3_prepare_v2(db, "INSERT INTO frame (file, stream_idx, pts, offset, size) VALUES (?, ?, ?, ?, ?)", -1, &stmt, 0);
 215         if (ret != SQLITE_OK) {
 216                 fprintf(stderr, "INSERT prepare: %s\n", sqlite3_errmsg(db));
 217                 exit(1);
 218         }
 219
 220         sqlite3_bind_int64(stmt, 1, rowid);
 221
 222         for (const FrameOnDiskAndStreamIdx &frame : frames) {
 223                 sqlite3_bind_int64(stmt, 2, frame.stream_idx);
 224                 sqlite3_bind_int64(stmt, 3, frame.frame.pts);
 225                 sqlite3_bind_int64(stmt, 4, frame.frame.offset);
 226                 sqlite3_bind_int(stmt, 5, frame.frame.size);
 227
 228                 ret = sqlite3_step(stmt);
 229                 if (ret == SQLITE_ROW) {
 230                         fprintf(stderr, "INSERT step: %s\n", sqlite3_errmsg(db));
 231                         exit(1);
 232                 }
 233
 234                 ret = sqlite3_reset(stmt);
 235                 if (ret != SQLITE_OK) {
 236                         fprintf(stderr, "INSERT reset: %s\n", sqlite3_errmsg(db));
 237                         exit(1);
 238                 }
 239         }
 240
 241         ret = sqlite3_finalize(stmt);
 242         if (ret != SQLITE_OK) {
 243                 fprintf(stderr, "INSERT finalize: %s\n", sqlite3_errmsg(db));
 244                 exit(1);
 245         }
 246
 247         // Commit.
 248         ret = sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
 249         if (ret != SQLITE_OK) {
 250                 fprintf(stderr, "COMMIT: %s\n", sqlite3_errmsg(db));
 251                 exit(1);
 252         }
 253 }