1  # Nim-RocksDB
  2  # Copyright 2024 Status Research & Development GmbH
  3  # Licensed under either of
  4  #
  5  #  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0)
  6  #  * GPL license, version 2.0, ([LICENSE-GPLv2](LICENSE-GPLv2) or https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html)
  7  #
  8  # at your option. This file may not be copied, modified, or distributed except according to those terms.
  9  
 10  {.used.}
 11  
 12  import std/os, tempfile, unittest2, ../rocksdb/[rocksdb, writebatchwi], ./test_helper
 13  
 14  suite "WriteBatchWIRef Tests":
 15    const
 16      CF_DEFAULT = "default"
 17      CF_OTHER = "other"
 18  
 19    let
 20      key1 = @[byte(1)]
 21      val1 = @[byte(1)]
 22      key2 = @[byte(2)]
 23      val2 = @[byte(2)]
 24      key3 = @[byte(3)]
 25      val3 = @[byte(3)]
 26  
 27    setup:
 28      let
 29        dbPath = mkdtemp() / "data"
 30        db = initReadWriteDb(dbPath, columnFamilyNames = @[CF_DEFAULT, CF_OTHER])
 31        defaultCfHandle = db.getColFamilyHandle(CF_DEFAULT).get()
 32        otherCfHandle = db.getColFamilyHandle(CF_OTHER).get()
 33  
 34    teardown:
 35      db.close()
 36      removeDir($dbPath)
 37  
 38    test "Test writing batch to the default column family":
 39      let batch = db.openWriteBatchWithIndex()
 40      defer:
 41        batch.close()
 42      check not batch.isClosed()
 43  
 44      check:
 45        batch.put(key1, val1).isOk()
 46        batch.put(key2, val2).isOk()
 47        batch.put(key3, val3).isOk()
 48        batch.count() == 3
 49  
 50        batch.delete(key2).isOk()
 51        batch.count() == 4
 52        not batch.isClosed()
 53  
 54        batch.getFromBatch(key1).get() == val1
 55        batch.getFromBatch(key2).isErr()
 56        batch.getFromBatch(key3).get() == val3
 57  
 58      let res = db.write(batch)
 59      check:
 60        res.isOk()
 61        db.write(batch).isOk() # test that it's idempotent
 62        db.get(key1).get() == val1
 63        db.keyExists(key2).get() == false
 64        db.get(key3).get() == val3
 65  
 66        batch.getFromBatch(key1).get() == val1
 67        batch.getFromBatch(key2).isErr()
 68        batch.getFromBatch(key3).get() == val3
 69  
 70      batch.clear()
 71      check:
 72        batch.count() == 0
 73        not batch.isClosed()
 74  
 75    test "Test writing batch to column family":
 76      let batch = db.openWriteBatchWithIndex()
 77      defer:
 78        batch.close()
 79      check not batch.isClosed()
 80  
 81      check:
 82        batch.put(key1, val1, otherCfHandle).isOk()
 83        batch.put(key2, val2, otherCfHandle).isOk()
 84        batch.put(key3, val3, otherCfHandle).isOk()
 85        batch.count() == 3
 86  
 87        batch.delete(key2, otherCfHandle).isOk()
 88        batch.count() == 4
 89        not batch.isClosed()
 90  
 91        batch.getFromBatch(key1, otherCfHandle).get() == val1
 92        batch.getFromBatch(key2, otherCfHandle).isErr()
 93        batch.getFromBatch(key3, otherCfHandle).get() == val3
 94  
 95      let res = db.write(batch)
 96      check:
 97        res.isOk()
 98        db.get(key1, otherCfHandle).get() == val1
 99        db.keyExists(key2, otherCfHandle).get() == false
100        db.get(key3, otherCfHandle).get() == val3
101  
102        batch.getFromBatch(key1, otherCfHandle).get() == val1
103        batch.getFromBatch(key2, otherCfHandle).isErr()
104        batch.getFromBatch(key3, otherCfHandle).get() == val3
105  
106      batch.clear()
107      check:
108        batch.count() == 0
109        not batch.isClosed()
110  
111    test "Test writing to multiple column families in single batch":
112      let batch = db.openWriteBatchWithIndex()
113      defer:
114        batch.close()
115      check not batch.isClosed()
116  
117      check:
118        batch.put(key1, val1, defaultCfHandle).isOk()
119        batch.put(key1, val1, otherCfHandle).isOk()
120        batch.put(key2, val2, otherCfHandle).isOk()
121        batch.put(key3, val3, otherCfHandle).isOk()
122        batch.count() == 4
123  
124        batch.delete(key2, otherCfHandle).isOk()
125        batch.count() == 5
126        not batch.isClosed()
127  
128      let res = db.write(batch)
129      check:
130        res.isOk()
131        db.get(key1, defaultCfHandle).get() == val1
132        db.get(key1, otherCfHandle).get() == val1
133        db.keyExists(key2, otherCfHandle).get() == false
134        db.get(key3, otherCfHandle).get() == val3
135  
136      batch.clear()
137      check:
138        batch.count() == 0
139        not batch.isClosed()
140  
141    test "Test writing to multiple column families in multiple batches":
142      let
143        batch1 = db.openWriteBatchWithIndex()
144        batch2 = db.openWriteBatchWithIndex()
145      defer:
146        batch1.close()
147        batch2.close()
148  
149      check:
150        not batch1.isClosed()
151        not batch2.isClosed()
152        batch1.put(key1, val1).isOk()
153        batch1.delete(key2, otherCfHandle).isOk()
154        batch1.put(key3, val3, otherCfHandle).isOk()
155        batch2.put(key1, val1, otherCfHandle).isOk()
156        batch2.delete(key1, otherCfHandle).isOk()
157        batch2.put(key3, val3).isOk()
158        batch1.count() == 3
159        batch2.count() == 3
160  
161      let res1 = db.write(batch1)
162      let res2 = db.write(batch2)
163      check:
164        res1.isOk()
165        res2.isOk()
166        db.get(key1).get() == val1
167        db.keyExists(key2).get() == false
168        db.get(key3).get() == val3
169        db.keyExists(key1, otherCfHandle).get() == false
170        db.keyExists(key2, otherCfHandle).get() == false
171        db.get(key3, otherCfHandle).get() == val3
172  
173        # Write batch is unchanged after write
174        batch1.count() == 3
175        batch2.count() == 3
176        not batch1.isClosed()
177        not batch2.isClosed()
178  
179    test "Test write empty batch":
180      let batch = db.openWriteBatchWithIndex()
181      defer:
182        batch.close()
183      check not batch.isClosed()
184  
185      check batch.count() == 0
186      let res1 = db.write(batch)
187      check:
188        res1.isOk()
189        batch.count() == 0
190        not batch.isClosed()
191  
192    test "Test multiple writes to same key":
193      let
194        batch1 = db.openWriteBatchWithIndex(overwriteKey = false)
195        batch2 = db.openWriteBatchWithIndex(overwriteKey = true)
196      defer:
197        batch1.close()
198        batch2.close()
199      check:
200        not batch1.isClosed()
201        not batch2.isClosed()
202  
203      check:
204        batch1.put(key1, val1).isOk()
205        batch1.delete(key1).isOk()
206        batch1.put(key1, val3).isOk()
207        batch1.count() == 3
208        batch1.getFromBatch(key1).get() == val3
209  
210        batch2.put(key1, val3).isOk()
211        batch2.put(key1, val2).isOk()
212        batch2.put(key1, val1).isOk()
213        batch2.count() == 3
214        batch2.getFromBatch(key1).get() == val1
215  
216    test "Put, get and delete empty key":
217      let batch = db.openWriteBatchWithIndex()
218      defer:
219        batch.close()
220  
221      let empty: seq[byte] = @[]
222      check:
223        batch.put(empty, val1).isOk()
224        batch.getFromBatch(empty).get() == val1
225        batch.delete(empty).isOk()
226        batch.getFromBatch(empty).isErr()
227  
228    test "Test close":
229      let batch = db.openWriteBatchWithIndex()
230  
231      check not batch.isClosed()
232      batch.close()
233      check batch.isClosed()
234      batch.close()
235      check batch.isClosed()