include/caffe/util/db_leveldb.hpp

   1 #ifdef USE_LEVELDB
   2 #ifndef CAFFE_UTIL_DB_LEVELDB_HPP
   3 #define CAFFE_UTIL_DB_LEVELDB_HPP
   4
   5 #include <string>
   6
   7 #include "leveldb/db.h"
   8 #include "leveldb/write_batch.h"
   9
  10 #include "caffe/util/db.hpp"
  11
  12 namespace caffe { namespace db {
  13
  14 class LevelDBCursor : public Cursor {
  15  public:
  16   explicit LevelDBCursor(leveldb::Iterator* iter)
  17     : iter_(iter) { SeekToFirst(); }
  18   ~LevelDBCursor() { delete iter_; }
  19   virtual void SeekToFirst() { iter_->SeekToFirst(); }
  20   virtual void Next() { iter_->Next(); }
  21   virtual string key() { return iter_->key().ToString(); }
  22   virtual string value() { return iter_->value().ToString(); }
  23   virtual bool valid() { return iter_->Valid(); }
  24
  25  private:
  26   leveldb::Iterator* iter_;
  27 };
  28
  29 class LevelDBTransaction : public Transaction {
  30  public:
  31   explicit LevelDBTransaction(leveldb::DB* db) : db_(db) { CHECK_NOTNULL(db_); }
  32   virtual void Put(const string& key, const string& value) {
  33     batch_.Put(key, value);
  34   }
  35   virtual void Commit() {
  36     leveldb::Status status = db_->Write(leveldb::WriteOptions(), &batch_);
  37     CHECK(status.ok()) << "Failed to write batch to leveldb "
  38                        << std::endl << status.ToString();
  39   }
  40
  41  private:
  42   leveldb::DB* db_;
  43   leveldb::WriteBatch batch_;
  44
  45   DISABLE_COPY_AND_ASSIGN(LevelDBTransaction);
  46 };
  47
  48 class LevelDB : public DB {
  49  public:
  50   LevelDB() : db_(NULL) { }
  51   virtual ~LevelDB() { Close(); }
  52   virtual void Open(const string& source, Mode mode);
  53   virtual void Close() {
  54     if (db_ != NULL) {
  55       delete db_;
  56       db_ = NULL;
  57     }
  58   }
  59   virtual LevelDBCursor* NewCursor() {
  60     return new LevelDBCursor(db_->NewIterator(leveldb::ReadOptions()));
  61   }
  62   virtual LevelDBTransaction* NewTransaction() {
  63     return new LevelDBTransaction(db_);
  64   }
  65
  66  private:
  67   leveldb::DB* db_;
  68 };
  69
  70
  71 }  // namespace db
  72 }  // namespace caffe
  73
  74 #endif  // CAFFE_UTIL_DB_LEVELDB_HPP
  75 #endif  // USE_LEVELDB