include/SqliteKV.hpp

   1 #ifndef SCX_SQLITEKV_HPP
   2 #define SCX_SQLITEKV_HPP
   3
   4 #include <sqlite3.h>
   5 #include <zlib.h>
   6 #include <string>
   7 #include <cstdio>
   8 #include <cstring>
   9
  10 namespace scx {
  11
  12 class SqliteKV
  13 {
  14 public:
  15     SqliteKV(const std::string& dbName = ""):
  16         BEGIN_TRANSACTION("BEGIN;"),
  17         COMMIT_TRANSACTION("COMMIT;"),
  18         CREATE_TABLE("create table if not exists %s(KEY text primary key, VALUE blob);"),
  19         DROP_TABLE("drop table if exists %s;"),
  20         INSERT_PAIR("insert into %s values('%s', ?);"),
  21         SELECT_VALUE("select VALUE from %s where KEY='%s';"),
  22         UPDATE_VALUE("update %s set VALUE=? where KEY='%s';"),
  23         DROP_PAIR("drop from %s where KEY='%s';"),
  24         COUNT_PAIRS("select count(key) from %s;"),
  25         COUNT_RESULTS("select count(key) from %s where KEY='%s';"),
  26         mDbName(dbName),
  27         mTableName(""),
  28         mDb(NULL),
  29         mSqlBufLen(200),
  30         mSqlBuf(new char[200])
  31     {
  32         if (!mDbName.empty())
  33             OpenDb(mDbName);
  34     }
  35
  36     ~SqliteKV()
  37     {
  38         if (!mDbName.empty())
  39             CloseDb();
  40         if (mSqlBuf != NULL)
  41             delete mSqlBuf;
  42     }
  43
  44     void OpenDb(const std::string& dbName)
  45     {
  46         mDbName = dbName;
  47         sqlite3_open(mDbName.c_str(), &mDb);
  48
  49         sqlite3_exec(mDb, "PRAGMA journal_mode=off;", NULL, 0, 0);
  50         sqlite3_exec(mDb, "PRAGMA synchronous=off;", NULL, 0, 0);
  51         sqlite3_exec(mDb, "PRAGMA cache_size=8000;", NULL, 0, 0);
  52         //sqlite3_exec(mDb, "PRAGMA page_size=2048;", NULL, 0, 0);
  53         //sqlite3_exec(mDb, "PRAGMA page_count=200;", NULL, 0, 0);
  54     }
  55
  56     void CloseDb()
  57     {
  58         mDbName.clear();
  59         sqlite3_finalize(mStmt);
  60         sqlite3_close(mDb);
  61     }
  62
  63     // switch table (if the targe table doesn't exist, I'll create it before switch)
  64     void SwitchTable(const std::string& tableName)
  65     {
  66         mTableName = tableName;
  67         CheckSqlBuf(CREATE_TABLE.size() + mTableName.size());
  68         snprintf(mSqlBuf, mSqlBufLen, CREATE_TABLE.c_str(), tableName.c_str());
  69         sqlite3_exec(mDb, mSqlBuf, NULL, 0, 0);
  70     }
  71
  72     // drop the current table
  73     void DropTable()
  74     {
  75         CheckSqlBuf(DROP_TABLE.size() + mTableName.size());
  76         snprintf(mSqlBuf, mSqlBufLen, DROP_TABLE.c_str(), mTableName.c_str());
  77         sqlite3_exec(mDb, mSqlBuf, NULL, 0, 0);
  78     }
  79
  80     void Begin()
  81     {
  82         sqlite3_exec(mDb, BEGIN_TRANSACTION.c_str(), NULL, 0, 0);
  83     }
  84
  85     void Commit()
  86     {
  87         sqlite3_exec(mDb, COMMIT_TRANSACTION.c_str(), NULL, 0, 0);
  88     }
  89
  90     /* put family */
  91
  92     // put fixed-size data, by sizeof(T)
  93     template<typename T>
  94     bool PutData(const std::string& key, const T& data)
  95     {
  96         return PutRaw(key, (char*)&data, sizeof(T));
  97     }
  98
  99     // put string data
 100     bool PutString(const std::string& key, const std::string& text)
 101     {
 102         return PutRaw(key, text.c_str(), text.size());
 103     }
 104
 105     // put raw data
 106     bool PutRaw(const std::string& key, const char* raw, int size)
 107     {
 108         CheckSqlBuf(INSERT_PAIR.size() + mTableName.size() + key.size());
 109         snprintf(mSqlBuf, mSqlBufLen, INSERT_PAIR.c_str(), mTableName.c_str(), key.c_str());
 110         sqlite3_prepare_v2(mDb, mSqlBuf, -1, &mStmt, 0);
 111         sqlite3_bind_blob(mStmt, 1, raw, size, SQLITE_STATIC);
 112         sqlite3_step(mStmt);
 113         sqlite3_finalize(mStmt);
 114         return sqlite3_changes(mDb) == 0;
 115     }
 116
 117     /* get family */
 118
 119     // get type data
 120     template<class T>
 121     bool GetData(const std::string& key, T& data)
 122     {
 123         const void* blob;
 124         int size;
 125         if (GetBlob(key, blob, size))
 126         {
 127             memcpy(&data, blob, size);
 128             return true;
 129         }
 130         else
 131         {
 132             return false;
 133         }
 134     }
 135
 136     // get string data
 137     bool GetString(const std::string& key, std::string& text)
 138     {
 139         const void* blob;
 140         int size;
 141         if (GetBlob(key, blob, size))
 142         {
 143             text.assign((const char*)blob, size);
 144             return true;
 145         }
 146         else
 147         {
 148             return false;
 149         }
 150     }
 151
 152     // get raw data, false if null or not found
 153     bool GetRaw(const std::string& key, char*& raw, int& size)
 154     {
 155         const void* blob;
 156         if (GetBlob(key, blob, size))
 157         {
 158             raw = new char[size];
 159             memcpy(raw, blob, size);
 160             return true;
 161         }
 162         else
 163         {
 164             return false;
 165         }
 166     }
 167
 168     /* update family */
 169
 170     template<typename T>
 171     bool UpdateData(const std::string& key, const T& data)
 172     {
 173         return UpdateRaw(key, (char*)&data, sizeof(T));
 174     }
 175
 176     bool UpdateString(const std::string& key, const std::string& text)
 177     {
 178         return UpdateRaw(key, text.c_str(), text.size());
 179     }
 180
 181     bool UpdateRaw(const std::string& key, const char* raw, int size)
 182     {
 183         //CheckSqlBuf(UPDATE_VALUE.size() + mTableName.size() + key.size());
 184         snprintf(mSqlBuf, mSqlBufLen, UPDATE_VALUE.c_str(), mTableName.c_str(), key.c_str());
 185         sqlite3_prepare_v2(mDb, mSqlBuf, -1, &mStmt, 0);
 186         sqlite3_bind_blob(mStmt, 1, raw, size, SQLITE_STATIC);
 187         sqlite3_step(mStmt);
 188         return true;
 189     }
 190
 191     // delete pair by key
 192     bool Delete(const std::string& key)
 193     {
 194         CheckSqlBuf(DROP_TABLE.size() + mTableName.size());
 195         snprintf(mSqlBuf, mSqlBufLen, DROP_TABLE.c_str(), mTableName.c_str());
 196         sqlite3_exec(mDb, mSqlBuf, NULL, 0, 0);
 197         return true;
 198     }
 199
 200     // number of pairs in table
 201     int Count()
 202     {
 203         CheckSqlBuf(COUNT_PAIRS.size() + mTableName.size());
 204         snprintf(mSqlBuf, mSqlBufLen, COUNT_PAIRS.c_str(), mTableName.c_str());
 205         sqlite3_prepare_v2(mDb, mSqlBuf, -1, &mStmt, 0);
 206         sqlite3_step(mStmt);
 207         return sqlite3_column_int(mStmt, 0);
 208     }
 209
 210 private:
 211     void CheckSqlBuf(size_t len)
 212     {
 213     /*
 214     if (mSqlBufLen < len)
 215     {
 216         delete mSqlBuf;
 217         mSqlBuf = new char[len];
 218         mSqlBufLen = len;
 219     }
 220     */
 221     }
 222
 223     bool GetBlob(const std::string& key, const void*& blob, int& size)
 224     {
 225         CheckSqlBuf(SELECT_VALUE.size() + mTableName.size() + key.size());
 226         snprintf(mSqlBuf, mSqlBufLen, SELECT_VALUE.c_str(), mTableName.c_str(), key.c_str());
 227         sqlite3_prepare_v2(mDb, mSqlBuf, -1, &mStmt, 0);
 228         sqlite3_step(mStmt);
 229         blob = sqlite3_column_blob(mStmt, 0);
 230         size = sqlite3_column_bytes(mStmt, 0);
 231         return (blob != NULL);
 232     }
 233
 234 private:
 235     static int OnRow(void*, int argc, char** argv, char** colName)
 236     {
 237         return 0;
 238     }
 239
 240 private:
 241     const std::string BEGIN_TRANSACTION;
 242     const std::string COMMIT_TRANSACTION;
 243     const std::string CREATE_TABLE;
 244     const std::string DROP_TABLE;
 245     const std::string INSERT_PAIR;
 246     const std::string SELECT_VALUE;
 247     const std::string UPDATE_VALUE;
 248     const std::string DROP_PAIR;
 249     const std::string COUNT_PAIRS;
 250     const std::string COUNT_RESULTS;
 251
 252 private:
 253     std::string mDbName;
 254     std::string mTableName;
 255     sqlite3* mDb;
 256     sqlite3_stmt* mStmt;
 257     size_t mSqlBufLen;
 258     char* mSqlBuf;
 259 };
 260
 261 }
 262
 263 #endif