#!/usr/bin/python # -*- coding: utf-8 -*- # # Copyright 2012-2020 Spotify AB # # 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, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from builtins import object import ctypes import ctypes.util import future libsparkey = ctypes.cdll.LoadLibrary(ctypes.util.find_library("sparkey")) # Some constants class Compression(object): NONE = 0 SNAPPY = 1 class IterState(object): NEW = 0 ACTIVE = 1 CLOSED = 2 INVALID = 3 class IterType(object): PUT = 0 DELETE = 1 class SparkeyException(Exception): pass def _format(func, ret, *args): func.restype = ret func.argtypes = tuple(args) return func def _ctypes_wrapper(func, ret, *args): fn = _format(func, ret, *args) def wrapper(*a, **kw): code = fn(*a, **kw) if code != 0: raise SparkeyException(_errstring(code)) return wrapper _ptr = ctypes.c_void_p _str = ctypes.c_char_p _byref = ctypes.byref _create_string_buffer = ctypes.create_string_buffer _c_ulonglong = ctypes.c_ulonglong _errstring = _format(libsparkey.sparkey_errstring, _str, ctypes.c_int) _logwriter_create = _ctypes_wrapper(libsparkey.sparkey_logwriter_create, ctypes.c_int, _ptr, _str, ctypes.c_int, ctypes.c_int) _logwriter_append = _ctypes_wrapper(libsparkey.sparkey_logwriter_append, ctypes.c_int, _ptr, _str) _logwriter_close = _ctypes_wrapper(libsparkey.sparkey_logwriter_close, ctypes.c_int, _ptr) _logwriter_flush = _ctypes_wrapper(libsparkey.sparkey_logwriter_flush, ctypes.c_int, _ptr) _logwriter_put = _ctypes_wrapper(libsparkey.sparkey_logwriter_put, ctypes.c_int, _ptr, _c_ulonglong, _str, _c_ulonglong, _str) _logwriter_delete = _ctypes_wrapper(libsparkey.sparkey_logwriter_delete, ctypes.c_int, _ptr, _c_ulonglong, _str) _logreader_open = _ctypes_wrapper(libsparkey.sparkey_logreader_open, ctypes.c_int, _ptr, _str) _logreader_close = _format(libsparkey.sparkey_logreader_close, None, _ptr) _logiter_close = _format(libsparkey.sparkey_logiter_close, None, _ptr) _logiter_create = _ctypes_wrapper(libsparkey.sparkey_logiter_create, ctypes.c_int, _ptr, _ptr) _logiter_next = _ctypes_wrapper(libsparkey.sparkey_logiter_next, ctypes.c_int, _ptr, _ptr) _logiter_state = _format(libsparkey.sparkey_logiter_state, ctypes.c_int, _ptr) _logiter_type = _format(libsparkey.sparkey_logiter_type, ctypes.c_int, _ptr) _logiter_keylen = _format(libsparkey.sparkey_logiter_keylen, _c_ulonglong, _ptr) _logiter_valuelen = _format(libsparkey.sparkey_logiter_valuelen, _c_ulonglong, _ptr) _logiter_fill_key = _ctypes_wrapper(libsparkey.sparkey_logiter_fill_key, ctypes.c_int, _ptr, _ptr, _c_ulonglong, _str, ctypes.POINTER(_c_ulonglong)) _logiter_fill_value = _ctypes_wrapper(libsparkey.sparkey_logiter_fill_value, ctypes.c_int, _ptr, _ptr, _c_ulonglong, _str, ctypes.POINTER(_c_ulonglong)) _hash_write = _ctypes_wrapper(libsparkey.sparkey_hash_write, ctypes.c_int, _str, _str, ctypes.c_int) _hash_open = _ctypes_wrapper(libsparkey.sparkey_hash_open, ctypes.c_int, _ptr, _str, _str) _hash_close = _format(libsparkey.sparkey_hash_close, None, _ptr) _hash_getreader = _format(libsparkey.sparkey_hash_getreader, _ptr, _ptr) _logiter_hashnext = _ctypes_wrapper(libsparkey.sparkey_logiter_hashnext, ctypes.c_int, _ptr, _ptr) _hash_get = _ctypes_wrapper(libsparkey.sparkey_hash_get, ctypes.c_int, _ptr, _str, _c_ulonglong, _ptr) _hash_numentries = _format(libsparkey.sparkey_hash_numentries, _c_ulonglong, _ptr) if str == bytes: def _to_bytes(s, name): t = type(s) if t != str and t != future.types.newstr: raise SparkeyException(name + " must be a string") return s def _to_str(b, name): if b is None: return None if type(b) != str: raise SparkeyException(name + " must be a string") return b else: def _to_bytes(s, name): t = type(s) if t == bytes: return s if t != str: raise SparkeyException(name + " must be a string") return s.encode('utf-8') def _to_str(b, name): if b is None: return None t = type(b) if t == str: return b if t != bytes: raise SparkeyException(name + " must be bytes") return b.decode('utf-8') class LogWriter(object): def __init__(self, filename, mode='NEW', compression_type=Compression.NONE, compression_block_size=0): """Creates or appends a log file. Types of keys and values can be strings or bytes. For Python 2, this is the same thing. For Python 3, strings will be encoded as UTF-8 This is not threadsafe, don't write to the same file from multiple threads or processes. @param filename: file to create or append to. @param mode: one of two modes: - NEW: creates the file regardless of whether it already exists or not. - APPEND: appends to the log if it exists, otherwise raises an exception. @param compression_type: one of two types: - NONE: keys and values are written as is, and each key-value pair is considered a block of its own. - SNAPPY: compression is done on a block level of at most compression_block_size uncompressed bytes. Each block may contain multiple key/value pairs and it may split keys or values over block borders. @param compression_block_size: mandatory unless compression is NONE. This indicates how large the maximum block may be. To get good compression and performance, this should be a fairly small multiple of expected key + value size. """ filename = _to_bytes(filename, "filename") log = _ptr() self._log = log if mode == 'NEW': _logwriter_create(_byref(log), filename, compression_type, compression_block_size) elif mode == 'APPEND': _logwriter_append(_byref(log), filename) else: raise SparkeyException("Invalid mode %s, expected 'NEW' or " "'APPEND'" % (mode)) def __del__(self): self.close() def close(self): """Closes the writer (if not already closed). Also flushes all pending changes from memory to file. """ log = self._log if log is not None: self._log = None _logwriter_close(_byref(log)) def _assert_open(self): if self._log is None: raise SparkeyException("Writer is closed") def flush(self): """Flushes all pending changes from memory to file.""" self._assert_open() _logwriter_flush(self._log) def __setitem__(self, key, value): """Equivalent to put(key, value)""" self.put(key, value) def put(self, key, value): """Append the key-value pair to the log. @param key: type must be bytes or string @param value: type must be bytes or string """ self._assert_open() key = _to_bytes(key, "key") value = _to_bytes(value, "value") _logwriter_put(self._log, len(key), key, len(value), value) def __delitem__(self, key): """del writer[key] is equivalent to delete(key) (see L{delete})""" self.delete(key) def delete(self, key): """Appends a delete operation of key to the log. @param key: type must be bytes or string """ self._assert_open() key = _to_bytes(key, "key") _logwriter_delete(self._log, len(key), key) class LogReader(object): def __init__(self, filename): """Opens a file for log iteration. @param filename: file to open. """ filename = _to_bytes(filename, "filename") log = _ptr() self._log = log _logreader_open(_byref(log), filename) def __del__(self): self.close() def close(self): """Safely closes the log reader.""" log = self._log if log is not None: self._log = None _logreader_close(_byref(log)) def __iter__(self): """Creates a new iterator for this log reader. @returntype: L{LogIter} """ return LogIter(self) def _assert_open(self): if self._log is None: raise SparkeyException("Reader is closed") def _iter_res(iterator, log): state = _logiter_state(iterator) if state != IterState.ACTIVE: raise StopIteration() type_ = _logiter_type(iterator) keylen = _logiter_keylen(iterator) string_buffer = _create_string_buffer(keylen) length = _c_ulonglong() _logiter_fill_key(iterator, log, keylen, string_buffer, _byref(length)) if length.value != keylen: raise SparkeyException("Invalid keylen, expected %s but got %s" % (keylen, length.value)) key = string_buffer.raw valuelen = _logiter_valuelen(iterator) string_buffer = _create_string_buffer(valuelen) _logiter_fill_value(iterator, log, valuelen, string_buffer, _byref(length)) if length.value != valuelen: raise SparkeyException("Invalid valuelen, expected %s but got %s" % (valuelen, length.value)) value = string_buffer.raw return key, value, type_ class LogIter(object): def __init__(self, logreader): """Internal function. Use iter(logreader) or just "for key, value, type in logreader:" instead. """ logreader._assert_open() self._iter = _ptr() self._log = logreader _logiter_create(_byref(self._iter), logreader._log) def __del__(self): self.close() def close(self): """Safely closes the iterator.""" if self._iter is not None: _logiter_close(_byref(self._iter)) self._iter = None def __iter__(self): return self def _assert_open(self): if self._iter is None or self._log is None: raise SparkeyException("Iterator is closed") self._log._assert_open() def next(self): """Return next element in the log. @return: (key, value, type) if there are remaining elements. key and value are strings and type is a L{IterType}. @raise StopIteration: if there are no more entries in the log. """ self._assert_open() _logiter_next(self._iter, self._log._log) return _iter_res(self._iter, self._log._log) def __next__(self): return self.next() def writehash(hashfile, logfile, hash_size=0): """Write a hash file based on the contents in the log file. If the log file hasn't been changed since the existing hashfile was created, this is a no-op. @param hashfile: file to create. If it already exists, it will atomically be updated. @param logfile: file to read from. It must exist. @param hash_size: Valid values are 0, 4, 8. 0 means autoselect hash size. 4 is 32 bit hash, 8 is 64 bit hash. """ hashfile = _to_bytes(hashfile, "hashfile") logfile = _to_bytes(logfile, "logfile") _hash_write(hashfile, logfile, hash_size) class HashReader(object): """This is a reader that supports both iteration and random lookups.""" def __init__(self, hashfile, logfile): """Opens a hash file and log file for reading. @param hashfile: Hash file to open, must exist and be associated with the log file. @param logfile: Log file to open, must exist. """ hashfile = _to_bytes(hashfile, "hashfile") logfile = _to_bytes(logfile, "logfile") reader = _ptr() self._reader = reader self._iter = None _hash_open(_byref(reader), hashfile, logfile) self._iter = HashIterator(self) def __del__(self): self.close() def close(self): """Safely close the reader.""" reader = self._reader if reader is not None: _hash_close(_byref(reader)) self._reader = None if self._iter is not None: self._iter.close() self._iter = None def __iter__(self): """Equivalent to L{iteritems}""" return self.iteritems() def iteritems(self): """Iterate through all live entries. @returntype: L{HashIterator} """ return HashIterator(self) def _assert_open(self): if self._reader is None: raise SparkeyException("HashReader is closed") def __getitem__(self, key): """reader[key] throws KeyError exception when key doesn't exist, otherwise is equivalent to reader.get(key) (see L{get}) @param key: for the item **Note** in python 3 this always returns a bytes object, use getAsString(key) to return a String version. """ value = self.get(key) if value is None: raise KeyError(key) return value def __contains__(self, key): self._assert_open() iterator = self._iter._iter _hash_get(self._reader, key, len(key), iterator) res = True state = _logiter_state(iterator) if state != IterState.ACTIVE: res = False return res def has_key(self, key): return self.__contains__(key) def get(self, key): """Retrieve the value associated with the key @param key: type must be bytes or string @returns: bytes representing the value associated with the key, or None if the key does not exist. """ return self._iter.get(key) def getAsString(self, key): """Retrieve the value associated with the key @param key: type must be bytes or string @returns: a string representing the value associated with the key, or None if the key does not exist. """ return _to_str(self.get(key), "value") def __len__(self): return _hash_numentries(self._reader) class HashIterator(object): def __init__(self, hashreader): """Internal function: use iter(hashreader) instead.""" hashreader._assert_open() self._iter = _ptr() self._log = _hash_getreader(hashreader._reader) self._hashreader = hashreader _logiter_create(_byref(self._iter), self._log) def __del__(self): self.close() def close(self): """Safely closes the iterator.""" if self._iter is not None: _logiter_close(_byref(self._iter)) self._hashreader = None self._log = None self._iter = None def __iter__(self): return self def next(self): """Return next live entry in the log. @return: (key, value) if there are remaining elements. key and value are strings. @raise StopIteration: if there are no more live entries in the log. """ self._assert_open() _logiter_hashnext(self._iter, self._hashreader._reader) t = _iter_res(self._iter, self._log) if t: key, value, type = t return key, value def __next__(self): return self.next() def _assert_open(self): if self._hashreader is None: raise SparkeyException("Iterator is closed") self._hashreader._assert_open() def __getitem__(self, key): value = self.get(key) if value is None: raise KeyError return value def get(self, key): """Get the value associated with the key @param key: type must be bytes or string @returns: bytes representing the value associated with the key, or None if the key does not exist. """ key = _to_bytes(key, "key") self._assert_open() iterator = self._iter log = self._log _hash_get(self._hashreader._reader, key, len(key), iterator) state = _logiter_state(iterator) if state != IterState.ACTIVE: return None type_ = _logiter_type(iterator) assert type_ == IterType.PUT valuelen = _logiter_valuelen(iterator) string_buffer = _create_string_buffer(valuelen) clen = _c_ulonglong() _logiter_fill_value(iterator, log, valuelen, string_buffer, _byref(clen)) if clen.value != valuelen: raise SparkeyException("Invalid valuelen, expected %s but got %s" % (valuelen, clen.value)) value = string_buffer.raw return value def getAsString(self, key): """Retrieve the value associated with the key @param key: type must be bytes or string @returns: a string representing the value associated with the key, or None if the key does not exist. """ return _to_str(self.get(key), "value") class HashWriter(object): def __init__(self, hashfile, logfile, mode='NEW', compression_type=Compression.NONE, compression_block_size=0, hash_size=0): """Creates a new writer. Does everything that L{LogWriter} does, but also writes the hash file. @param hashfile: filename of hash file @param logfile: filename of log file @param mode: Same as in L{LogWriter.__init__} @param compression_type: Same as in L{LogWriter.__init__} @param compression_block_size: Same as in L{LogWriter.__init__} @param hash_size: Valid values are 0, 4, 8. 0 means autoselect hash size . 4 is 32 bit hash, 8 is 64 bit hash. """ self._logwriter = LogWriter(logfile, mode, compression_type, compression_block_size) self._hashfile = hashfile self._logfile = logfile self._reader = None self._hash_size = hash_size def _assert_open(self): if self._logwriter is None: raise SparkeyException("Writer is closed") def __setitem__(self, key, value): """Equivalent to writer.put(key, value), see L{put}""" self.put(key, value) def put(self, k, v): """Append the key-value pair to the log. @param key: type must be bytes or string @param value: type must be bytes or string """ self._logwriter.put(k, v) def __delitem__(self, key): """Equivalent to writer.delete(key), see L{delete}""" self.delete(key) def delete(self, k): """Appends a delete operation of key to the log. @param key: type must be bytes or string """ self._assert_open() self._logwriter.delete(k) def flush(self): """Flushes all log writes, and also rebuilds the hash.""" self._assert_open() self._logwriter.flush() writehash(self._hashfile, self._logfile, self._hash_size) def __del__(self): self.destroy() def destroy(self): """Closes the writer, but does not flush anything. All writes before the previous flush will be gone. """ if self._logwriter is not None: self._logwriter.close() self._logwriter = None self._close_reader() self._hashfile = None self._logfile = None def finish(self): """Equivalent to L{close}""" self.close() def close(self): """Flushes pending log writes from memory to disk, rewrites the hash file and closes the writer. """ if self._logwriter is not None: self.flush() self.destroy() # Reader related code def _close_reader(self): if self._reader is not None: self._reader.close() self._reader = None def _init_reader(self): if self._reader is None: self._reader = HashReader(self._hashfile, self._logfile) return self._reader def __iter__(self): """Equivalent to L{iteritems}""" return self.iteritems() def iteritems(self): """Iterate through all entries that have been flushed. @returns: L{HashIterator} """ self._assert_open() return self._init_reader().iteritems() def __getitem__(self, key): """Equivalent to writer.get(key), see L{get}""" self._assert_open() return self._init_reader().get(key) def get(self, key): """Performs a hash lookup of a key. Only finds things that were flushed to the hash. @param key: type must be bytes or string @returns: bytes representing the value associated with the key, or None if the key does not exist in the hash. """ self._assert_open() return self._init_reader().get(key) def getAsString(self, key): """Performs a hash lookup of a key. Only finds things that were flushed to the hash. @param key: type must be bytes or string @returns: a string representing the value associated with the key, or None if the key does not exist in the hash. """ return _to_str(self.get(key), "value")