| | #ifdef USE_LEVELDB |
| | #ifndef CAFFE_UTIL_DB_LEVELDB_HPP |
| | #define CAFFE_UTIL_DB_LEVELDB_HPP |
| |
|
| | #include <string> |
| |
|
| | #include "leveldb/db.h" |
| | #include "leveldb/write_batch.h" |
| |
|
| | #include "caffe/util/db.hpp" |
| |
|
| | namespace caffe { namespace db { |
| |
|
| | class LevelDBCursor : public Cursor { |
| | public: |
| | explicit LevelDBCursor(leveldb::Iterator* iter) |
| | : iter_(iter) { |
| | SeekToFirst(); |
| | CHECK(iter_->status().ok()) << iter_->status().ToString(); |
| | } |
| | ~LevelDBCursor() { delete iter_; } |
| | virtual void SeekToFirst() { iter_->SeekToFirst(); } |
| | virtual void Next() { iter_->Next(); } |
| | virtual string key() { return iter_->key().ToString(); } |
| | virtual string value() { return iter_->value().ToString(); } |
| | virtual bool valid() { return iter_->Valid(); } |
| |
|
| | private: |
| | leveldb::Iterator* iter_; |
| | }; |
| |
|
| | class LevelDBTransaction : public Transaction { |
| | public: |
| | explicit LevelDBTransaction(leveldb::DB* db) : db_(db) { CHECK_NOTNULL(db_); } |
| | virtual void Put(const string& key, const string& value) { |
| | batch_.Put(key, value); |
| | } |
| | virtual void Commit() { |
| | leveldb::Status status = db_->Write(leveldb::WriteOptions(), &batch_); |
| | CHECK(status.ok()) << "Failed to write batch to leveldb " |
| | << std::endl << status.ToString(); |
| | } |
| |
|
| | private: |
| | leveldb::DB* db_; |
| | leveldb::WriteBatch batch_; |
| |
|
| | DISABLE_COPY_AND_ASSIGN(LevelDBTransaction); |
| | }; |
| |
|
| | class LevelDB : public DB { |
| | public: |
| | LevelDB() : db_(NULL) { } |
| | virtual ~LevelDB() { Close(); } |
| | virtual void Open(const string& source, Mode mode); |
| | virtual void Close() { |
| | if (db_ != NULL) { |
| | delete db_; |
| | db_ = NULL; |
| | } |
| | } |
| | virtual LevelDBCursor* NewCursor() { |
| | return new LevelDBCursor(db_->NewIterator(leveldb::ReadOptions())); |
| | } |
| | virtual LevelDBTransaction* NewTransaction() { |
| | return new LevelDBTransaction(db_); |
| | } |
| |
|
| | private: |
| | leveldb::DB* db_; |
| | }; |
| |
|
| |
|
| | } |
| | } |
| |
|
| | #endif |
| | #endif |
| |
|