add cf argument for SstFileWriter constructor (#22)

librocksdb_sys/crocksdb/c.cc

@@ -2458,19 +2458,21 @@ crocksdb_envoptions_t* crocksdb_envoptions_create() {
 void crocksdb_envoptions_destroy(crocksdb_envoptions_t* opt) { delete opt; }
 
 crocksdb_sstfilewriter_t* crocksdb_sstfilewriter_create(
-    const crocksdb_envoptions_t* env, const crocksdb_options_t* io_options) {
+    const crocksdb_envoptions_t* env, const crocksdb_options_t* io_options,
+    crocksdb_column_family_handle_t* column_family) {
   crocksdb_sstfilewriter_t* writer = new crocksdb_sstfilewriter_t;
   writer->rep =
-      new SstFileWriter(env->rep, io_options->rep, io_options->rep.comparator);
+      new SstFileWriter(env->rep, io_options->rep, io_options->rep.comparator, column_family->rep);
   return writer;
 }
 
 crocksdb_sstfilewriter_t* crocksdb_sstfilewriter_create_with_comparator(
     const crocksdb_envoptions_t* env, const crocksdb_options_t* io_options,
-    const crocksdb_comparator_t* comparator) {
+    const crocksdb_comparator_t* comparator,
+    crocksdb_column_family_handle_t* column_family) {
   crocksdb_sstfilewriter_t* writer = new crocksdb_sstfilewriter_t;
   writer->rep =
-      new SstFileWriter(env->rep, io_options->rep, io_options->rep.comparator);
+      new SstFileWriter(env->rep, io_options->rep, comparator, column_family->rep);
   return writer;
 }
 

librocksdb_sys/crocksdb/rocksdb/c.h

@@ -986,11 +986,13 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_envoptions_destroy(
 
 extern C_ROCKSDB_LIBRARY_API crocksdb_sstfilewriter_t*
 crocksdb_sstfilewriter_create(const crocksdb_envoptions_t* env,
-                             const crocksdb_options_t* io_options);
+                             const crocksdb_options_t* io_options,
+                             crocksdb_column_family_handle_t* column_family);
 extern C_ROCKSDB_LIBRARY_API crocksdb_sstfilewriter_t*
 crocksdb_sstfilewriter_create_with_comparator(
     const crocksdb_envoptions_t* env, const crocksdb_options_t* io_options,
-    const crocksdb_comparator_t* comparator);
+    const crocksdb_comparator_t* comparator,
+    crocksdb_column_family_handle_t* column_family);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_sstfilewriter_open(
     crocksdb_sstfilewriter_t* writer, const char* name, char** errptr);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_sstfilewriter_add(

librocksdb_sys/src/lib.rs

@@ -629,11 +629,13 @@ extern "C" {
 
     // SstFileWriter
     pub fn crocksdb_sstfilewriter_create(env: *mut EnvOptions,
-                                         io_options: *const DBOptions)
+                                         io_options: *const DBOptions,
+                                         cf: *mut DBCFHandle)
                                          -> *mut SstFileWriter;
     pub fn crocksdb_sstfilewriter_create_with_comparator(env: *mut EnvOptions,
                                                          io_options: *const DBOptions,
-                                                         comparator: *const DBComparator)
+                                                         comparator: *const DBComparator,
+                                                         cf: *mut DBCFHandle)
                                                          -> *mut SstFileWriter;
     pub fn crocksdb_sstfilewriter_open(writer: *mut SstFileWriter,
                                        name: *const c_char,

src/external_file.rs

-// Copyright 2017 PingCAP, Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use librocksdb_sys;
-use rocksdb_options::{Options, EnvOptions};
-
-use std::ffi::CString;
-
-/// SstFileWriter is used to create sst files that can be added to database later
-/// All keys in files generated by SstFileWriter will have sequence number = 0
-pub struct SstFileWriter {
-    inner: *mut librocksdb_sys::SstFileWriter,
-}
-
-unsafe impl Send for SstFileWriter {}
-
-impl SstFileWriter {
-    pub fn new(env_opt: &EnvOptions, opt: &Options) -> SstFileWriter {
-        unsafe {
-            SstFileWriter {
-                inner: librocksdb_sys::crocksdb_sstfilewriter_create(env_opt.inner, opt.inner),
-            }
-        }
-    }
-
-    /// Prepare SstFileWriter to write into file located at "file_path".
-    pub fn open(&mut self, name: &str) -> Result<(), String> {
-        let path = match CString::new(name.to_owned()) {
-            Err(e) => return Err(format!("invalid path {}: {:?}", name, e)),
-            Ok(p) => p,
-        };
-        unsafe { Ok(ffi_try!(crocksdb_sstfilewriter_open(self.inner, path.as_ptr()))) }
-    }
-
-    /// Add key, value to currently opened file
-    /// REQUIRES: key is after any previously added key according to comparator.
-    pub fn add(&mut self, key: &[u8], val: &[u8]) -> Result<(), String> {
-        unsafe {
-            ffi_try!(crocksdb_sstfilewriter_add(self.inner,
-                                                key.as_ptr(),
-                                                key.len(),
-                                                val.as_ptr(),
-                                                val.len()));
-            Ok(())
-        }
-    }
-
-    /// Finalize writing to sst file and close file.
-    pub fn finish(&mut self) -> Result<(), String> {
-        unsafe {
-            ffi_try!(crocksdb_sstfilewriter_finish(self.inner));
-            Ok(())
-        }
-    }
-}
-
-impl Drop for SstFileWriter {
-    fn drop(&mut self) {
-        unsafe { librocksdb_sys::crocksdb_sstfilewriter_destroy(self.inner) }
-    }
-}

src/lib.rs

@@ -25,17 +25,15 @@ pub mod rocksdb_options;
 pub mod merge_operator;
 pub mod comparator;
 mod compaction_filter;
-mod external_file;
 mod slice_transform;
 
 pub use compaction_filter::CompactionFilter;
-pub use external_file::SstFileWriter;
 pub use librocksdb_sys::{DBCompactionStyle, DBCompressionType, DBRecoveryMode,
                          DBStatisticsTickerType, DBStatisticsHistogramType, new_bloom_filter,
                          self as crocksdb_ffi};
 pub use merge_operator::MergeOperands;
 pub use rocksdb::{DB, DBIterator, DBVector, Kv, SeekKey, Writable, WriteBatch, CFHandle, Range,
-                  BackupEngine};
+                  BackupEngine, SstFileWriter};
 pub use rocksdb_options::{BlockBasedOptions, Options, ReadOptions, WriteOptions, RestoreOptions,
                           IngestExternalFileOptions, EnvOptions, HistogramData};
 pub use slice_transform::SliceTransform;

src/rocksdb.rs

@@ -16,7 +16,7 @@
 use crocksdb_ffi::{self, DBWriteBatch, DBCFHandle, DBInstance, DBBackupEngine,
                    DBStatisticsTickerType, DBStatisticsHistogramType};
 use libc::{self, c_int, c_void, size_t};
-use rocksdb_options::{Options, ReadOptions, UnsafeSnap, WriteOptions, FlushOptions,
+use rocksdb_options::{Options, ReadOptions, UnsafeSnap, WriteOptions, FlushOptions, EnvOptions,
                       RestoreOptions, IngestExternalFileOptions, HistogramData};
 use std::{fs, ptr, slice};
 use std::collections::BTreeMap;
@@ -1301,6 +1301,61 @@ impl Drop for BackupEngine {
     }
 }
 
+/// SstFileWriter is used to create sst files that can be added to database later
+/// All keys in files generated by SstFileWriter will have sequence number = 0
+pub struct SstFileWriter {
+    inner: *mut crocksdb_ffi::SstFileWriter,
+}
+
+unsafe impl Send for SstFileWriter {}
+
+impl SstFileWriter {
+    pub fn new(env_opt: &EnvOptions, opt: &Options, cf: &CFHandle) -> SstFileWriter {
+        unsafe {
+            SstFileWriter {
+                inner: crocksdb_ffi::crocksdb_sstfilewriter_create(env_opt.inner,
+                                                                   opt.inner,
+                                                                   cf.inner),
+            }
+        }
+    }
+
+    /// Prepare SstFileWriter to write into file located at "file_path".
+    pub fn open(&mut self, name: &str) -> Result<(), String> {
+        let path = match CString::new(name.to_owned()) {
+            Err(e) => return Err(format!("invalid path {}: {:?}", name, e)),
+            Ok(p) => p,
+        };
+        unsafe { Ok(ffi_try!(crocksdb_sstfilewriter_open(self.inner, path.as_ptr()))) }
+    }
+
+    /// Add key, value to currently opened file
+    /// REQUIRES: key is after any previously added key according to comparator.
+    pub fn add(&mut self, key: &[u8], val: &[u8]) -> Result<(), String> {
+        unsafe {
+            ffi_try!(crocksdb_sstfilewriter_add(self.inner,
+                                                key.as_ptr(),
+                                                key.len(),
+                                                val.as_ptr(),
+                                                val.len()));
+            Ok(())
+        }
+    }
+
+    /// Finalize writing to sst file and close file.
+    pub fn finish(&mut self) -> Result<(), String> {
+        unsafe {
+            ffi_try!(crocksdb_sstfilewriter_finish(self.inner));
+            Ok(())
+        }
+    }
+}
+
+impl Drop for SstFileWriter {
+    fn drop(&mut self) {
+        unsafe { crocksdb_ffi::crocksdb_sstfilewriter_destroy(self.inner) }
+    }
+}
 
 #[cfg(test)]
 mod test {

tests/test_ingest_external_file.rs

@@ -16,10 +16,10 @@ use rocksdb::*;
 use std::fs;
 use tempdir::TempDir;
 
-fn gen_sst(opt: &Options, path: &str, data: &[(&[u8], &[u8])]) {
+fn gen_sst(opt: &Options, cf: &CFHandle, path: &str, data: &[(&[u8], &[u8])]) {
     let _ = fs::remove_file(path);
     let env_opt = EnvOptions::new();
-    let mut writer = SstFileWriter::new(&env_opt, opt);
+    let mut writer = SstFileWriter::new(&env_opt, opt, cf);
     writer.open(path).unwrap();
     for &(k, v) in data {
         writer.add(k, v).unwrap();
@@ -44,6 +44,7 @@ fn test_ingest_external_file() {
     let test_sstfile_str = test_sstfile.to_str().unwrap();
 
     gen_sst(db.get_options(),
+            db.cf_handle("default").unwrap(),
             test_sstfile_str,
             &[(b"k1", b"v1"), (b"k2", b"v2")]);
 
@@ -54,6 +55,7 @@ fn test_ingest_external_file() {
     assert_eq!(db.get(b"k2").unwrap().unwrap(), b"v2");
 
     gen_sst(&cf_opts,
+            handle,
             test_sstfile_str,
             &[(b"k1", b"v3"), (b"k2", b"v4")]);
     db.ingest_external_file_cf(handle, &ingest_opt, &[test_sstfile_str]).unwrap();
@@ -63,6 +65,7 @@ fn test_ingest_external_file() {
     let snap = db.snapshot();
 
     gen_sst(db.get_options(),
+            handle,
             test_sstfile_str,
             &[(b"k2", b"v5"), (b"k3", b"v6")]);
     ingest_opt = ingest_opt.move_files(true);