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

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