├── logo.png ├── .gitignore ├── .github └── workflows │ └── tests.yaml ├── test_g.py ├── setup.py ├── README.md ├── greenquery.py ├── tests.py └── greendb.py /logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/coleifer/greendb/HEAD/logo.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | data/ 2 | config.json 3 | MANIFEST 4 | build 5 | dist 6 | *.pyc 7 | -------------------------------------------------------------------------------- /.github/workflows/tests.yaml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | on: [push] 3 | jobs: 4 | tests: 5 | name: ${{ matrix.python-version }} 6 | runs-on: ubuntu-latest 7 | strategy: 8 | fail-fast: false 9 | matrix: 10 | python-version: [3.7, "3.10", "3.11"] 11 | steps: 12 | - uses: actions/checkout@v2 13 | - uses: actions/setup-python@v2 14 | with: 15 | python-version: ${{ matrix.python-version }} 16 | - name: pip deps 17 | run: pip install gevent lmdb msgpack-python 18 | - name: runtests 19 | run: python tests.py 20 | -------------------------------------------------------------------------------- /test_g.py: -------------------------------------------------------------------------------- 1 | """ 2 | Test simpledb with many concurrent connections. 3 | """ 4 | from gevent import monkey; monkey.patch_all() 5 | 6 | import time 7 | 8 | import gevent 9 | from greendb import Client 10 | 11 | client = Client() 12 | 13 | def get_sleep_set(k, v, n=1): 14 | client.set(k, v) 15 | #time.sleep(n) 16 | client._sleep(n) 17 | assert client.get(k) == v 18 | data = {b'%s-%032d' % (k, i): 'v%01024d' % i for i in range(100)} 19 | client.mset(data) 20 | resp = client.mget([b'%s-%032d' % (k, i) for i in range(100)]) 21 | assert resp == data 22 | # client.close() # If this were a real app, we would put this here. 23 | 24 | n = 1 25 | t = 100 26 | start = time.time() 27 | 28 | greenlets = [] 29 | for i in range(t): 30 | greenlets.append(gevent.spawn(get_sleep_set, b'k%d' % i, b'v%d' % i, n)) 31 | 32 | for g in greenlets: 33 | g.join() 34 | 35 | client.flush() 36 | stop = time.time() 37 | print('done. slept=%s, took %.2f for %s threads.' % (n, stop - start, t)) 38 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from setuptools import setup 3 | 4 | 5 | with open(os.path.join(os.path.dirname(__file__), 'README.md')) as fh: 6 | readme = fh.read() 7 | 8 | try: 9 | version = __import__('greendb').__version__ 10 | except ImportError: 11 | version = '0.2.3' 12 | 13 | 14 | setup( 15 | name='greendb', 16 | version=version, 17 | description='greendb', 18 | long_description='greendb', 19 | author='Charles Leifer', 20 | author_email='coleifer@gmail.com', 21 | url='http://github.com/coleifer/greendb/', 22 | packages=[], 23 | py_modules=['greendb', 'greenquery'], 24 | install_requires=['lmdb', 'gevent', 'msgpack-python'], 25 | classifiers=[ 26 | 'Development Status :: 4 - Beta', 27 | 'Environment :: Web Environment', 28 | 'Intended Audience :: Developers', 29 | 'License :: OSI Approved :: MIT License', 30 | 'Operating System :: OS Independent', 31 | 'Programming Language :: Python', 32 | ], 33 | scripts=['greendb.py'], 34 | test_suite='tests') 35 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![](http://media.charlesleifer.com/blog/photos/logo-0.png) 2 | 3 | ### greendb 4 | 5 | async server frontend for symas lmdb. 6 | 7 | #### description 8 | 9 | greendb is a lightweight server (and Python client) for symas lmdb. The server 10 | uses the new Redis RESPv3 protocol. 11 | 12 | greendb supports multiple independent databases, much like Redis. Values are 13 | serialized using `msgpack`, so the server is capable of storing all the 14 | data-types supported by msgpack. greendb also provides per-database 15 | configuration of multi-value support, allowing you to efficiently store 16 | multiple values at a given key, in sorted order (e.g. for secondary indexes). 17 | 18 | With greendb, database keys are always bytestrings. Values may be: 19 | 20 | * dict 21 | * list 22 | * set 23 | * bytestrings 24 | * unicode strings 25 | * integers 26 | * floating-point 27 | * boolean 28 | * `None` 29 | 30 | #### installing 31 | 32 | ``` 33 | $ pip install greendb 34 | ``` 35 | 36 | Alternatively, you can install from git: 37 | 38 | ``` 39 | $ git clone https://github.com/coleifer/greendb 40 | $ cd greendb 41 | $ python setup.py install 42 | ``` 43 | 44 | Dependencies: 45 | 46 | * [gevent](http://www.gevent.org/) 47 | * [lmdb](https://github.com/dw/py-lmdb) 48 | * [msgpack-python](https://github.com/msgpack/msgpack-python) 49 | 50 | #### running 51 | 52 | ``` 53 | $ greendb.py -h 54 | 55 | Usage: greendb.py [options] 56 | 57 | Options: 58 | -h, --help show this help message and exit 59 | -c CONFIG, --config=CONFIG 60 | Config file (default="config.json") 61 | -D DATA_DIR, --data-dir=DATA_DIR 62 | Directory to store db environment and data. 63 | -d, --debug Log debug messages. 64 | -e, --errors Log error messages only. 65 | -H HOST, --host=HOST Host to listen on. 66 | -l LOG_FILE, --log-file=LOG_FILE 67 | Log file. 68 | -m MAP_SIZE, --map-size=MAP_SIZE 69 | Maximum size of memory-map used for database. The 70 | default value is 256M and should be increased. Accepts 71 | value in bytes or file-size using "M" or "G" suffix. 72 | --max-clients=MAX_CLIENTS 73 | Maximum number of clients. 74 | -n MAX_DBS, --max-dbs=MAX_DBS 75 | Number of databases in environment. Default=16. 76 | -p PORT, --port=PORT Port to listen on. 77 | -r, --reset Reset database and config. All data will be lost. 78 | -s, --sync Flush system buffers to disk when committing a 79 | transaction. Durable but much slower. 80 | -u DUPSORT, --dupsort=DUPSORT 81 | db index(es) to support dupsort 82 | -M, --no-metasync Flush system buffers to disk only once per 83 | transaction, omit the metadata flush. 84 | -W, --writemap Use a writeable memory map. 85 | -A, --map-async When used with "--writemap" (-W), use asynchronous 86 | flushes to disk. 87 | ``` 88 | 89 | Complete config file example with default values: 90 | 91 | ```javascript 92 | 93 | { 94 | "host": "127.0.0.1", 95 | "port": 31337, 96 | "max_clients": 1024, 97 | "path": "data", // Directory for data storage, default is "data" in CWD. 98 | "map_size": "256M", // Default map size is 256MB. INCREASE THIS! 99 | "read_only": false, // Open the database in read-only mode. 100 | "metasync": true, // Sync metadata changes (recommended). 101 | "sync": false, // Sync all changes (durable, but much slower). 102 | "writemap": false, // Use a writable map (probably safe to do). 103 | "map_async": false, // Asynchronous writable map. 104 | "meminit": true, // Initialize new memory pages to zero. 105 | "max_dbs": 16, // Maximum number of DBs. 106 | "max_spare_txns": 64, 107 | "lock": true, // Lock the database when opening environment. 108 | "dupsort": false // Either a boolean or a list of DB indexes. 109 | } 110 | ``` 111 | 112 | Example custom configuration: 113 | 114 | * 1GB max database size 115 | * dupsort enabled on databases 13, 14 and 15 116 | * data stored in /var/lib/greendb/data 117 | 118 | ```javascript 119 | { 120 | "map_size": "1G", 121 | "dupsort": [13, 14, 15], 122 | "path": "/var/lib/greendb/data/" 123 | } 124 | ``` 125 | 126 | Equivalent configuration using command-line arguments: 127 | 128 | ``` 129 | $ greendb.py -m 1G -u 13 -u 14 -u 15 -D /var/lib/greendb/data/ 130 | ``` 131 | 132 | ### client 133 | 134 | A Python client is included in the `greendb` module. All server commands are 135 | implemented as client methods using the lower-case command name, for example: 136 | 137 | ```python 138 | 139 | from greendb import Client 140 | client = Client(host='10.0.0.3') 141 | 142 | # Execute the ENVINFO command. 143 | print(client.envinfo()) 144 | 145 | # Set multiple key/value pairs, read a key, then delete two keys. 146 | client.mset({'k1': 'v1', 'k2': 'v2'}) # MSET 147 | print(client.get('k1')) # GET 148 | client.mdelete(['k1', 'k2']) # MDELETE 149 | ``` 150 | 151 | Additionally, the `Client` implements much of the Python `dict` interface, such 152 | as item get/set/delete, iteration, length, contains, etc. 153 | 154 | If an error occurs, either due to a malformed command (e.g., missing required 155 | parameters) or for any other reason (e.g., attempting to write to a read-only 156 | database), then a `CommandError` will be raised by the client with a message 157 | indicating what caused the error. 158 | 159 | **A note about connections**: the greendb client will automatically connect the 160 | first time you issue a command to the server. The client maintains its own 161 | thread-safe (and greenlet-safe) connection pool. If you wish to explicitly 162 | connect, the client may be used as a context manager. 163 | 164 | In a multi-threaded or multi-greenlet application (e.g. a web app), the client 165 | will maintain a separate connection for each thread/greenlet. 166 | 167 | ### command reference 168 | 169 | Below is the list of supported commands. **Commands are available on the client 170 | using the lower-case command name as the method name**. 171 | 172 | 173 | 174 | 175 | 176 | 177 | 178 | 179 | 180 | 181 | 182 | 183 | 184 | 185 | 186 | 187 | 188 | 189 | 190 | 191 | 192 | 193 | 194 | 195 | 196 | 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | 214 | 215 | 216 | 217 | 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 | 230 | 231 | 232 | 233 | 234 | 235 | 236 | 237 | 238 | 239 | 240 | 241 | 242 | 243 | 244 | 245 | 246 | 247 | 248 | 249 | 250 | 251 | 252 | 253 | 254 | 255 | 256 | 257 | 258 | 259 | 260 | 261 | 262 | 263 | 264 | 265 | 266 | 267 | 268 | 270 | 271 | 272 | 273 | 274 | 275 | 276 | 277 | 278 | 279 | 280 | 281 | 282 | 283 | 284 | 285 | 286 | 287 | 289 | 290 | 291 | 292 | 293 | 294 | 295 | 296 | 297 | 298 | 299 | 300 | 302 | 303 | 304 | 305 | 306 | 307 | 310 | 311 | 312 | 313 | 314 | 315 | 318 | 319 | 320 | 321 | 322 | 323 | 326 | 327 | 329 | 330 | 331 | 332 | 335 | 336 | 337 | 338 | 339 | 340 | 343 | 344 | 345 | 346 | 347 | 348 | 349 | 350 | 351 | 352 | 353 | 354 | 355 | 356 | 357 | 358 | 359 | 360 | 361 | 362 | 363 | 364 | 365 | 367 | 368 | 369 | 370 | 371 | 372 | 374 | 375 | 376 | 377 | 380 | 381 | 382 | 383 | 384 | 385 | 388 | 389 | 391 | 392 | 393 | 394 | 395 | 396 | 397 | 398 | 399 | 400 | 403 | 404 | 405 | 406 | 407 | 408 | 409 | 410 | 411 | 412 | 413 | 414 | 415 | 416 | 417 | 418 | 419 | 420 | 421 | 422 | 423 | 425 | 426 | 427 | 428 | 429 | 430 | 432 | 433 | 434 | 435 | 436 | 437 | 438 | 439 | 440 | 441 | 442 | 443 | 444 | 445 | 446 | 447 | 448 | 449 | 450 | 451 | 452 | 453 | 454 | 455 | 456 | 457 | 458 | 459 | 460 | 461 | 462 | 463 | 464 | 465 | 466 | 467 | 468 |
commanddescriptionargumentsreturn value
ENVINFOmetadata and storage configuration settings(none)dict
ENVSTATmetadata related to the global b-tree(none)dict
FLUSHdelete all records in the currently-selected database(none)boolean indicating success
FLUSHALLdelete all records in all databases(none)dict mapping database index to boolean
PINGping the server(none)"pong"
STATmetadata related to the currently-selected database b-tree(none)dict
SYNCsynchronize database to disk (use when sync=False)(none)(none)
USEselect the given databasedatabase index, 0 through (max_dbs - 1)int: active database index
KV commands
COUNTget the number of key/value pairs in active database(none)int
DECRdecrement the value at the given keyamount to decrement by (optional, default is 1)int or float
INCRincrement the value at the given keyamount to increment by (optional, default is 1)int or float
CAScompare-and-setkey, original value, new valueboolean indicating success or failure
DELETEdelete a key and any value(s) associatedkeyint: number of keys removed (1 on success, 0 if key not found)
DELETEDUPdelete a particular key/value pair when dupsort is enabledkey, value to deleteint: number of key+value removed (1 on success, 0 if key+value not found)
DELETEDUPRAWdelete a particular key/value pair when dupsort is enabled using an 269 | unserialized bytestring as the valuekey, value to deleteint: number of key+value removed (1 on success, 0 if key+value not found)
DUPCOUNTget number of values stored at the given key (requires dupsort)keyint: number of values, or None if key does not exist
EXISTSdetermine if the given key existskeybool
GETget the value associated with a given key. If dupsort is enabled and 288 | multiple values are present, the one that is sorted first will be returned.keyvalue or None
GETDUPget all values associated with a given key (requires dupsort)keylist of values or None if key does not exist
LENGTHget the length of the value stored at a given key, e.g. for a string 301 | this returns the number of characters, for a list the number of items, etc.keylength of value or None if key does not exist
POPatomically get and delete the value at a given key. If dupsort is 308 | enabled and multiple values are present, the one that is sorted first will 309 | be removed and returned.keyvalue or None
REPLACEatomically get and set the value at a given key. If dupsort is enabled, 316 | the first value will be returned (if exists) and ALL values will be removed 317 | so that only the new value is stored.keyprevious value or None
SETstore a key/value pair. If dupsort is enabled, duplicate values will be 324 | stored at the given key in sorted-order. Additionally, if dupsort is 325 | enabled and the exact key/value pair already exist, no changes are made.key, valueint: 1 if new key/value added, 0 if dupsort is enabled and the 328 | key/value already exist
SETDUPstore a key/value pair, treating duplicates as successful writes 333 | (requires dupsort). Unlike SET, if the exact key/value pair already exists, 334 | this command will return 1 indicating success.key, valueint: 1 on success
SETDUPRAWstore a key/value pair, treating duplicates as successful writes 341 | (requires dupsort). Additionally, the value is not serialized, but is 342 | stored as a raw bytestring.key, value (bytes)int: 1 on success
SETNXstore a key/value pair only if the key does not existkey, valueint: 1 on success, 0 if key already exists
Bulk KV commands
MDELETEdelete multiple keyslist of keysint: number of keys deleted
MGETget the value of multiple keyslist of keysdict of key and value. Keys that were requested, but which do not 366 | exist are not included in the response.
MGETDUPget all values of multiple keys (requires dupsort)list of keysdict of key to list of values. Keys that were requested, but which do 373 | not exist, are not included in the response.
MPOPatomically get and delete the value of multiple keys. If dupsort is 378 | enabled and multiple values are stored at a given key, only the first value 379 | will be removed.list of keysdict of key to value
MREPLACEatomically get and set the value of multiple keys. If dupsort is 386 | enabled and multiple values are stored at a given key, only the first value 387 | will be returned and all remaining values discarded.dict of key to valuedict of key to previous value. Keys that did not exist previously will 390 | not be included in the response.
MSETset the value of multiple keys.dict of key to valueint: number of key / value pairs set
MSETDUPstore multiple key/value pairs, treating duplicates as successful 401 | writes (requires dupsort). Unlike MSET, if the exact key/value pair already 402 | exists, this command will treat the write as a success.dict of key to valueint: number of key / value pairs set
MSETNXstore multiple key/value pair only if the key does not existdict of key to valueint: number of key / value pairs set
Cursor / range commands
DELETERANGEdelete a range of keys using optional inclusive start/end-pointsstart key (optional), end key (optional), count (optional)int: number of keys deleted
GETRANGEretrieve a range of key/value pairs using optional inclusive 424 | start/end-pointsstart key (optional), end key (optional), count (optional)list of [key, value] lists
GETRANGEDUPRAWretrieve a range of duplicate values stored in a given key, using 431 | optional (inclusive) start/end-pointskey, start value (optional), end value (optional), count (optional)list of values (as bytestrings)
KEYSretrieve a range of keys using optional inclusive start/end-pointsstart key (optional), end key (optional), count (optional)list of keys
PREFIXretrieve a range of key/value pairs which match the given prefixprefix, count (optional)list of [key, value] lists
VALUESretrieve a range of values using optional inclusive start/end-pointsstart key (optional), end key (optional), count (optional)list of values
Client commands
QUITdisconnect from server(none)int: 1
SHUTDOWNterminate server process from client (be careful!)(none)(none)
469 | 470 | #### protocol 471 | 472 | The protocol is the Redis RESPv3 protocol. I have extended the protocol with 473 | one additional data-type, a dedicated type for representing UTF8-encoded 474 | unicode text (notably absent from RESPv3). This type is denoted by the leading 475 | `^` byte in responses. 476 | 477 | Details can be found here: https://github.com/antirez/RESP3/blob/master/spec.md 478 | -------------------------------------------------------------------------------- /greenquery.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | import struct 3 | import sys 4 | import time 5 | 6 | from greendb import Client 7 | 8 | 9 | __all__ = [ 10 | 'Client', 11 | 'Field', 12 | 'IntegerField', 13 | 'LongField', 14 | 'DateTimeField', 15 | 'TimestampField', 16 | 'BooleanField', 17 | 'Model', 18 | ] 19 | 20 | 21 | if sys.version_info[0] == 2: 22 | unicode_type = unicode 23 | else: 24 | unicode_type = str 25 | 26 | 27 | class Node(object): 28 | def __init__(self): 29 | self.negated = False 30 | 31 | def _e(op, inv=False): 32 | """ 33 | Lightweight factory which returns a method that builds an Expression 34 | consisting of the left-hand and right-hand operands, using `op`. 35 | """ 36 | def inner(self, rhs): 37 | if inv: 38 | return Expression(rhs, op, self) 39 | return Expression(self, op, rhs) 40 | return inner 41 | 42 | __and__ = _e('AND') 43 | __or__ = _e('OR') 44 | __rand__ = _e('AND', inv=True) 45 | __ror__ = _e('OR', inv=True) 46 | __eq__ = _e('=') 47 | __ne__ = _e('!=') 48 | __lt__ = _e('<') 49 | __le__ = _e('<=') 50 | __gt__ = _e('>') 51 | __ge__ = _e('>=') 52 | 53 | def between(self, start, stop, start_inclusive=True, stop_inclusive=False): 54 | return Expression(self, 'between', (start, stop, start_inclusive, 55 | stop_inclusive)) 56 | 57 | def startswith(self, prefix): 58 | return Expression(self, 'startswith', prefix) 59 | 60 | 61 | class Expression(Node): 62 | def __init__(self, lhs, op, rhs): 63 | self.lhs = lhs 64 | self.op = op 65 | self.rhs = rhs 66 | 67 | def __repr__(self): 68 | return '' % (self.lhs, self.op, self.rhs) 69 | 70 | 71 | def encode(s): 72 | if isinstance(s, bytes): 73 | return s 74 | elif isinstance(s, unicode_type): 75 | return s.encode('utf8') 76 | return str(s).encode('utf8') 77 | 78 | 79 | class Field(Node): 80 | _counter = 0 81 | 82 | def __init__(self, index=False, default=None): 83 | self.index = index 84 | self.default = default 85 | self.model = None 86 | self.name = None 87 | self._order = Field._counter 88 | Field._counter += 1 89 | 90 | def __repr__(self): 91 | return '<%s: %s.%s>' % ( 92 | type(self), 93 | self.model._meta.model_name, 94 | self.name) 95 | 96 | def bind(self, model, name): 97 | self.model = model 98 | self.name = name 99 | setattr(self.model, self.name, FieldDescriptor(self)) 100 | 101 | def clone(self): 102 | field = type(self)(index=self.index, default=self.default) 103 | field.model = self.model 104 | field.name = self.name 105 | return field 106 | 107 | def serialize(self, value): 108 | return value 109 | 110 | def deserialize(self, value): 111 | return value 112 | 113 | def index_value(self, value): 114 | return encode(value) 115 | 116 | 117 | class IntegerField(Field): 118 | def index_value(self, value): 119 | return struct.pack('>H', value) 120 | 121 | 122 | class LongField(Field): 123 | def index_value(self, value): 124 | return struct.pack('>Q', value) 125 | 126 | 127 | class DateTimeField(Field): 128 | def serialize(self, value): 129 | return value.strftime('%Y-%m-%d %H:%M:%S.%f') 130 | 131 | def deserialize(self, value): 132 | return datetime.datetime.strptime(value, '%Y-%m-%d %H:%M:%S.%f') 133 | 134 | def index_value(self, value): 135 | timestamp = time.mktime(value.timetuple()) 136 | timestamp += value.microsecond * .000001 137 | timestamp = int(timestamp * 1e6) 138 | return struct.pack('>Q', timestamp) 139 | 140 | 141 | class TimestampField(Field): 142 | def __init__(self, index=False, default=datetime.datetime.now): 143 | super(TimestampField, self).__init__(index=index, default=default) 144 | 145 | def serialize(self, value): 146 | timestamp = time.mktime(value.timetuple()) 147 | timestamp += value.microsecond * .000001 148 | timestamp = int(timestamp * 1e6) 149 | return struct.pack('>Q', timestamp) 150 | 151 | def deserialize(self, value): 152 | raw_ts, = struct.unpack('>Q', value) 153 | timestamp, micro = divmod(raw_ts, 1e6) 154 | return (datetime.datetime 155 | .fromtimestamp(timestamp) 156 | .replace(microsecond=int(micro))) 157 | 158 | index_value = serialize 159 | 160 | 161 | class BooleanField(Field): 162 | def index_value(self, value): 163 | return b'\x01' if value else b'\x00' 164 | 165 | 166 | class FieldDescriptor(object): 167 | def __init__(self, field): 168 | self.field = field 169 | self.name = self.field.name 170 | 171 | def __get__(self, instance, instance_type=None): 172 | if instance is not None: 173 | return instance._data.get(self.name) 174 | return self.field 175 | 176 | def __set__(self, instance, value): 177 | instance._data[self.name] = value 178 | 179 | 180 | class DeclarativeMeta(type): 181 | def __new__(cls, name, bases, attrs): 182 | if bases == (object,): 183 | return super(DeclarativeMeta, cls).__new__(cls, name, bases, attrs) 184 | 185 | client = None 186 | db = None 187 | index_db = None 188 | fields = {} 189 | 190 | # Inherit fields from parent classes. 191 | for base in bases: 192 | if not hasattr(base, '_meta'): 193 | continue 194 | 195 | for field in base._meta.sorted_fields: 196 | if field.name not in fields: 197 | fields[field.name] = field.clone() 198 | 199 | if client is None and base._meta.client is not None: 200 | client = base._meta.client 201 | if db is None: 202 | db = base._meta.db 203 | if index_db is None: 204 | index_db = base._meta.index_db 205 | 206 | # Apply defaults if no value was found. 207 | db = 0 if db is None else db 208 | index_db = 15 if index_db is None else index_db 209 | 210 | # Introspect all declared fields. 211 | for key, value in attrs.items(): 212 | if isinstance(value, Field): 213 | fields[key] = value 214 | 215 | # Read metadata configuration. 216 | declared_meta = attrs.pop('Meta', None) 217 | if declared_meta: 218 | if getattr(declared_meta, 'client', None) is not None: 219 | client = declared_meta.client 220 | if getattr(declared_meta, 'db', None) is not None: 221 | db = declared_meta.db 222 | if getattr(declared_meta, 'index_db', None) is not None: 223 | index_db = declared_meta.index_db 224 | 225 | # Always have an `id` field. 226 | if 'id' not in fields: 227 | fields['id'] = LongField() 228 | 229 | attrs['_meta'] = Metadata(name, client, db, index_db, fields) 230 | model = super(DeclarativeMeta, cls).__new__(cls, name, bases, attrs) 231 | 232 | # Bind fields to model. 233 | for name, field in fields.items(): 234 | field.bind(model, name) 235 | 236 | # Process 237 | model._meta.prepared() 238 | 239 | return model 240 | 241 | 242 | class Metadata(object): 243 | def __init__(self, model_name, client, db, index_db, fields): 244 | self.model_name = model_name 245 | self.client = client 246 | self.db = db 247 | self.index_db = index_db 248 | self.fields = fields 249 | 250 | self.name = model_name.lower() 251 | self.sequence = 'id_seq:%s' % self.name 252 | 253 | self.defaults = {} 254 | self.defaults_callable = {} 255 | 256 | def set_client(self, client): 257 | self.client = client 258 | 259 | def prepared(self): 260 | self.sorted_fields = sorted( 261 | [field for field in self.fields.values()], 262 | key=lambda field: field._order) 263 | 264 | # Populate index attributes. 265 | self.indexed_fields = set() 266 | self.indexed_field_objects = [] 267 | self.indexes = {} 268 | for field in self.sorted_fields: 269 | if field.index: 270 | self.indexed_fields.add(field.name) 271 | self.indexed_field_objects.append(field) 272 | self.indexes[field.name] = Index(self.client, self.index_db, 273 | field) 274 | 275 | for field in self.sorted_fields: 276 | if callable(field.default): 277 | self.defaults_callable[field.name] = field.default 278 | elif field.default: 279 | self.defaults[field.name] = field.default 280 | 281 | def next_id(self): 282 | return self.client.incr(self.sequence) 283 | 284 | def get_instance_key(self, instance_id): 285 | return '%s:%s' % (self.name, instance_id) 286 | 287 | 288 | def with_metaclass(meta, base=object): 289 | return meta('newbase', (base,), {}) 290 | 291 | 292 | class Model(with_metaclass(DeclarativeMeta)): 293 | def __init__(self, **kwargs): 294 | self._data = self._meta.defaults.copy() 295 | for key, value in self._meta.defaults_callable.items(): 296 | self._data[key] = value() 297 | self._data.update(kwargs) 298 | 299 | @classmethod 300 | def create(cls, **kwargs): 301 | instance = cls(**kwargs) 302 | instance.save() 303 | return instance 304 | 305 | @classmethod 306 | def load(cls, primary_key): 307 | data = cls._read_model_data(primary_key) 308 | if data is None: 309 | raise KeyError('%s with id=%s not found' % 310 | (cls._meta.name, primary_key)) 311 | return cls(**data) 312 | 313 | def save(self): 314 | if self.id and self._meta.indexes: 315 | original_data = self._read_model_data(primary_key) 316 | else: 317 | original_data = None 318 | 319 | # Save the actual model data. 320 | self._save_model_data() 321 | 322 | # Update any secondary indexes. 323 | self._update_indexes(original_data) 324 | 325 | def _save_model_data(self): 326 | # Generate the next ID in sequence if no ID is set. 327 | if not self.id: 328 | self.id = self._meta.next_id() 329 | 330 | # Retrieve the primary key identifying this model instance. 331 | key = self._meta.get_instance_key(self.id) 332 | 333 | # Prepare the data for storage. Some Python data-types, e.g. datetimes, 334 | # cannot be serialized natively by the greendb protocol, so we 335 | # serialize them before writing. 336 | accum = {} 337 | for field in self._meta.sorted_fields: 338 | value = self._data.get(field.name) 339 | if value is not None: 340 | accum[field.name] = field.serialize(value) 341 | 342 | self._meta.client.set(key, accum) # Alt: .set(key, self._data). 343 | 344 | def _update_indexes(self, original_data=None): 345 | primary_key = self.id 346 | 347 | # Store current model data in the index. 348 | for field_name, index in self._meta.indexes.items(): 349 | field = self._meta.fields[field_name] 350 | new_value = self._data.get(field_name) 351 | if original_data is not None and field_name in original_data: 352 | old_value = original_data[field_name] 353 | if old_value != new_value: 354 | index.delete(old_value, primary_key) 355 | 356 | if new_value is not None: 357 | index.store(new_value, primary_key) 358 | 359 | @classmethod 360 | def _read_model_data(cls, primary_key): 361 | key = cls._meta.get_instance_key(primary_key) 362 | data = cls._meta.client.get(key) 363 | if data is not None: 364 | return cls._deserialize_raw_data(data) 365 | 366 | @classmethod 367 | def _deserialize_raw_data(cls, data): 368 | accum = {} 369 | for key, value in data.items(): 370 | field = cls._meta.fields.get(key) 371 | if field is not None: 372 | accum[key] = field.deserialize(value) 373 | else: 374 | accum[key] = value 375 | return accum 376 | 377 | def delete(self): 378 | primary_key = self.id 379 | 380 | # Load up original data if we have indexes to clean up. 381 | if self._meta.indexes: 382 | data = self._read_model_data(primary_key) 383 | else: 384 | data = None 385 | 386 | # Delete the model data. 387 | key = self._meta.get_instance_key(primary_key) 388 | self._meta.client.delete(key) 389 | 390 | # Clear out the indexes. 391 | if data is not None: 392 | for field_name, index in self._meta.indexes.items(): 393 | value = data.get(field_name) or self._data.get(field_name) 394 | if value is not None: 395 | field = self._meta.fields[field_name] 396 | index.delete(value, self.id) 397 | 398 | @classmethod 399 | def get(cls, expr): 400 | results = cls.query(expr) 401 | if results: 402 | return results[0] 403 | 404 | @classmethod 405 | def all(cls, count=None): 406 | accum = [] 407 | prefix = encode(cls._meta.name) 408 | start = prefix + b':\x00' 409 | stop = prefix + b':\xff' 410 | for k, v in cls._meta.client.getrange(start, stop, count): 411 | data = cls._deserialize_raw_data(v) 412 | accum.append(cls(**data)) 413 | return accum 414 | 415 | @classmethod 416 | def query(cls, expr): 417 | def dfs(expr): 418 | lhs = expr.lhs 419 | rhs = expr.rhs 420 | if isinstance(lhs, Expression): 421 | lhs = dfs(lhs) 422 | if isinstance(rhs, Expression): 423 | rhs = dfs(rhs) 424 | 425 | if isinstance(lhs, Field): 426 | index = cls._meta.indexes[lhs.name] 427 | return index.query(rhs, expr.op) 428 | elif expr.op == 'AND': 429 | return set(lhs) & set(rhs) 430 | elif expr.op == 'OR': 431 | return set(lhs) | set(rhs) 432 | else: 433 | raise ValueError('Unable to execute query, unexpected type.') 434 | 435 | # Collect raw list of integer IDs. 436 | id_list = dfs(expr) 437 | 438 | # Bulk-load the model data, creating a mapping of model key -> int id. 439 | keys = [encode(cls._meta.get_instance_key(pk)) for pk in id_list] 440 | key_to_data = cls._meta.client.mget(keys) 441 | 442 | deserialize = cls._deserialize_raw_data 443 | return [cls(**deserialize(key_to_data[key])) for key in keys] 444 | 445 | 446 | class Index(object): 447 | def __init__(self, client, index_db, field): 448 | self.client = client 449 | self.db = index_db 450 | self.field = field 451 | self.key = 'idx:%s:%s' % (field.model._meta.name, field.name) 452 | 453 | # Obtain references to serialization routine for query values, e.g. 454 | # datetimes are stored as 64-bit unsigned integer microseconds. So when 455 | # we receive a datetime, we need to convert the incoming user value. 456 | self.convert = field.index_value 457 | 458 | def get_value(self, value, primary_key): 459 | return b'%s\x00%s' % (self.convert(value), encode(primary_key)) 460 | 461 | def store(self, value, primary_key): 462 | value = self.get_value(value, primary_key) 463 | return self.client.setdupraw(self.key, value, db=self.db) 464 | 465 | def delete(self, value, primary_key): 466 | value = self.get_value(value, primary_key) 467 | return self.client.deletedupraw(self.key, value, db=self.db) 468 | 469 | def _range_query(self, start=None, stop=None): 470 | return self.client.getrangedupraw(self.key, start, stop, db=self.db) 471 | 472 | def get_range_values(self, start=None, stop=None): 473 | accum = [] 474 | for raw_value in self._range_query(start, stop): 475 | delim_idx = raw_value.rfind(b'\x00') 476 | accum.append(int(raw_value[delim_idx + 1:].decode('ascii'))) 477 | return accum 478 | 479 | def query(self, value, operation): 480 | if operation == '=': 481 | bval = self.convert(value) 482 | return self.get_range_values(bval + b'\x00', bval + b'\x00\xff') 483 | elif operation in ('<', '<='): 484 | bval = self.convert(value) 485 | stop = bval + (b'\x00\x00' if operation == '<' else b'\x00\xff') 486 | return self.get_range_values(None, stop) 487 | elif operation in ('>', '>='): 488 | bval = self.convert(value) 489 | start = bval + (b'\x00\xff' if operation == '>' else b'\x00\x00') 490 | return self.get_range_values(start) 491 | elif operation == 'between': 492 | sstart, sstop, start_incl, stop_incl = value 493 | start = self.convert(sstart) 494 | stop = self.convert(sstop) 495 | start = start + (b'\x00\x00' if start_incl else b'\x00\xff') 496 | stop = stop + (b'\x00\xff' if stop_incl else b'\x00\x00') 497 | return self.get_range_values(start, stop) 498 | elif operation == '!=': 499 | bval = self.convert(value) 500 | accum = [] 501 | for raw_value in self._range_query(): 502 | idx_value, pk = raw_value.rsplit(b'\x00', 1) 503 | if bval != idx_value: 504 | accum.append(int(pk.decode('ascii'))) 505 | return accum 506 | elif operation == 'startswith': 507 | bval = self.convert(value) 508 | return self.get_range_values(bval, bval + b'\xff') 509 | else: 510 | raise ValueError('unrecognized operation: "%s"' % operation) 511 | -------------------------------------------------------------------------------- /tests.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from __future__ import unicode_literals # Required for 2.x compatability. 4 | 5 | import datetime 6 | import logging 7 | import os 8 | import shutil 9 | import sys 10 | import tempfile 11 | import unittest 12 | 13 | import gevent 14 | 15 | from greendb import Client 16 | from greendb import CommandError 17 | from greendb import Server 18 | from greendb import logger 19 | from greenquery import * 20 | 21 | 22 | TEST_HOST = '127.0.0.1' 23 | TEST_PORT = 31327 24 | 25 | 26 | def run_server(): 27 | logger.addHandler(logging.StreamHandler()) 28 | logger.setLevel(logging.ERROR) 29 | tmp_dir = tempfile.mkdtemp(suffix='greendb') 30 | data_dir = os.path.join(tmp_dir, 'data') 31 | server = Server(host=TEST_HOST, port=TEST_PORT, path=data_dir, 32 | max_dbs=4, dupsort=[3]) 33 | def run(): 34 | try: 35 | server.run() 36 | finally: 37 | if os.path.exists(tmp_dir): 38 | shutil.rmtree(tmp_dir) 39 | t = gevent.spawn(run) 40 | gevent.sleep(0) 41 | return t, server, tmp_dir 42 | 43 | 44 | class BaseTestCase(unittest.TestCase): 45 | @classmethod 46 | def setUpClass(cls): 47 | cls.c = Client(host=TEST_HOST, port=TEST_PORT, pool=False) 48 | 49 | @classmethod 50 | def tearDownClass(cls): 51 | cls.c.quit() 52 | 53 | def tearDown(self): 54 | super(BaseTestCase, self).tearDown() 55 | self.c.flushall() 56 | 57 | 58 | class TestBasicOperations(BaseTestCase): 59 | def setUp(self): 60 | super(TestBasicOperations, self).setUp() 61 | # By default we will use the 0th database. 62 | self.c.use(0) 63 | 64 | def test_ping(self): 65 | self.assertEqual(self.c.ping(), b'pong') 66 | 67 | def test_identity(self): 68 | test_data = ( 69 | b'foo', 70 | b'\xff\x00\xff', 71 | 'foo', 72 | '\u2012\u2013'.encode('utf8'), 73 | 0, 74 | 1337, 75 | -1, 76 | 3.14159, 77 | [b'foo', b'\xff\x00\xff', 31337, [b'bar']], 78 | {b'k1': b'v1', b'k2': 2, b'k3': {b'x3': b'y3'}}, 79 | True, 80 | False, 81 | None, 82 | b'', 83 | b'a' * (1024 * 1024), # 1MB value. 84 | ) 85 | 86 | for test in test_data: 87 | self.assertEqual(self.c.set('key', test), 1) 88 | self.assertEqual(self.c.get('key'), test) 89 | 90 | def test_dict_interface(self): 91 | self.c['k1'] = 'v1' 92 | self.assertEqual(self.c['k1'], 'v1') 93 | self.assertFalse('k2' in self.c) 94 | self.c['k2'] = 'v2' 95 | self.assertTrue('k2' in self.c) 96 | self.assertEqual(list(self.c), [b'k1', b'k2']) 97 | self.assertEqual(len(self.c), 2) 98 | del self.c['k2'] 99 | self.assertFalse('k2' in self.c) 100 | self.assertEqual(len(self.c), 1) 101 | 102 | self.c.update(k2='v2', k3='v3') 103 | self.assertEqual(self.c[:'k2'], [[b'k1', 'v1'], [b'k2', 'v2']]) 104 | self.assertEqual(self.c['k1':'k2x'], [[b'k1', 'v1'], [b'k2', 'v2']]) 105 | self.assertEqual(self.c['k1x':], [[b'k2', 'v2'], [b'k3', 'v3']]) 106 | self.assertEqual(self.c['k1x'::1], [[b'k2', 'v2']]) 107 | 108 | del self.c['k1x'::1] 109 | self.assertEqual(list(self.c), [b'k1', b'k3']) 110 | del self.c['k1', 'k2', 'k3', 'k4'] 111 | self.assertEqual(list(self.c), []) 112 | 113 | def test_env_info(self): 114 | info = self.c.envinfo() 115 | self.assertEqual(info['clients'], 1) 116 | self.assertEqual(info['host'], '127.0.0.1') 117 | self.assertEqual(info['port'], TEST_PORT) 118 | sc = info['storage'] 119 | self.assertEqual(sc['dupsort'], [3]) 120 | self.assertEqual(sc['lock'], 1) 121 | self.assertEqual(sc['max_dbs'], 4) 122 | self.assertEqual(sc['sync'], 0) 123 | 124 | rc = self.c.reader_check() 125 | self.assertEqual(rc, 0) 126 | 127 | def test_decode_keys(self): 128 | c = Client(host=TEST_HOST, port=TEST_PORT, decode_keys=True) 129 | c.mset({'k1': b'v1', 'k2': {'x1': 'y1', 'x2': {'y2': b'z2'}}}) 130 | c.mset({'i1': {1: 2}, 'k3': ['foo', {b'a': 'b'}, b'bar']}) 131 | 132 | # Dicts are decoded recursively. 133 | self.assertEqual(c.mget(['k1', 'k2']), { 134 | 'k1': b'v1', 135 | 'k2': {'x1': 'y1', 'x2': {'y2': b'z2'}}}) 136 | 137 | # Nested dicts are not decoded. 138 | self.assertEqual(c.get('k3'), ['foo', {b'a': 'b'}, b'bar']) 139 | 140 | # ints and other values are modified, as keys must be strings. 141 | self.assertEqual(c.get('i1'), {'1': 2}) 142 | c.close() 143 | 144 | def test_crud(self): 145 | # Setting a key/value returns number of keys set. 146 | self.assertEqual(self.c.set('k1', b'v1'), 1) 147 | 148 | # We can verify the key exists, and that we can retrieve our value. 149 | self.assertTrue(self.c.exists('k1')) 150 | self.assertEqual(self.c.get('k1'), b'v1') 151 | self.assertEqual(self.c.count(), 1) 152 | 153 | # Deleting returns the number of keys deleted. 154 | self.assertEqual(self.c.delete('k1'), 1) 155 | self.assertFalse(self.c.exists('k1')) 156 | self.assertEqual(self.c.count(), 0) 157 | 158 | # Subsequent call to delete returns 0. 159 | self.assertEqual(self.c.delete('k1'), 0) 160 | 161 | # Getting a nonexistant key returns None. 162 | self.assertTrue(self.c.get('k1') is None) 163 | 164 | # Let's set and then update a key. 165 | self.assertEqual(self.c.set('key', 'ccc'), 1) 166 | self.assertEqual(self.c.get('key'), 'ccc') 167 | 168 | # dupsort is disabled for this database, so the value is replaced. 169 | self.assertEqual(self.c.set('key', 'ddd'), 1) 170 | self.assertEqual(self.c.get('key'), 'ddd') 171 | self.assertEqual(self.c.set('key', 'bbb'), 1) 172 | self.assertEqual(self.c.get('key'), 'bbb') 173 | self.assertEqual(self.c.getdup('key'), ['bbb']) 174 | self.assertRaises(CommandError, self.c.dupcount, 'key') 175 | 176 | # We can set to the same value and the db returns 1. When dupsort is 177 | # enabled, this returns 0. 178 | self.assertEqual(self.c.set('key', 'bbb'), 1) 179 | self.assertEqual(self.c.getdup('key'), ['bbb']) 180 | 181 | self.assertEqual(self.c.pop('key'), 'bbb') 182 | self.assertTrue(self.c.pop('key') is None) 183 | 184 | self.assertTrue(self.c.replace('key', 'aaa') is None) 185 | self.assertEqual(self.c.get('key'), 'aaa') 186 | self.assertEqual(self.c.replace('key', 'bbb'), 'aaa') 187 | self.assertEqual(self.c.get('key'), 'bbb') 188 | self.assertTrue(self.c.replace('keyz', 'xxx') is None) 189 | self.assertEqual(self.c.replace('keyz', 'yyy'), 'xxx') 190 | 191 | self.assertEqual(self.c.setnx('key', 'ccc'), 0) 192 | self.assertEqual(self.c.setnx('key2', 'xxx'), 1) 193 | self.assertEqual(self.c.get('key'), 'bbb') 194 | self.assertEqual(self.c.get('key2'), 'xxx') 195 | 196 | def test_get_set_raw(self): 197 | self.assertEqual(self.c.setraw('key', b'abc'), 1) 198 | self.assertEqual(self.c.getraw('key'), b'abc') 199 | self.assertTrue(self.c.getraw('x') is None) 200 | 201 | # Unicode values are encoded as UTF8 automatically. 202 | self.assertEqual(self.c.setraw('key', u'\u2021'), 1) 203 | self.assertEqual(self.c.getraw('key'), b'\xe2\x80\xa1') 204 | 205 | data = {'k1': b'v1', 'k2': u'\u2021'} 206 | self.assertEqual(self.c.msetraw(data), 2) 207 | self.assertEqual(self.c.mgetraw(['k1', 'k2', 'kx']), { 208 | b'k1': b'v1', 209 | b'k2': b'\xe2\x80\xa1'}) 210 | 211 | def test_crud_dupsort(self): 212 | # Use the DB with dupsort enabled. 213 | self.c.use(3) 214 | 215 | # Setting a key/value returns number of keys set. 216 | self.assertEqual(self.c.set('k1', b'v1'), 1) 217 | 218 | # We can verify the key exists, and that we can retrieve our value. 219 | self.assertTrue(self.c.exists('k1')) 220 | self.assertEqual(self.c.get('k1'), b'v1') 221 | self.assertEqual(self.c.count(), 1) 222 | 223 | # We can add another value with dupsort enabled. 224 | self.assertEqual(self.c.set('k1', b'v1-x'), 1) 225 | 226 | # Deleting returns the number of keys deleted. 227 | self.assertEqual(self.c.delete('k1'), 1) 228 | self.assertFalse(self.c.exists('k1')) 229 | self.assertEqual(self.c.count(), 0) 230 | 231 | # Subsequent call to delete returns 0. 232 | self.assertEqual(self.c.delete('k1'), 0) 233 | 234 | # Set multiple values and then use deletedup to verify the old values 235 | # are preserved. 236 | self.c.set('k1', b'v1-a') 237 | self.c.set('k1', b'v1-b') 238 | self.c.set('k1', b'v1-c') 239 | self.assertEqual(self.c.deletedup('k1', b'v1-b'), 1) 240 | self.assertEqual(self.c.deletedup('k1', b'v1-x'), 0) 241 | self.assertTrue(self.c.exists('k1')) 242 | self.assertEqual(self.c.count(), 2) 243 | self.assertEqual(self.c.getdup('k1'), [b'v1-a', b'v1-c']) 244 | self.assertEqual(self.c.delete('k1'), 1) 245 | self.assertEqual(self.c.count(), 0) 246 | 247 | # Getting a nonexistant key returns None. 248 | self.assertTrue(self.c.get('k1') is None) 249 | 250 | # Let's set and then update a key. 251 | self.assertEqual(self.c.set('key', b'ccc'), 1) 252 | self.assertEqual(self.c.get('key'), b'ccc') 253 | 254 | # Because our databases use dupsort and multi-value, we actually get 255 | # "ccc" here, because "ccc" sorts before "ddd". 256 | self.assertEqual(self.c.set('key', b'ddd'), 1) 257 | self.assertEqual(self.c.get('key'), b'ccc') 258 | self.assertEqual(self.c.set('key', b'bbb'), 1) 259 | self.assertEqual(self.c.get('key'), b'bbb') 260 | self.assertEqual(self.c.getdup('key'), [b'bbb', b'ccc', b'ddd']) 261 | self.assertEqual(self.c.dupcount('key'), 3) 262 | 263 | # However we can't set the same exact key/data, except via setdup: 264 | self.assertEqual(self.c.set('key', b'ccc'), 0) 265 | self.assertEqual(self.c.set('key', b'bbb'), 0) 266 | self.assertEqual(self.c.getdup('key'), [b'bbb', b'ccc', b'ddd']) 267 | self.assertEqual(self.c.dupcount('key'), 3) 268 | 269 | self.assertEqual(self.c.pop('key'), b'bbb') 270 | self.assertEqual(self.c.pop('key'), b'ccc') 271 | self.assertEqual(self.c.pop('key'), b'ddd') 272 | self.assertTrue(self.c.pop('key') is None) 273 | 274 | self.assertTrue(self.c.replace('key', b'aaa') is None) 275 | self.assertEqual(self.c.get('key'), b'aaa') 276 | self.assertEqual(self.c.replace('key', b'bbb'), b'aaa') 277 | self.assertEqual(self.c.get('key'), b'bbb') 278 | self.assertEqual(self.c.getdup('key'), [b'bbb']) 279 | self.assertEqual(self.c.dupcount('key'), 1) 280 | self.assertTrue(self.c.replace('keyz', b'xxx') is None) 281 | self.assertEqual(self.c.replace('keyz', b'yyy'), b'xxx') 282 | self.assertEqual(self.c.getdup('keyz'), [b'yyy']) 283 | 284 | self.assertEqual(self.c.setnx('key', b'ccc'), 0) 285 | self.assertEqual(self.c.setnx('key2', b'xxx'), 1) 286 | self.assertEqual(self.c.get('key'), b'bbb') 287 | self.assertEqual(self.c.getdup('key'), [b'bbb']) 288 | self.assertEqual(self.c.get('key2'), b'xxx') 289 | 290 | def test_bulk_operations(self): 291 | self.assertEqual(self.c.mset({'k1': b'v1', 'k2': b'v2', 'k3': b'v3'}), 292 | 3) 293 | self.assertEqual(self.c.mset({ 294 | 'k1': b'v1-x', 295 | 'k2': b'v2', 296 | 'k4': b'v4', 297 | 'k5': b'v5'}), 4) 298 | 299 | self.assertEqual(self.c.mget(['k1', 'k3', 'k5']), { 300 | b'k1': b'v1-x', 301 | b'k3': b'v3', 302 | b'k5': b'v5'}) 303 | self.assertEqual(self.c.mget(['k0', 'k2', 'kx']), {b'k2': b'v2'}) 304 | self.assertEqual(self.c.mget(['kx', 'ky', 'kz']), {}) 305 | 306 | # Bulk delete returns number actually deleted. 307 | self.assertEqual(self.c.mdelete(['k2', 'k5', 'kx']), 2) 308 | self.assertEqual(self.c.mdelete(['k2', 'k5', 'kx']), 0) 309 | 310 | # Bulk pop. 311 | self.assertEqual(self.c.mpop(['k3', 'k4', 'kz']), 312 | {b'k3': b'v3', b'k4': b'v4'}) 313 | self.assertEqual(self.c.mpop(['k3', 'k4', 'kz']), {}) 314 | 315 | # Bulk replace, returns the previous value (if it exists). 316 | res = self.c.mreplace({'k3': b'v3-x', 'k4': b'v4-x', 'k1': b'v1-z'}) 317 | self.assertEqual(res, {b'k1': b'v1-x'}) 318 | 319 | res = self.c.mreplace({'k3': b'v3-y', 'k4': b'v4-y'}) 320 | self.assertEqual(res, {b'k3': b'v3-x', b'k4': b'v4-x'}) 321 | 322 | # Set NX. 323 | self.assertEqual(self.c.msetnx({'k2': b'v2', 'k4': b'v4'}), 1) 324 | self.assertEqual(self.c.msetnx({'k2': b'v2', 'k3': b'v3'}), 0) 325 | self.assertEqual(self.c.mget(['k1', 'k2', 'k3', 'k4', 'k5']), { 326 | b'k1': b'v1-z', 327 | b'k2': b'v2', 328 | b'k3': b'v3-y', 329 | b'k4': b'v4-y'}) 330 | 331 | def test_bulk_operations_dupsort(self): 332 | # We need to specify the DB that supports dupsort. 333 | self.c.use(3) 334 | 335 | self.assertEqual(self.c.mset({'k1': b'v1', 'k2': b'v2', 'k3': b'v3'}), 336 | 3) 337 | self.assertEqual(self.c.mset({ 338 | 'k1': b'v1-x', 339 | 'k2': b'v2', 340 | 'k4': b'v4', 341 | 'k5': b'v5'}), 3) # Only k1, k4 and k5 are counted: k2 is same. 342 | self.assertEqual(self.c.mset({'k5': b'v5-x'}), 1) 343 | 344 | self.assertEqual(self.c.mget(['k1', 'k3', 'k5']), { 345 | b'k1': b'v1', # v1-x sorts *after* v1. 346 | b'k3': b'v3', 347 | b'k5': b'v5'}) 348 | self.assertEqual(self.c.mget(['k0', 'k2', 'kx']), {b'k2': b'v2'}) 349 | self.assertEqual(self.c.mget(['kx', 'ky', 'kz']), {}) 350 | 351 | # We can get all dupes using mgetdup. 352 | self.assertEqual(self.c.mgetdup(['k1', 'k3', 'k5', 'kx']), { 353 | b'k1': [b'v1', b'v1-x'], 354 | b'k3': [b'v3'], 355 | b'k5': [b'v5', b'v5-x']}) 356 | 357 | # Bulk delete returns number actually deleted. Event though k5 has two 358 | # values, it is only counted once here. 359 | self.assertEqual(self.c.mdelete(['k2', 'k5', 'kx']), 2) 360 | self.assertEqual(self.c.mdelete(['k2', 'k5', 'kx']), 0) 361 | 362 | # Bulk pop. 363 | self.assertEqual(self.c.mpop(['k3', 'k4', 'kz']), 364 | {b'k3': b'v3', b'k4': b'v4'}) 365 | self.assertEqual(self.c.mpop(['k3', 'k4', 'kz']), {}) 366 | 367 | # Bulk pop with duplicates. Only the first value is popped off. 368 | self.c.mset({'k5': b'v5-a'}) 369 | self.c.mset({'k5': b'v5-b'}) 370 | self.assertEqual(self.c.mpop(['k1', 'k5']), 371 | {b'k1': b'v1', b'k5': b'v5-a'}) 372 | self.assertEqual(self.c.mpop(['k1', 'k5']), 373 | {b'k1': b'v1-x', b'k5': b'v5-b'}) 374 | self.assertEqual(self.c.mpop(['k1', 'k5']), {}) 375 | 376 | # Restore the values to k1. 377 | self.c.set('k1', b'v1-x') 378 | 379 | # Bulk replace, returns the previous value (if it exists). 380 | res = self.c.mreplace({'k3': b'v3-x', 'k4': b'v4-x', 'k1': b'v1-z'}) 381 | self.assertEqual(res, {b'k1': b'v1-x'}) 382 | 383 | # Although we overwrote k1, the original value is not preserved! 384 | self.assertEqual(self.c.getdup('k1'), [b'v1-z']) 385 | self.assertEqual(self.c.dupcount('k1'), 1) 386 | 387 | # Verify the original values are returned. 388 | res = self.c.mreplace({'k3': b'v3-y', 'k4': b'v4-y'}) 389 | self.assertEqual(res, {b'k3': b'v3-x', b'k4': b'v4-x'}) 390 | 391 | # Set NX works the same as non-dupsort. 392 | self.assertEqual(self.c.msetnx({'k2': b'v2', 'k4': b'v4'}), 1) 393 | self.assertEqual(self.c.msetnx({'k2': b'v2', 'k3': b'v3'}), 0) 394 | self.assertEqual(self.c.mget(['k1', 'k2', 'k3', 'k4', 'k5']), { 395 | b'k1': b'v1-z', 396 | b'k2': b'v2', 397 | b'k3': b'v3-y', 398 | b'k4': b'v4-y'}) 399 | 400 | def test_match_prefix(self): 401 | indices = list(range(0, 100, 5)) 402 | 403 | # k000, k005, k010, k015 ... k090, k095. 404 | self.c.mset({'k%03d' % i: b'v%d' % i for i in indices}) 405 | 406 | def assertPrefix(prefix, indices, count=None): 407 | r = self.c.prefix(prefix, count) 408 | if count is not None: 409 | indices = indices[:count] 410 | self.assertEqual(r, [[b'k%03d' % i, b'v%d' % i] for i in indices]) 411 | 412 | for count in (None, 1, 2, 100): 413 | # Prefix scan works as expected. 414 | assertPrefix('k01', [10, 15]) 415 | assertPrefix('k00', [0, 5]) 416 | assertPrefix('k020', [20]) 417 | assertPrefix('k025', [25]) 418 | 419 | # No keys match these prefixes. 420 | assertPrefix('k021', []) 421 | assertPrefix('k001', []) 422 | assertPrefix('k1', []) 423 | assertPrefix('j', []) 424 | 425 | # These prefixes match all keys. 426 | assertPrefix('', indices) 427 | assertPrefix('k', indices) 428 | assertPrefix('k0', indices) 429 | 430 | def test_match_prefix_dupsort(self): 431 | self.c.use(3) 432 | self.c.mset({'aaa': b'a1', 'aab': b'b1', 'aac': b'c1'}) 433 | self.c.mset({'aaa': b'a2', 'aac': b'c2'}) 434 | items = [ 435 | [b'aaa', b'a1'], [b'aaa', b'a2'], 436 | [b'aab', b'b1'], 437 | [b'aac', b'c1'], [b'aac', b'c2']] 438 | 439 | self.assertEqual(self.c.prefix('aa'), items) 440 | self.assertEqual(self.c.prefix('aa', 3), items[:3]) 441 | self.assertEqual(self.c.prefix('aaa', 3), items[:2]) 442 | self.assertEqual(self.c.prefix('aaa', 1), items[:1]) 443 | 444 | def test_deleterange(self): 445 | self.c.mset(dict(('k%d' % i, b'v%d' % i) for i in range(20))) 446 | self.assertEqual(self.c.deleterange('k18', 'k3'), 4) # 18, 19, 2, 3. 447 | self.assertEqual(self.c.deleterange('k17x', 'k3x'), 0) 448 | 449 | self.assertEqual(self.c.deleterange('k10', 'k15', 2), 2) # 10, 11. 450 | self.assertEqual(self.c.deleterange('k10', 'k15', 2), 2) # 12, 13. 451 | r = self.c.getrange('k0', 'k15') 452 | self.assertEqual([k for k, _ in r], [b'k0', b'k1', b'k14', b'k15']) 453 | self.assertEqual(self.c.deleterange('k0', 'k15'), 4) 454 | 455 | self.assertEqual(self.c.deleterange(None, 'k0'), 0) 456 | self.assertEqual(self.c.deleterange('k9z', None), 0) 457 | self.assertEqual(self.c.deleterange('a0', 'a9'), 0) 458 | self.assertEqual(self.c.deleterange('z0', 'z9'), 0) 459 | self.assertEqual(self.c.deleterange(), 8) 460 | 461 | def test_deleterange_dupsort(self): 462 | self.assertEqual(self.c.use(3), 3) # Use dupsort db. 463 | self.c.mset({'k0': b'v0', 'k1': b'v1', 'k10': b'v10', 'k11': b'v11'}) 464 | self.c.mset({'k0': b'v0-x', 'k10': b'v10-x', 'k11': b'v11-x'}) 465 | self.c.mset({'k0': b'v0-y', 'k11': b'v11-y'}) 466 | 467 | # k0 [v0,v0-x,v0-y], k1 [v1], k10 [v10,v10-x], k11 [v11,v11-x,v11-y] 468 | self.assertEqual(self.c.count(), 9) 469 | self.assertEqual(self.c.deleterange('k1', 'k9', 2), 2) 470 | self.assertEqual([v for _, v in self.c.getrange('k1', 'k9')], 471 | [b'v10-x', b'v11', b'v11-x', b'v11-y']) 472 | 473 | self.assertEqual(self.c.deleterange('k1', 'k9', 3), 3) 474 | self.assertEqual([v for _, v in self.c.getrange()], 475 | [b'v0', b'v0-x', b'v0-y', b'v11-y']) 476 | 477 | self.assertEqual(self.c.deleterange(None, None, 2), 2) 478 | self.assertEqual([v for _, v in self.c.getrange()], 479 | [b'v0-y', b'v11-y']) 480 | 481 | self.assertEqual(self.c.deleterange(None, None, 4), 2) 482 | self.assertEqual(self.c.count(), 0) 483 | 484 | def test_getrange_keys_values(self): 485 | self.c.mset(dict(('k%d' % i, b'v%d' % i) for i in range(20))) 486 | sorted_values = list(map(int, sorted(map(str, range(20))))) 487 | 488 | def assertRange(start, end, indices, count=None): 489 | if count is not None: 490 | indices = indices[:count] 491 | 492 | # Test KEYS command. 493 | res = self.c.keys(start, end, count) 494 | self.assertEqual(res, [b'k%d' % i for i in indices]) 495 | 496 | # Test VALUES command. 497 | res = self.c.values(start, end, count) 498 | self.assertEqual(res, [b'v%d' % i for i in indices]) 499 | 500 | # Test GETRANGE command. 501 | res = self.c.getrange(start, end, count) 502 | self.assertEqual(res, [[b'k%d' % i, b'v%d' % i] for i in indices]) 503 | 504 | for count in (None, 1, 2, 100): 505 | assertRange('k3', 'k6', [3, 4, 5, 6], count) 506 | assertRange('k18', 'k3', [18, 19, 2, 3], count) 507 | assertRange('k3x', 'k6x', [4, 5, 6], count) 508 | assertRange('k0', 'k12', [0, 1, 10, 11, 12], count) 509 | assertRange('k01', 'k121', [1, 10, 11, 12], count) 510 | 511 | # Test boundaries. 512 | assertRange(None, None, sorted_values, count) 513 | assertRange(None, 'kz', sorted_values, count) 514 | assertRange('a0', None, sorted_values, count) 515 | assertRange('k0', None, sorted_values, count) 516 | assertRange('k0', 'k9', sorted_values, count) 517 | 518 | # Test out-of-bounds. 519 | assertRange(None, 'a0', [], count) 520 | assertRange('z0', None, [], count) 521 | assertRange('a0', 'a99', [], count) 522 | assertRange('z0', 'z99', [], count) 523 | 524 | def test_getrange_keys_values_dupsort(self): 525 | self.c.use(3) 526 | nums = [0, 1, 10, 2, 3, 4] 527 | self.c.mset(dict(('k%d' % i, b'v%d' % i) for i in nums)) 528 | self.c.mset(dict(('k%d' % i, b'v%d-x' % i) for i in nums if not i % 2)) 529 | 530 | def assertRange(start, end, indices, count=None): 531 | accum = [] 532 | for i in indices: 533 | accum.append([b'k%d' % i, b'v%d' % i]) 534 | if i % 2 == 0: 535 | accum.append([b'k%d' % i, b'v%d-x' % i]) 536 | if count is not None: 537 | accum = accum[:count] 538 | 539 | res = self.c.getrange(start, end, count) 540 | self.assertEqual(res, accum) 541 | 542 | res = self.c.keys(start, end, count) 543 | self.assertEqual(res, [k for k, _ in accum]) 544 | 545 | res = self.c.values(start, end, count) 546 | self.assertEqual(res, [v for _, v in accum]) 547 | 548 | for count in (None, 1, 2, 100): 549 | assertRange('k2', 'k4', [2, 3, 4], count) 550 | assertRange('k1', 'k3', [1, 10, 2, 3], count) 551 | assertRange('k2x', 'k4x', [3, 4], count) 552 | assertRange('k0', 'k12', [0, 1, 10], count) 553 | assertRange('k01', 'k101', [1, 10], count) 554 | 555 | # Test boundaries. 556 | assertRange(None, None, nums, count) 557 | assertRange(None, 'kz', nums, count) 558 | assertRange('a0', None, nums, count) 559 | assertRange('k0', None, nums, count) 560 | assertRange('k0', 'k9', nums, count) 561 | 562 | # Test out-of-bounds. 563 | assertRange(None, 'a0', [], count) 564 | assertRange('z0', None, [], count) 565 | assertRange('a0', 'a99', [], count) 566 | assertRange('z0', 'z99', [], count) 567 | 568 | def test_incr_decr(self): 569 | self.assertEqual(self.c.incr('i0'), 1) 570 | self.assertEqual(self.c.incr('i0', 2), 3) 571 | self.assertEqual(self.c.incr('i1', 2), 2) 572 | self.assertEqual(self.c.incr('i1', 2.5), 4.5) 573 | self.assertEqual(self.c.decr('i1', 3.5), 1.0) 574 | self.assertEqual(self.c.decr('i0'), 2) 575 | self.assertEqual(self.c.decr('i2'), -1) 576 | self.assertEqual(self.c.get('i0'), 2) 577 | self.assertEqual(self.c.get('i1'), 1.) 578 | self.assertEqual(self.c.get('i2'), -1) 579 | self.c.set('i2', -2) 580 | self.assertEqual(self.c.decr('i2'), -3) 581 | self.c.set('i1', 2.0) 582 | self.assertEqual(self.c.incr('i1'), 3.) 583 | 584 | def test_cas(self): 585 | self.c.set('k1', b'v1') 586 | self.c.set('k2', b'v2') 587 | self.assertTrue(self.c.cas('k1', b'v1', b'v1-x')) 588 | self.assertFalse(self.c.cas('k1', b'v1', b'v1-y')) 589 | self.assertEqual(self.c.get('k1'), b'v1-x') 590 | 591 | self.assertFalse(self.c.cas('k2', b'v1-x', b'v1-z')) 592 | self.assertEqual(self.c.get('k2'), b'v2') 593 | self.assertTrue(self.c.cas('k2', b'v2', b'v2-x')) 594 | self.assertTrue(self.c.cas('k2', b'v2-x', b'v2-y')) 595 | self.assertFalse(self.c.cas('k2', b'v2', b'v2-z')) 596 | self.assertEqual(self.c.get('k2'), b'v2-y') 597 | 598 | self.assertFalse(self.c.cas('k3', b'', b'v3')) 599 | self.assertFalse(self.c.cas('k3', b'x3', b'y3')) 600 | self.assertTrue(self.c.cas('k3', None, b'v3')) 601 | self.assertEqual(self.c.get('k3'), b'v3') 602 | 603 | def test_cas_dupsort(self): 604 | self.c.use(3) 605 | 606 | self.c.set('k1', b'v1-b') 607 | self.c.set('k1', b'v1-a') 608 | self.c.set('k1', b'v1-c') 609 | self.c.set('k2', b'v2-a') 610 | 611 | # v1-a is the first element, so that is what we compare. 612 | self.assertFalse(self.c.cas('k1', b'v1-b', b'v1-x')) 613 | 614 | # We will successfully match value, so v1-a is deleted and v1-y is 615 | # added. The next compare will have to be against v1-b, however. 616 | self.assertTrue(self.c.cas('k1', b'v1-a', b'v1-y')) 617 | self.assertEqual(self.c.getdup('k1'), [b'v1-b', b'v1-c', b'v1-y']) 618 | self.assertFalse(self.c.cas('k1', b'v1-y', b'v1-z')) 619 | self.assertTrue(self.c.cas('k1', b'v1-b', b'v1-0')) 620 | self.assertEqual(self.c.getdup('k1'), [b'v1-0', b'v1-c', b'v1-y']) 621 | 622 | self.assertFalse(self.c.cas('k2', b'v1-0', b'v2-z')) 623 | self.assertEqual(self.c.getdup('k2'), [b'v2-a']) 624 | self.assertTrue(self.c.cas('k2', b'v2-a', b'v2-x')) 625 | self.assertTrue(self.c.cas('k2', b'v2-x', b'v2-y')) 626 | self.assertEqual(self.c.getdup('k2'), [b'v2-y']) 627 | 628 | self.assertFalse(self.c.cas('k3', b'', b'v3')) 629 | self.assertFalse(self.c.cas('k3', b'x3', b'y3')) 630 | self.assertTrue(self.c.cas('k3', None, b'v3')) 631 | self.assertEqual(self.c.getdup('k3'), [b'v3']) 632 | 633 | def test_length(self): 634 | self.c.set('k1', b'abc') 635 | self.c.set('k2', 'defghijkl') 636 | self.c.set('k3', b'') 637 | self.c.set('k4', [0, 1, 2, 3]) 638 | self.c.set('k5', {'x': 'y', 'z': 'w'}) 639 | self.assertEqual(self.c.length('k1'), 3) 640 | self.assertEqual(self.c.length('k2'), 9) 641 | self.assertEqual(self.c.length('k3'), 0) 642 | self.assertEqual(self.c.length('k4'), 4) 643 | self.assertEqual(self.c.length('k5'), 2) 644 | self.assertTrue(self.c.length('kx') is None) 645 | 646 | def test_raw_operations_dupsort(self): 647 | # We need to specify the DB that supports dupsort. 648 | self.c.use(3) 649 | 650 | data = ( 651 | # id, username, status 652 | (1, 'huey', 1), 653 | (2, 'zaizee', 1), 654 | (3, 'mickey', 0), 655 | (4, 'beanie', 1), 656 | (5, 'gracie', 0), 657 | (6, 'rocky', 0), 658 | (7, 'rocky-2', 0), 659 | (8, 'rocky-3', 0)) 660 | for pk, username, status in data: 661 | self.c.setdupraw('u:username', '%s %s' % (username, pk)) 662 | self.c.setdupraw('u:status', '%s %s' % (status, pk)) 663 | 664 | # We'll use chr(32) " " for the lower-bound, and chr(126) "~" for the 665 | # upper-bound. 666 | def assertR(start, stop, expected): 667 | res = self.c.getrangedupraw('u:username', start, stop) 668 | pks = [int(r.decode('utf8').rsplit(' ', 1)[1]) for r in res] 669 | self.assertEqual(pks, expected) 670 | 671 | # Verify inclusiveness of both boundaries in this scheme. 672 | assertR('zaizee ', 'zaizee ', []) 673 | assertR('zaizee ', 'zaizee 0', []) 674 | assertR('zaizee ', 'zaizee 2', [2]) 675 | assertR('zaizee ', 'zaizee ~', [2]) 676 | assertR('zaizee 0', 'zaizee ~', [2]) 677 | assertR('zaizee 2', 'zaizee ~', [2]) 678 | 679 | # Test equality comparison. 680 | assertR('zaizee ', 'zaizee ~', [2]) 681 | 682 | # Less-than, and less-than-or-equal. 683 | assertR(None, 'rocky ', [4, 5, 1, 3]) 684 | assertR(None, 'rocky ~', [4, 5, 1, 3, 6]) 685 | 686 | # Greater-than, and greater-than-or-equal. 687 | assertR('rocky ~', None, [7, 8, 2]) 688 | assertR('rocky ', None, [6, 7, 8, 2]) 689 | 690 | # Startswith (prefix search). 691 | assertR('rocky', 'rocky~', [6, 7, 8]) 692 | assertR('b', 'b~', [4]) 693 | 694 | # Between inclusive, between exclusive. 695 | assertR('huey ', 'rocky ~', [1, 3, 6]) 696 | assertR('huey ~', 'rocky ', [3]) 697 | 698 | # Delete a few items and verify deletedupraw works. 699 | self.c.deletedupraw('u:username', 'huey 1') 700 | self.c.deletedupraw('u:username', 'rocky-2 7') 701 | self.c.deletedupraw('u:username', 'rocky-3 x') # Does not match! 702 | assertR('gracie', 'zaizee ', [5, 3, 6, 8]) 703 | 704 | # Status test. 705 | def assertS(start, stop, expected): 706 | res = self.c.getrangedupraw('u:status', start, stop) 707 | pks = [int(r.decode('utf8').rsplit(' ', 1)[1]) for r in res] 708 | self.assertEqual(pks, expected) 709 | 710 | assertS('0 ', '0 ~', [3, 5, 6, 7, 8]) 711 | assertS('1 ', '1 ~', [1, 2, 4]) 712 | assertS('0 ', '1 ~', [3, 5, 6, 7, 8, 1, 2, 4]) 713 | 714 | def test_processing_instruction_use_db(self): 715 | # Verify we can specify the database for a one-off operation. 716 | for i in range(4): 717 | self.c.set('k1', 'v1-%s' % i, db=i) 718 | for i in range(4): 719 | self.assertEqual(self.c.get('k1', db=i), 'v1-%s' % i) 720 | 721 | # The default db is 0. 722 | self.assertEqual(self.c.get('k1'), 'v1-0') 723 | 724 | # We can explicitly switch the db. 725 | self.assertEqual(self.c.use(2), 2) 726 | self.assertEqual(self.c.get('k1'), 'v1-2') 727 | 728 | # We can specify the db for a one-off command. 729 | for i in range(4): 730 | self.assertEqual(self.c.get('k1', db=i), 'v1-%s' % i) 731 | 732 | # The value from our call to "use()" is preserved. 733 | self.assertEqual(self.c.get('k1'), 'v1-2') 734 | 735 | def test_sleep_attributes(self): 736 | resp, attrs = self.c._sleep(0.1) 737 | self.assertEqual(resp, 0.1) 738 | self.assertTrue('start' in attrs.data) 739 | self.assertTrue('stop' in attrs.data) 740 | 741 | 742 | class Base(Model): 743 | class Meta: 744 | client = Client(host=TEST_HOST, port=TEST_PORT) 745 | database = 0 746 | index_db = 3 747 | 748 | class User(Base): 749 | username = Field(index=True) 750 | status = IntegerField(index=True) 751 | 752 | class Misc(Base): 753 | f = Field() 754 | f_i = IntegerField(index=True) 755 | f_l = LongField(index=True) 756 | f_dt = DateTimeField(index=True) 757 | f_ts = TimestampField(index=True) 758 | f_b = BooleanField(index=True) 759 | 760 | 761 | class TestGreenQuery(BaseTestCase): 762 | def test_field_type_serialization(self): 763 | make_dt = lambda d: datetime.datetime(2018, 1, d, d, 0, 0) 764 | for i in range(1, 11): 765 | dt = make_dt(i) 766 | Misc.create(f='f%s' % i, f_i=i, f_l=i * 1000, f_dt=dt, f_ts=dt, 767 | f_b=(i % 3 == 1)) 768 | 769 | def assertQ(expr, nums): 770 | query = Misc.query(expr) 771 | rows = [(m.f, m.f_i, m.f_l, m.f_dt, m.f_ts, m.f_b) for m in query] 772 | expected = [('f%s' % n, n, n * 1000, make_dt(n), make_dt(n), 773 | int(n % 3 == 1)) for n in nums] 774 | self.assertEqual(rows, expected) 775 | 776 | assertQ((Misc.f_i == 5), [5]) 777 | assertQ((Misc.f_i > 8), [9, 10]) 778 | assertQ((Misc.f_i >= 8), [8, 9, 10]) 779 | assertQ((Misc.f_i < 3), [1, 2]) 780 | assertQ((Misc.f_i <= 3), [1, 2, 3]) 781 | assertQ(Misc.f_i.between(3, 6), [3, 4, 5]) 782 | assertQ(Misc.f_i.between(3, 6, False, True), [4, 5, 6]) 783 | 784 | assertQ((Misc.f_l == 5000), [5]) 785 | assertQ((Misc.f_l > 8000), [9, 10]) 786 | assertQ((Misc.f_l <= 3000), [1, 2, 3]) 787 | 788 | for field in (Misc.f_dt, Misc.f_ts): 789 | assertQ((field == make_dt(5)), [5]) 790 | assertQ((field > make_dt(8)), [9, 10]) 791 | assertQ((field >= make_dt(8)), [8, 9, 10]) 792 | assertQ((field < make_dt(3)), [1, 2]) 793 | assertQ((field <= make_dt(3)), [1, 2, 3]) 794 | assertQ(field.between(make_dt(3), make_dt(6)), [3, 4, 5]) 795 | assertQ(field.between(make_dt(3), make_dt(6), False, True), 796 | [4, 5, 6]) 797 | 798 | assertQ((Misc.f_b == True), [1, 10, 4, 7]) 799 | assertQ((Misc.f_b != False), [1, 10, 4, 7]) 800 | assertQ((Misc.f_b == False), [2, 3, 5, 6, 8, 9]) 801 | assertQ((Misc.f_b != True), [2, 3, 5, 6, 8, 9]) 802 | 803 | def test_field_data_types(self): 804 | # As long as the field is not indexed, we can store arbitrary data. 805 | # When the field *is* indexed, the field class must implement an 806 | # "index_value()" method to handle the appropriate conversion. 807 | Misc.create(f={'k1': 'v1', 'k2': 'v2'}, f_i=1) 808 | Misc.create(f=['foo', 'bar', 'baz'], f_i=2) 809 | 810 | m1 = Misc.get(Misc.f_i == 1) 811 | m2 = Misc.get(Misc.f_i == 2) 812 | 813 | self.assertEqual(m1.f, {'k1': 'v1', 'k2': 'v2'}) 814 | self.assertEqual(m2.f, ['foo', 'bar', 'baz']) 815 | 816 | def _create_test_users(self): 817 | username_status = ( 818 | ('huey', 1), 819 | ('mickey', 0), 820 | ('zaizee', 1), 821 | ('beanie', 2), 822 | ('gracie', 0), 823 | ('rocky', 0), 824 | ('rocky-2', 0), 825 | ) 826 | for username, status in username_status: 827 | User.create(username=username, status=status) 828 | return username_status 829 | 830 | def test_basic_operations(self): 831 | username_status = self._create_test_users() 832 | 833 | u_db = User.load(2) 834 | self.assertEqual(u_db.id, 2) 835 | self.assertEqual(u_db.username, 'mickey') 836 | self.assertEqual(u_db.status, 0) 837 | 838 | self.assertRaises(KeyError, User.load, 99) 839 | 840 | # We can retrieve all users, sorted by primary key. 841 | all_users = [user.username for user in User.all()] 842 | self.assertEqual(all_users, [u for u, _ in username_status]) 843 | 844 | # We can apply a limit. 845 | some_users = [user.username for user in User.all(3)] 846 | self.assertEqual(some_users, ['huey', 'mickey', 'zaizee']) 847 | 848 | # Query a single object using equality. 849 | huey = User.get(User.username == 'huey') 850 | self.assertEqual(huey.id, 1) 851 | self.assertEqual(huey.username, 'huey') 852 | 853 | def assertQ(expr, usernames): 854 | query = User.query(expr) 855 | self.assertEqual([u.username for u in query], usernames) 856 | 857 | # Query using ranges. 858 | assertQ(User.username.between('huey', 'rocky'), ['huey', 'mickey']) 859 | assertQ(User.username.between('huey', 'rocky', True, True), 860 | ['huey', 'mickey', 'rocky']) # Inclusive. 861 | assertQ((User.username >= 'rocky'), ['rocky', 'rocky-2', 'zaizee']) 862 | assertQ((User.username > 'rocky'), ['rocky-2', 'zaizee']) 863 | assertQ((User.username <= 'rocky'), 864 | ['beanie', 'gracie', 'huey', 'mickey', 'rocky']) 865 | assertQ((User.username < 'rocky'), 866 | ['beanie', 'gracie', 'huey', 'mickey']) 867 | assertQ((User.username == 'mickey'), ['mickey']) 868 | assertQ(User.username.startswith('ro'), ['rocky', 'rocky-2']) 869 | assertQ(User.username.startswith('rocky'), ['rocky', 'rocky-2']) 870 | assertQ(User.username.startswith('rocky '), []) 871 | 872 | assertQ((User.status > 1), ['beanie']) 873 | assertQ((User.status >= 1), ['huey', 'zaizee', 'beanie']) 874 | assertQ((User.status < 1), ['mickey', 'gracie', 'rocky', 'rocky-2']) 875 | assertQ((User.status != 0), ['huey', 'zaizee', 'beanie']) 876 | 877 | # Delete objects. 878 | for username in ('rocky', 'rocky-2', 'gracie', 'beanie'): 879 | user = User.get(User.username == username) 880 | user.delete() 881 | 882 | # Validate query results after deletions. 883 | self.assertEqual([u.username for u in User.all()], 884 | ['huey', 'mickey', 'zaizee']) 885 | assertQ(User.username.between('huey', 'rocky', True, True), 886 | ['huey', 'mickey']) 887 | assertQ((User.username >= 'rocky'), ['zaizee']) 888 | assertQ((User.username > 'rocky'), ['zaizee']) 889 | assertQ((User.username <= 'rocky'), ['huey', 'mickey']) 890 | assertQ((User.username < 'rocky'), ['huey', 'mickey']) 891 | assertQ((User.username == 'mickey'), ['mickey']) 892 | assertQ(User.username.startswith('h'), ['huey']) 893 | assertQ(User.username.startswith('r'), []) 894 | 895 | assertQ((User.status > 1), []) 896 | assertQ((User.status >= 1), ['huey', 'zaizee']) 897 | assertQ((User.status < 1), ['mickey']) 898 | assertQ((User.status != 0), ['huey', 'zaizee']) 899 | 900 | def test_compound_queries(self): 901 | username_status = self._create_test_users() 902 | 903 | def assertQ(expr, usernames): 904 | query = User.query(expr) 905 | self.assertEqual([u.username for u in query], usernames) 906 | 907 | assertQ((User.username == 'huey') | (User.username == 'zaizee'), 908 | ['huey', 'zaizee']) 909 | assertQ((User.username == 'huey') | (User.status == 1), 910 | ['huey', 'zaizee']) 911 | assertQ((User.username == 'huey') | (User.status == 2), 912 | ['huey', 'beanie']) 913 | assertQ((User.username == 'huey') & (User.status == 1), ['huey']) 914 | assertQ((User.username == 'huey') & (User.status != 1), []) 915 | 916 | assertQ((User.username.startswith('h') | 917 | User.username.startswith('b') | 918 | User.username.startswith('r')), 919 | ['huey', 'beanie', 'rocky', 'rocky-2']) 920 | 921 | 922 | if __name__ == '__main__': 923 | server_t, server, tmp_dir = run_server() 924 | unittest.main(argv=sys.argv) 925 | -------------------------------------------------------------------------------- /greendb.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from __future__ import unicode_literals # Required for 2.x compatability. 4 | 5 | import gevent 6 | from gevent import socket 7 | from gevent.local import local as greenlet_local 8 | from gevent.pool import Pool 9 | from gevent.server import StreamServer 10 | from gevent.thread import get_ident 11 | 12 | import lmdb 13 | from msgpack import packb as msgpack_packb 14 | from msgpack import unpackb as msgpack_unpackb 15 | 16 | from collections import namedtuple 17 | from contextlib import contextmanager 18 | from functools import wraps 19 | from io import BytesIO 20 | from socket import error as socket_error 21 | import datetime 22 | import heapq 23 | import json 24 | import logging 25 | import operator 26 | import optparse 27 | import os 28 | import re 29 | import shutil 30 | import sys 31 | import struct 32 | import time 33 | 34 | 35 | __version__ = '0.2.3' 36 | 37 | logger = logging.getLogger(__name__) 38 | 39 | 40 | if sys.version_info[0] == 3: 41 | unicode = str 42 | basestring = (bytes, str) 43 | int_types = int 44 | use_buffers = True 45 | else: 46 | int_types = (int, long) 47 | use_buffers = False 48 | 49 | def encode(s): 50 | if isinstance(s, unicode): 51 | return s.encode('utf-8') 52 | elif isinstance(s, bytes): 53 | return s 54 | else: 55 | return str(s).encode('utf-8') 56 | 57 | def encode_bulk_dict(d): 58 | accum = {} 59 | for key, value in d.items(): 60 | accum[encode(key)] = value 61 | return accum 62 | 63 | def encode_bulk_list(l): 64 | return [encode(k) for k in l] 65 | 66 | def decode(s): 67 | if isinstance(s, unicode): 68 | return s 69 | elif isinstance(s, bytes): 70 | return s.decode('utf-8') 71 | else: 72 | return str(s) 73 | 74 | def decode_bulk_dict(d): 75 | accum = {} 76 | for key, value in d.items(): 77 | accum[decode(key)] = value 78 | return accum 79 | 80 | 81 | class ClientQuit(Exception): pass 82 | class Shutdown(Exception): pass 83 | 84 | class ServerError(Exception): pass 85 | class ConnectionError(ServerError): pass 86 | class ServerInternalError(ServerError): pass 87 | 88 | class CommandError(Exception): 89 | def __init__(self, message): 90 | self.message = message 91 | super(CommandError, self).__init__() 92 | def __str__(self): 93 | return self.message 94 | __unicode__ = __str__ 95 | 96 | 97 | Error = namedtuple('Error', ('message',)) 98 | class Attributes(object): 99 | __slots__ = ('data',) 100 | def __init__(self, data=None): 101 | self.data = data or {} 102 | def __repr__(self): 103 | return '' % self.data 104 | def __getitem__(self, key): 105 | return self.data[key] 106 | def __iter__(self): 107 | return iter(self.data.items()) 108 | def __len__(self): 109 | return len(self.data) 110 | ProcessingInstruction = namedtuple('ProcessingInstruction', ('op', 'value')) 111 | VerbatimString = namedtuple('VerbatimString', ('text', 'prefix')) 112 | 113 | # Composite type for representing a response plus some attributes. 114 | _AttributeResponse = namedtuple('_AttributeResponse', ('data', 'attributes')) 115 | class AttributeResponse(_AttributeResponse): 116 | def __new__(cls, data, **attrs): 117 | attributes = Attributes(attrs) 118 | return super(AttributeResponse, cls).__new__(cls, data, attributes) 119 | 120 | PI_USE_DB = b'\x01' 121 | PROCESSING_INSTRUCTIONS = set((PI_USE_DB,)) 122 | 123 | CRLF = b'\r\n' 124 | MAXINT = 1 << 63 125 | READSIZE = 4 * 1024 126 | 127 | 128 | class _Socket(object): 129 | def __init__(self, s): 130 | self._socket = s 131 | self.is_closed = False 132 | self.buf = BytesIO() 133 | self.bytes_read = self.bytes_written = 0 134 | self.recvbuf = bytearray(READSIZE) 135 | 136 | self.sendbuf = BytesIO() 137 | self.bytes_pending = 0 138 | self._write_pending = False 139 | 140 | def __del__(self): 141 | if not self.is_closed: 142 | self.buf.close() 143 | self.sendbuf.close() 144 | try: 145 | self._socket.shutdown(socket.SHUT_RDWR) 146 | except socket.error: 147 | pass 148 | self._socket.close() 149 | 150 | def _read_from_socket(self, length): 151 | l = marker = 0 152 | recvptr = memoryview(self.recvbuf) 153 | self.buf.seek(self.bytes_written) 154 | 155 | try: 156 | while True: 157 | # Read up to READSIZE bytes into our dedicated recvbuf (using 158 | # buffer protocol to avoid copy). 159 | l = self._socket.recv_into(recvptr, READSIZE) 160 | if not l: 161 | self.close() 162 | raise ConnectionError('client went away') 163 | 164 | # Copy bytes read into file-like socket buffer. 165 | self.buf.write(recvptr[:l]) 166 | self.bytes_written += l 167 | 168 | marker += l 169 | if length > 0 and length > marker: 170 | continue 171 | break 172 | except socket.timeout: 173 | raise ConnectionError('timed out reading from socket') 174 | except socket.error: 175 | raise ConnectionError('error while reading from socket') 176 | 177 | def _read(self, length): 178 | buflen = self.bytes_written - self.bytes_read 179 | 180 | # Are we requesting more data than is available in the socket buffer? 181 | # If so, read from the socket into our socket buffer. This operation 182 | # may block or fail if the sender disconnects. 183 | if length > buflen: 184 | self._read_from_socket(length - buflen) 185 | 186 | # Move the cursor to the last-read byte of the socket buffer and read 187 | # the requested length. Update the last-read byte for subsequent reads. 188 | self.buf.seek(self.bytes_read) 189 | data = self.buf.read(length) 190 | self.bytes_read += length 191 | 192 | # If we have read all the data available in the socket buffer, truncate 193 | # so that it does not grow endlessly. 194 | if self.bytes_read == self.bytes_written: 195 | self.purge() 196 | return data 197 | 198 | def read(self, length): 199 | # Convenience function for reading a number of bytes followed by CRLF. 200 | # The CRLF is thrown away but will appear to have been read/consumed. 201 | data = self._read(length) 202 | self._read(2) # Read the CR/LF... assert self._read(2) == CRLF. 203 | return data 204 | 205 | def readline(self): 206 | # Move the cursor to the last-read byte of the socket buffer and 207 | # attempt to read up to the first CRLF. 208 | self.buf.seek(self.bytes_read) 209 | data = self.buf.readline() 210 | 211 | # If the data did not end with a CRLF, then get more from the socket 212 | # until we can read a line. 213 | while not data.endswith(CRLF): 214 | self._read_from_socket(0) 215 | self.buf.seek(self.bytes_read) 216 | data = self.buf.readline() 217 | 218 | # Update the last-read byte, marking our line as having been read. 219 | self.bytes_read += len(data) 220 | 221 | # If we've read all the available data in the socket buffer, truncate 222 | # so that it does not grow endlessly. 223 | if self.bytes_read == self.bytes_written: 224 | self.purge() 225 | 226 | return data[:-2] # Strip CRLF. 227 | 228 | def write(self, data): 229 | self.bytes_pending += self.sendbuf.write(data) 230 | 231 | def send(self, blocking=False): 232 | if self.bytes_pending and not self._write_pending: 233 | self._write_pending = True 234 | if blocking: 235 | self._send_to_socket() 236 | else: 237 | gevent.spawn(self._send_to_socket) 238 | return True 239 | return False 240 | 241 | def _send_to_socket(self): 242 | if self.is_closed: 243 | return 244 | 245 | data = self.sendbuf.getvalue() 246 | self.sendbuf.seek(0) 247 | self.sendbuf.truncate() 248 | self.bytes_pending = 0 249 | try: 250 | self._socket.sendall(data) 251 | except socket.error: 252 | self.close() 253 | raise ConnectionError('connection went away while sending data') 254 | finally: 255 | self._write_pending = False 256 | 257 | def purge(self): 258 | self.buf.seek(0) 259 | self.buf.truncate() 260 | self.bytes_read = self.bytes_written = 0 261 | 262 | def close(self): 263 | if self.is_closed: 264 | return False 265 | 266 | self.purge() 267 | self.buf.close() 268 | self.buf = None 269 | self.sendbuf.close() 270 | self.sendbuf = None 271 | 272 | try: 273 | self._socket.shutdown(socket.SHUT_RDWR) 274 | except: 275 | pass 276 | self._socket.close() 277 | self.is_closed = True 278 | return True 279 | 280 | 281 | class ProtocolHandler(object): 282 | def __init__(self): 283 | self.handlers = { 284 | b'*': self.handle_array, 285 | b'$': self.handle_blob, 286 | b'^': self.handle_unicode_string, 287 | b'+': self.handle_simple_string, 288 | b'-': self.handle_simple_error, 289 | b':': self.handle_number, 290 | b'_': self.handle_null, 291 | b',': self.handle_float, 292 | b'#': self.handle_boolean, 293 | b'!': self.handle_error, 294 | b'=': self.handle_verbatim_string, 295 | b'(': self.handle_number, 296 | b'%': self.handle_dict, 297 | b'~': self.handle_set, 298 | b'|': self.handle_attributes, 299 | b'>': self.handle_push, 300 | b'.': self.handle_processing_instruction, 301 | } 302 | 303 | def handle(self, sock): 304 | first_byte = sock._read(1) 305 | try: 306 | return self.handlers[first_byte](sock) 307 | except KeyError: 308 | # Handle un-encoded data as a raw simple string. 309 | rest = sock.readline() 310 | return first_byte + rest 311 | 312 | def handle_array(self, sock): 313 | num_elements = int(sock.readline()) 314 | return [self.handle(sock) for _ in range(num_elements)] 315 | 316 | def handle_blob(self, sock): 317 | length = int(sock.readline()) 318 | if length >= 0: 319 | return sock.read(length) 320 | 321 | def handle_unicode_string(self, sock): 322 | length = int(sock.readline()) 323 | if length >= 0: 324 | return sock.read(length).decode('utf8') 325 | 326 | def handle_simple_string(self, sock): 327 | return sock.readline() 328 | 329 | def handle_simple_error(self, sock): 330 | return Error(sock.readline()) 331 | 332 | def handle_number(self, sock): 333 | return int(sock.readline()) 334 | 335 | def handle_null(self, sock): 336 | sock.readline() 337 | return 338 | 339 | def handle_float(self, sock): 340 | return float(sock.readline()) 341 | 342 | def handle_boolean(self, sock): 343 | next_byte = sock.read(1) 344 | if next_byte == b't' or next_byte == b'T': 345 | return True 346 | elif next_byte == b'f' or next_byte == b'F': 347 | return False 348 | raise ValueError('unsupported value for boolean type') 349 | 350 | def handle_error(self, sock): 351 | length = int(sock.readline()) 352 | if length >= 0: 353 | return Error(sock.read(length)) 354 | 355 | def handle_verbatim_string(self, sock): 356 | length = int(sock.readline()) 357 | if length >= 0: 358 | data = sock.read(length).decode('utf8') 359 | prefix, rest = data[:3], data[4:] 360 | return VerbatimString(rest, prefix) 361 | 362 | def handle_dict(self, sock): 363 | accum = {} 364 | num_items = int(sock.readline()) 365 | for _ in range(num_items): 366 | key = self.handle(sock) 367 | accum[key] = self.handle(sock) 368 | return accum 369 | 370 | def handle_set(self, sock): 371 | return set(self.handle_array(sock)) 372 | 373 | def handle_attributes(self, sock): 374 | accum = {} 375 | num_items = int(sock.readline()) 376 | for _ in range(num_items): 377 | key = self.handle(sock) 378 | accum[key] = self.handle(sock) 379 | return Attributes(accum) 380 | 381 | def handle_push(self, sock): 382 | # Push data is represented same as Array type. First value is a string, 383 | # indicating the type of data being pushed. 384 | return self.handle_array(sock) 385 | 386 | def handle_processing_instruction(self, sock): 387 | instruction = sock.read(1) 388 | if instruction not in PROCESSING_INSTRUCTIONS: 389 | raise ValueError('unrecognized processing instruction, refusing to' 390 | ' process.') 391 | return ProcessingInstruction(instruction, self.handle(sock)) 392 | 393 | def write_response(self, sock, data, blocking=False): 394 | self._write(sock, data) 395 | sock.send(blocking) 396 | 397 | def _write(self, sock, data): 398 | if isinstance(data, (bytes, memoryview)): 399 | sock.write(b'$%d\r\n%s\r\n' % (len(data), data)) 400 | elif isinstance(data, unicode): 401 | data = encode(data) 402 | sock.write(b'^%d\r\n%s\r\n' % (len(data), data)) 403 | elif data is True or data is False: 404 | sock.write(b'#%s\r\n' % (b't' if data else b'f')) 405 | elif isinstance(data, int_types): 406 | if data > MAXINT: 407 | sock.write(b'(%d\r\n' % data) 408 | else: 409 | sock.write(b':%d\r\n' % data) 410 | elif data is None: 411 | sock.write(b'_\r\n') 412 | elif isinstance(data, Error): 413 | error = encode(data.message) 414 | sock.write(b'!%d\r\n%s\r\n' % (len(error), error)) 415 | elif isinstance(data, ProcessingInstruction): 416 | sock.write(b'.%s\r\n' % encode(data.op)) 417 | self._write(sock, data.value) 418 | elif isinstance(data, AttributeResponse): 419 | # Handle a composite response-type, consisting of attributes and 420 | # the actual response data. 421 | self._write(sock, data.attributes) 422 | self._write(sock, data.data) 423 | elif isinstance(data, VerbatimString): 424 | prefix = encode(data.prefix) 425 | data = encode(data.text) 426 | nchars = len(prefix) + len(data) + 1 427 | sock.write(b'=%d\r\n%s:%s\r\n' % (nchars, prefix, data)) 428 | elif isinstance(data, Attributes): 429 | attributes = data.data 430 | sock.write(b'|%d\r\n' % len(attributes)) 431 | for key in attributes: 432 | self._write(sock, key) 433 | self._write(sock, attributes[key]) 434 | elif isinstance(data, float): 435 | sock.write(b',%0.8f\r\n' % data) 436 | elif isinstance(data, (list, tuple)): 437 | sock.write(b'*%d\r\n' % len(data)) 438 | for item in data: 439 | self._write(sock, item) 440 | elif isinstance(data, set): 441 | sock.write(b'~%d\r\n' % len(data)) 442 | for item in data: 443 | self._write(sock, item) 444 | elif isinstance(data, dict): 445 | sock.write(b'%%%d\r\n' % len(data)) 446 | for key in data: 447 | self._write(sock, key) 448 | self._write(sock, data[key]) 449 | else: 450 | raise ValueError('unrecognized type') 451 | 452 | 453 | 454 | DEFAULT_MAP_SIZE = 1024 * 1024 * 256 # 256MB. 455 | 456 | 457 | class Storage(object): 458 | def __init__(self, path, map_size=DEFAULT_MAP_SIZE, read_only=False, 459 | metasync=True, sync=False, writemap=False, map_async=False, 460 | meminit=True, max_dbs=16, max_spare_txns=64, lock=True, 461 | dupsort=False): 462 | self._path = path 463 | self._config = { 464 | 'map_size': map_size, 465 | 'readonly': read_only, 466 | 'metasync': metasync, 467 | 'sync': sync, 468 | 'writemap': writemap, 469 | 'map_async': map_async, 470 | 'meminit': meminit, 471 | 'max_dbs': max_dbs, 472 | 'max_spare_txns': max_spare_txns, 473 | 'lock': lock, 474 | } 475 | self._dupsort = dupsort 476 | 477 | # Open LMDB environment and initialize data-structures. 478 | self.is_open = False 479 | self.open() 480 | 481 | def supports_dupsort(self, db): 482 | # Allow dupsort to be a list of db indexes or a simple boolean. 483 | if isinstance(self._dupsort, list): 484 | return db in self._dupsort 485 | return self._dupsort 486 | 487 | def open(self): 488 | if self.is_open: 489 | return False 490 | 491 | self.env = lmdb.open(self._path, **self._config) 492 | self.databases = {} 493 | 494 | for i in range(self._config['max_dbs']): 495 | # Allow databases to support duplicate values. 496 | self.databases[i] = self.env.open_db( 497 | encode('db%s' % i), 498 | dupsort=self.supports_dupsort(i)) 499 | 500 | self.is_open = True 501 | return True 502 | 503 | def close(self): 504 | if not self.is_open: 505 | return False 506 | 507 | self.sync(True) # Always sync before closing. 508 | self.env.close() 509 | return True 510 | 511 | def reset(self): 512 | self.close() 513 | if os.path.exists(self._path): 514 | shutil.rmtree(self._path) 515 | return self.open() 516 | 517 | def get_storage_config(self): 518 | config = {'path': self._path, 'dupsort': self._dupsort} 519 | config.update(self._config) # Add LMDB config. 520 | return config 521 | 522 | def stat(self): 523 | return self.env.stat() 524 | 525 | def count(self): 526 | return self.env.stat()['entries'] 527 | 528 | def info(self): 529 | return self.env.info() 530 | 531 | def reader_check(self): 532 | return self.env.reader_check() 533 | 534 | @contextmanager 535 | def db(self, db=0, write=False): 536 | if not self.is_open: 537 | raise ValueError('Cannot operate on closed environment.') 538 | 539 | if not isinstance(db, int_types): 540 | raise ValueError('database index must be integer') 541 | 542 | txn = self.env.begin(db=self.databases[db], write=write, 543 | buffers=use_buffers) 544 | try: 545 | yield txn 546 | except Exception as exc: 547 | txn.abort() 548 | raise 549 | else: 550 | txn.commit() 551 | 552 | def flush(self, db): 553 | db_handle = self.databases[db] 554 | with self.db(db, True) as txn: 555 | txn.drop(db_handle, delete=False) 556 | return True 557 | 558 | def flushall(self): 559 | return dict((db, self.flush(db)) for db in self.databases) 560 | 561 | def sync(self, force=False): 562 | return self.env.sync(force) 563 | 564 | 565 | class Connection(object): 566 | def __init__(self, storage, sock): 567 | self.storage = storage 568 | self.sock = sock 569 | self.db = 0 570 | 571 | def use_db(self, idx): 572 | if not isinstance(idx, int): 573 | raise CommandError('database index must be an integer') 574 | if idx not in self.storage.databases: 575 | raise CommandError('unrecognized database: %s' % idx) 576 | self.db = idx 577 | return self.db 578 | 579 | def ctx(self, write=False): 580 | return self.storage.db(self.db, write) 581 | 582 | @contextmanager 583 | def cursor(self, write=False): 584 | with self.ctx(write) as txn: 585 | cursor = txn.cursor() 586 | try: 587 | yield cursor 588 | finally: 589 | cursor.close() 590 | 591 | def close(self): 592 | self.sock.close() 593 | 594 | 595 | mpackb = lambda o: msgpack_packb(o, use_bin_type=True) 596 | munpackb = lambda b: msgpack_unpackb(b, raw=False) 597 | def mpackdict(d): 598 | for key, value in d.items(): 599 | yield (encode(key), mpackb(value)) 600 | 601 | def encodedict(d): 602 | for key, value in d.items(): 603 | yield (encode(key), encode(value)) 604 | 605 | 606 | def requires_dupsort(meth): 607 | @wraps(meth) 608 | def verify_dupsort(self, client, *args): 609 | if not self.storage.supports_dupsort(client.db): 610 | raise CommandError('Currently-selected database %s does not ' 611 | 'support dupsort.' % client.db) 612 | return meth(self, client, *args) 613 | return verify_dupsort 614 | 615 | 616 | class Server(object): 617 | def __init__(self, host='127.0.0.1', port=31337, max_clients=1024, 618 | path='data', **storage_config): 619 | self._host = host 620 | self._port = port 621 | self._max_clients = max_clients 622 | 623 | self._pool = Pool(max_clients) 624 | self._server = StreamServer( 625 | (self._host, self._port), 626 | self.connection_handler, 627 | spawn=self._pool) 628 | 629 | self._commands = self.get_commands() 630 | self._protocol = ProtocolHandler() 631 | self.storage = Storage(path, **storage_config) 632 | 633 | def get_commands(self): 634 | accum = {} 635 | commands = ( 636 | # Database / environment management. 637 | ('ENVINFO', self.envinfo), 638 | ('ENVSTAT', self.envstat), 639 | ('READERCHECK', self.reader_check), 640 | ('FLUSH', self.flush), 641 | ('FLUSHALL', self.flushall), 642 | ('PING', self.ping), 643 | ('STAT', self.stat), 644 | ('SYNC', self.sync), 645 | ('USE', self.use_db), 646 | 647 | # K/V operations. 648 | ('COUNT', self.count), 649 | ('DECR', self.decr), 650 | ('INCR', self.incr), 651 | ('CAS', self.cas), 652 | ('DELETE', self.delete), 653 | ('DELETEDUP', self.deletedup), 654 | ('DELETEDUPRAW', self.deletedupraw), 655 | ('DUPCOUNT', self.dupcount), 656 | ('EXISTS', self.exists), 657 | ('GET', self.get), 658 | ('GETDUP', self.getdup), 659 | ('GETRAW', self.getraw), 660 | ('LENGTH', self.length), 661 | ('POP', self.pop), 662 | ('REPLACE', self.replace), 663 | ('SET', self.set), 664 | ('SETDUP', self.setdup), 665 | ('SETDUPRAW', self.setdupraw), 666 | ('SETNX', self.setnx), 667 | ('SETRAW', self.setraw), 668 | 669 | # Bulk K/V operations. 670 | ('MDELETE', self.mdelete), 671 | ('MGET', self.mget), 672 | ('MGETDUP', self.mgetdup), 673 | ('MGETRAW', self.mgetraw), 674 | ('MPOP', self.mpop), 675 | ('MREPLACE', self.mreplace), 676 | ('MSET', self.mset), 677 | ('MSETDUP', self.msetdup), 678 | ('MSETNX', self.msetnx), 679 | ('MSETRAW', self.msetraw), 680 | 681 | # Cursor operations. 682 | ('DELETERANGE', self.deleterange), 683 | ('GETRANGE', self.getrange), 684 | ('GETRANGEDUPRAW', self.getrangedupraw), 685 | ('GETRANGERAW', self.getrangeraw), 686 | ('ITEMS', self.getrange), 687 | ('KEYS', self.keys), 688 | ('PREFIX', self.match_prefix), 689 | ('VALUES', self.values), 690 | 691 | # Client operations. 692 | ('SLEEP', self.client_sleep), 693 | ('QUIT', self.client_quit), 694 | ('SHUTDOWN', self.shutdown), 695 | ) 696 | for cmd, callback in commands: 697 | accum[encode(cmd)] = callback 698 | return accum 699 | 700 | # Database / environment management. 701 | def envinfo(self, client): 702 | info = self.storage.info() 703 | info.update( 704 | clients=len(self._pool), 705 | host=self._host, 706 | port=self._port, 707 | max_clients=self._max_clients, 708 | storage=self.storage.get_storage_config()) 709 | return info 710 | 711 | def envstat(self, client): 712 | return self.storage.stat() 713 | 714 | def reader_check(self, client): 715 | return self.storage.reader_check() 716 | 717 | def flush(self, client): 718 | return self.storage.flush(client.db) 719 | 720 | def flushall(self, client): 721 | return self.storage.flushall() 722 | 723 | def ping(self, client): 724 | return b'pong' 725 | 726 | def stat(self, client): 727 | with client.ctx() as txn: 728 | return txn.stat() 729 | 730 | def sync(self, client): 731 | return self.storage.sync() 732 | 733 | def use_db(self, client, idx): 734 | return client.use_db(idx) 735 | 736 | # K/V operations. 737 | def count(self, client): 738 | with client.ctx() as txn: 739 | stat = txn.stat() 740 | return stat['entries'] 741 | 742 | def decr(self, client, key, amount=1): 743 | return self._incr(client, encode(key), -amount) 744 | 745 | def incr(self, client, key, amount=1): 746 | return self._incr(client, encode(key), amount) 747 | 748 | def _incr(self, client, key, amount): 749 | with client.cursor(True) as cursor: 750 | # If the key does not exist, just set the desired value. 751 | if not cursor.set_key(key): 752 | cursor.put(key, mpackb(amount)) 753 | return amount 754 | 755 | orig = munpackb(cursor.value()) 756 | try: 757 | value = orig + amount 758 | except TypeError: 759 | raise CommandError('decr operation on wrong type of value') 760 | 761 | cursor.delete() 762 | cursor.put(key, mpackb(value)) 763 | return value 764 | 765 | def cas(self, client, key, old_value, new_value): 766 | key = encode(key) 767 | with client.ctx(True) as txn: 768 | value = txn.get(key) 769 | if value is not None and munpackb(value) == old_value: 770 | if self.storage.supports_dupsort(client.db): 771 | txn.delete(key, value) 772 | txn.put(key, mpackb(new_value)) 773 | return True 774 | elif value is None and old_value is None: 775 | txn.put(key, mpackb(new_value)) 776 | return True 777 | else: 778 | return False 779 | 780 | def delete(self, client, key): 781 | with client.ctx(True) as txn: 782 | return txn.delete(encode(key)) 783 | 784 | @requires_dupsort 785 | def deletedup(self, client, key, value): 786 | with client.ctx(True) as txn: 787 | return txn.delete(encode(key), mpackb(value)) 788 | 789 | @requires_dupsort 790 | def deletedupraw(self, client, key, value): 791 | with client.ctx(True) as txn: 792 | return txn.delete(encode(key), encode(value)) 793 | 794 | @requires_dupsort 795 | def dupcount(self, client, key): 796 | with client.cursor() as cursor: 797 | if not cursor.set_key(encode(key)): 798 | return 799 | return cursor.count() 800 | 801 | def exists(self, client, key): 802 | sentinel = object() 803 | with client.ctx() as txn: 804 | return txn.get(encode(key), sentinel) is not sentinel 805 | 806 | def get(self, client, key): 807 | with client.ctx() as txn: 808 | res = txn.get(encode(key)) 809 | if res is not None: 810 | return munpackb(res) 811 | 812 | def getdup(self, client, key): 813 | key = encode(key) 814 | with client.cursor() as cursor: 815 | if not cursor.set_key(key): 816 | return 817 | accum = [] 818 | while cursor.key() == key: 819 | accum.append(munpackb(cursor.value())) 820 | if not cursor.next_dup(): 821 | break 822 | return accum 823 | 824 | def getraw(self, client, key): 825 | with client.ctx() as txn: 826 | return txn.get(encode(key)) 827 | 828 | def length(self, client, key): 829 | value = self.get(client, key) 830 | if value is not None: 831 | try: 832 | return len(value) 833 | except TypeError: 834 | raise CommandError('incompatible type for LENGTH command') 835 | 836 | def pop(self, client, key): 837 | with client.ctx(True) as txn: 838 | res = txn.pop(encode(key)) 839 | if res is not None: 840 | return munpackb(res) 841 | 842 | def replace(self, client, key, value): 843 | with client.ctx(True) as txn: 844 | old_val = txn.replace(encode(key), mpackb(value)) 845 | if old_val is not None: 846 | return munpackb(old_val) 847 | 848 | def set(self, client, key, value): 849 | with client.ctx(True) as txn: 850 | return txn.put(encode(key), mpackb(value), dupdata=False) 851 | 852 | @requires_dupsort 853 | def setdup(self, client, key, value): 854 | with client.ctx(True) as txn: 855 | return txn.put(encode(key), mpackb(value)) 856 | 857 | @requires_dupsort 858 | def setdupraw(self, client, key, value): 859 | with client.ctx(True) as txn: 860 | return txn.put(encode(key), encode(value)) 861 | 862 | def setnx(self, client, key, value): 863 | with client.ctx(True) as txn: 864 | return txn.put(encode(key), mpackb(value), dupdata=False, 865 | overwrite=False) 866 | 867 | def setraw(self, client, key, value): 868 | with client.ctx(True) as txn: 869 | return txn.put(encode(key), encode(value), dupdata=False) 870 | 871 | # Bulk K/V operations. 872 | def mdelete(self, client, keys): 873 | n = 0 874 | with client.ctx(True) as txn: 875 | for key in map(encode, keys): 876 | if txn.delete(key): 877 | n += 1 878 | return n 879 | 880 | def mget(self, client, keys): 881 | accum = {} 882 | with client.ctx() as txn: 883 | for key in map(encode, keys): 884 | res = txn.get(key) 885 | if res is not None: 886 | accum[key] = munpackb(res) 887 | return accum 888 | 889 | def mgetdup(self, client, keys): 890 | accum = {} 891 | with client.cursor() as cursor: 892 | for key in map(encode, keys): 893 | if cursor.set_key(key): 894 | values = [] 895 | while cursor.key() == key: 896 | values.append(munpackb(cursor.value())) 897 | if not cursor.next_dup(): 898 | break 899 | accum[key] = values 900 | return accum 901 | 902 | def mgetraw(self, client, keys): 903 | accum = {} 904 | with client.ctx() as txn: 905 | for key in map(encode, keys): 906 | res = txn.get(key) 907 | if res is not None: 908 | accum[key] = res 909 | return accum 910 | 911 | def mpop(self, client, keys): 912 | accum = {} 913 | with client.cursor(True) as cursor: 914 | for key in map(encode, keys): 915 | res = cursor.pop(key) 916 | if res is not None: 917 | accum[key] = munpackb(res) 918 | return accum 919 | 920 | def mreplace(self, client, data): 921 | accum = {} 922 | with client.cursor(True) as cursor: 923 | for key, value in data.items(): 924 | key = encode(key) 925 | old_val = cursor.replace(key, mpackb(value)) 926 | if old_val is not None: 927 | accum[key] = munpackb(old_val) 928 | return accum 929 | 930 | def mset(self, client, data): 931 | with client.cursor(True) as cursor: 932 | consumed, added = cursor.putmulti(mpackdict(data), dupdata=False) 933 | return added 934 | 935 | def msetdup(self, client, data): 936 | with client.cursor(True) as cursor: 937 | consumed, added = cursor.putmulti(mpackdict(data)) 938 | return added 939 | 940 | def msetnx(self, client, data): 941 | with client.cursor(True) as cursor: 942 | consumed, added = cursor.putmulti(mpackdict(data), dupdata=False, 943 | overwrite=False) 944 | return added 945 | 946 | def msetraw(self, client, data): 947 | with client.cursor(True) as cursor: 948 | consumed, added = cursor.putmulti(encodedict(data), dupdata=False) 949 | return added 950 | 951 | # Cursor operations. 952 | def deleterange(self, client, start=None, stop=None, count=None): 953 | if count is None: 954 | count = 0 955 | stop = encode(stop) if stop is not None else stop 956 | n = 0 957 | 958 | with client.cursor(write=True) as cursor: 959 | if start is None: 960 | if not cursor.first(): 961 | return n 962 | elif not cursor.set_range(encode(start)): 963 | return n 964 | 965 | while True: 966 | key = cursor.key() 967 | if use_buffers: 968 | key = key.tobytes() 969 | if stop is not None and key > stop: 970 | break 971 | 972 | if not cursor.delete(): 973 | break 974 | 975 | n += 1 976 | count -= 1 977 | if count == 0: 978 | break 979 | 980 | return n 981 | 982 | def _cursor_op(self, client, start, stop, count, cb, stopcond=operator.gt): 983 | accum = [] 984 | if count is None: 985 | count = 0 986 | stop = encode(stop) if stop is not None else stop 987 | 988 | with client.cursor() as cursor: 989 | if start is None: 990 | if not cursor.first(): 991 | return [] 992 | elif not cursor.set_range(encode(start)): 993 | return [] 994 | 995 | while True: 996 | key, data = cb(cursor) 997 | if use_buffers: 998 | key = key.tobytes() 999 | if stop is not None and stopcond(key, stop): 1000 | break 1001 | accum.append(data) 1002 | count -= 1 1003 | if count == 0 or not cursor.next(): 1004 | break 1005 | 1006 | return accum 1007 | 1008 | def getrange(self, client, start=None, stop=None, count=None): 1009 | def cb(cursor): 1010 | key, value = cursor.item() 1011 | return key, (key, munpackb(value)) 1012 | return self._cursor_op(client, start, stop, count, cb) 1013 | 1014 | @requires_dupsort 1015 | def getrangedupraw(self, client, key, start=None, stop=None, count=None): 1016 | accum = [] 1017 | if count is None: 1018 | count = 0 1019 | stop = encode(stop) if stop is not None else stop 1020 | 1021 | with client.cursor() as cursor: 1022 | if start is None: 1023 | if not cursor.set_range(encode(key)): 1024 | return [] 1025 | elif not cursor.set_range_dup(encode(key), encode(start)): 1026 | return [] 1027 | 1028 | while True: 1029 | value = cursor.value() 1030 | if use_buffers: 1031 | value = value.tobytes() 1032 | if stop is not None and value > stop: 1033 | break 1034 | accum.append(value) 1035 | count -= 1 1036 | if count == 0 or not cursor.next_dup(): 1037 | break 1038 | 1039 | return accum 1040 | 1041 | def getrangeraw(self, client, start=None, stop=None, count=None): 1042 | def cb(cursor): 1043 | key, value = cursor.item() 1044 | return key, (key, value) 1045 | return self._cursor_op(client, start, stop, count, cb) 1046 | 1047 | def keys(self, client, start=None, stop=None, count=None): 1048 | def cb(cursor): 1049 | key = cursor.key() 1050 | return key, key 1051 | return self._cursor_op(client, start, stop, count, cb) 1052 | 1053 | def values(self, client, start=None, stop=None, count=None): 1054 | def cb(cursor): 1055 | key, value = cursor.item() 1056 | return key, munpackb(value) 1057 | return self._cursor_op(client, start, stop, count, cb) 1058 | 1059 | def match_prefix(self, client, prefix, count=None): 1060 | def cb(cursor): 1061 | key, value = cursor.item() 1062 | return key, (key, munpackb(value)) 1063 | stopcond = lambda k, p: not k.startswith(p) 1064 | return self._cursor_op(client, prefix, prefix, count, cb, stopcond) 1065 | 1066 | # Client operations. 1067 | def client_sleep(self, client, timeout=1): 1068 | start = time.time() 1069 | gevent.sleep(timeout) 1070 | stop = time.time() 1071 | return AttributeResponse(timeout, start=start, stop=stop) 1072 | 1073 | def client_quit(self, client): 1074 | raise ClientQuit('client closed connection') 1075 | 1076 | def shutdown(self, client): 1077 | raise Shutdown('shutting down') 1078 | 1079 | # Server implementation. 1080 | def run(self): 1081 | try: 1082 | self._server.serve_forever() 1083 | finally: 1084 | self.storage.close() 1085 | 1086 | def connection_handler(self, conn, address): 1087 | logger.info('Connection received: %s:%s' % address) 1088 | client = Connection(self.storage, _Socket(conn)) 1089 | while True: 1090 | try: 1091 | self.request_response(client) 1092 | except ConnectionError: 1093 | logger.info('Client went away: %s:%s' % address) 1094 | client.close() 1095 | break 1096 | except ClientQuit: 1097 | logger.info('Client exited: %s:%s.' % address) 1098 | break 1099 | except Exception as exc: 1100 | logger.exception('Error processing command.') 1101 | 1102 | def request_response(self, client): 1103 | data = self._protocol.handle(client.sock) 1104 | 1105 | # If we received a processing instruction, it will be handled here, as 1106 | # the next request will contain the relevant data. 1107 | if isinstance(data, ProcessingInstruction): 1108 | self.execute_processing_instruction(client, data) 1109 | return 1110 | 1111 | try: 1112 | resp = self.respond(client, data) 1113 | except Shutdown: 1114 | logger.info('Shutting down') 1115 | self._protocol.write_response(client.sock, 1, True) 1116 | raise KeyboardInterrupt 1117 | except ClientQuit: 1118 | self._protocol.write_response(client.sock, 1, True) 1119 | raise 1120 | except CommandError as command_error: 1121 | resp = Error(command_error.message) 1122 | except Exception as exc: 1123 | logger.exception('Unhandled error') 1124 | resp = Error('Unhandled server error: "%s"' % str(exc)) 1125 | 1126 | self._protocol.write_response(client.sock, resp) 1127 | 1128 | def execute_processing_instruction(self, client, pi): 1129 | if pi.op == PI_USE_DB and client.db != pi.value: 1130 | orig_db = client.db 1131 | try: 1132 | client.use_db(pi.value) 1133 | self.request_response(client) 1134 | finally: 1135 | client.use_db(orig_db) 1136 | 1137 | def respond(self, client, data): 1138 | if not isinstance(data, (list, tuple)): 1139 | try: 1140 | data = data.split() 1141 | except: 1142 | raise CommandError('Unrecognized request type.') 1143 | 1144 | if not isinstance(data[0], basestring): 1145 | raise CommandError('First parameter must be command name.') 1146 | 1147 | command = data[0].upper() 1148 | if command not in self._commands: 1149 | raise CommandError('Unrecognized command: %s' % command) 1150 | else: 1151 | logger.debug('Received %s', decode(command)) 1152 | 1153 | return self._commands[command](client, *data[1:]) 1154 | 1155 | 1156 | class SocketPool(object): 1157 | def __init__(self, host, port, timeout=60, max_age=None): 1158 | self.host = host 1159 | self.port = port 1160 | self.timeout = timeout 1161 | self.max_age = max_age or 3600 1162 | self.free = [] 1163 | self.in_use = {} 1164 | 1165 | def checkout(self): 1166 | now = time.time() 1167 | tid = get_ident() 1168 | if tid in self.in_use: 1169 | sock = self.in_use[tid] 1170 | if sock.is_closed: 1171 | del self.in_use[tid] 1172 | else: 1173 | return self.in_use[tid] 1174 | 1175 | while self.free: 1176 | ts, sock = heapq.heappop(self.free) 1177 | if ts < now - self.max_age: 1178 | sock.close() 1179 | else: 1180 | self.in_use[tid] = sock 1181 | return sock 1182 | 1183 | sock = self.create_socket() 1184 | self.in_use[tid] = sock 1185 | return sock 1186 | 1187 | def create_socket(self): 1188 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 1189 | sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) 1190 | sock.settimeout(self.timeout) 1191 | sock.connect((self.host, self.port)) 1192 | return _Socket(sock) 1193 | 1194 | def checkin(self): 1195 | tid = get_ident() 1196 | if tid in self.in_use: 1197 | sock = self.in_use.pop(tid) 1198 | if not sock.is_closed: 1199 | heapq.heappush(self.free, (time.time(), sock)) 1200 | return True 1201 | return False 1202 | 1203 | def close(self): 1204 | tid = get_ident() 1205 | sock = self.in_use.pop(tid, None) 1206 | if sock: 1207 | sock.close() 1208 | return True 1209 | return False 1210 | 1211 | 1212 | class _ConnectionState(object): 1213 | def __init__(self, **kwargs): 1214 | super(_ConnectionState, self).__init__(**kwargs) 1215 | self.reset() 1216 | def reset(self): self.conn = None 1217 | def set_connection(self, conn): self.conn = conn 1218 | 1219 | class _ConnectionLocal(_ConnectionState, greenlet_local): pass 1220 | 1221 | 1222 | class Client(object): 1223 | def __init__(self, host='127.0.0.1', port=31337, decode_keys=False, 1224 | timeout=60, pool=True, max_age=None): 1225 | self.host = host 1226 | self.port = port 1227 | self._decode_keys = decode_keys 1228 | self._timeout = timeout 1229 | self._pool = SocketPool(host, port, timeout, max_age) if pool else None 1230 | self._protocol = ProtocolHandler() 1231 | self._state = _ConnectionLocal() 1232 | 1233 | def is_closed(self): 1234 | return self._state.conn is None 1235 | 1236 | def close(self): 1237 | if self._state.conn is None: return False 1238 | if self._pool is not None: 1239 | self._pool.checkin() 1240 | else: 1241 | self._state.conn.close() 1242 | self._state.reset() 1243 | return True 1244 | 1245 | def connect(self): 1246 | if self._state.conn is not None: return False 1247 | if self._pool is None: 1248 | sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 1249 | sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) 1250 | sock.settimeout(self._timeout) 1251 | sock.connect((self.host, self.port)) 1252 | conn = _Socket(sock) 1253 | else: 1254 | conn = self._pool.checkout() 1255 | 1256 | self._state.conn = conn 1257 | return True 1258 | 1259 | def read_response(self, conn, close_conn=False): 1260 | try: 1261 | resp = self._protocol.handle(conn) 1262 | except EOFError: 1263 | self.close() 1264 | raise ConnectionError('server went away') 1265 | except Exception: 1266 | logger.exception('internal server error') 1267 | self.close() 1268 | raise ServerInternalError('internal server error') 1269 | else: 1270 | if close_conn: 1271 | self.close() 1272 | if isinstance(resp, Error): 1273 | raise CommandError(decode(resp.message)) 1274 | if self._decode_keys and isinstance(resp, dict): 1275 | resp = decode_bulk_dict(resp) 1276 | return resp 1277 | 1278 | def execute(self, cmd, args, close_conn=False, db=None): 1279 | if self._state.conn is None: 1280 | self.connect() 1281 | 1282 | conn = self._state.conn 1283 | 1284 | # If the db is explicitly specified for a one-off command, then 1285 | # handle this using a processing instruction. 1286 | if db is not None: 1287 | pi = ProcessingInstruction(PI_USE_DB, db) 1288 | self._protocol._write(conn, pi) 1289 | 1290 | # Execute the command. The command and its arguments (if it has any) 1291 | # are all packed into a tuple and written to the server as a list. 1292 | self._protocol.write_response(conn, (cmd,) + args) 1293 | 1294 | # If the response is of type `Attributes`, continue reading the next 1295 | # response. 1296 | resp = self.read_response(conn, close_conn) 1297 | if isinstance(resp, Attributes): 1298 | return self.read_response(conn, False), resp 1299 | else: 1300 | return resp 1301 | 1302 | def command(cmd, close_conn=False): 1303 | def method(self, *args, **kwargs): 1304 | return self.execute(encode(cmd), args, close_conn, **kwargs) 1305 | return method 1306 | 1307 | # Database/schema management. 1308 | envinfo = command('ENVINFO') 1309 | envstat = command('ENVSTAT') 1310 | reader_check = command('READERCHECK') 1311 | flush = command('FLUSH') 1312 | flushall = command('FLUSHALL') 1313 | ping = command('PING') 1314 | stat = command('STAT') 1315 | sync = command('SYNC') 1316 | use = command('USE') 1317 | 1318 | # Basic k/v operations. 1319 | count = command('COUNT') 1320 | decr = command('DECR') 1321 | incr = command('INCR') 1322 | cas = command('CAS') 1323 | delete = command('DELETE') 1324 | deletedup = command('DELETEDUP') 1325 | deletedupraw = command('DELETEDUPRAW') 1326 | dupcount = command('DUPCOUNT') 1327 | exists = command('EXISTS') 1328 | get = command('GET') 1329 | getdup = command('GETDUP') 1330 | getraw = command('GETRAW') 1331 | length = command('LENGTH') 1332 | pop = command('POP') 1333 | replace = command('REPLACE') 1334 | set = command('SET') 1335 | setdup = command('SETDUP') 1336 | setdupraw = command('SETDUPRAW') 1337 | setnx = command('SETNX') 1338 | setraw = command('SETRAW') 1339 | 1340 | # Bulk k/v operations. 1341 | mdelete = command('MDELETE') 1342 | mget = command('MGET') 1343 | mgetdup = command('MGETDUP') 1344 | mgetraw = command('MGETRAW') 1345 | mpop = command('MPOP') 1346 | mreplace = command('MREPLACE') 1347 | mset = command('MSET') 1348 | msetdup = command('MSETDUP') 1349 | msetnx = command('MSETNX') 1350 | msetraw = command('MSETRAW') 1351 | 1352 | # Cursor operations. 1353 | deleterange = command('DELETERANGE') 1354 | getrange = command('GETRANGE') 1355 | getrangedupraw = command('GETRANGEDUPRAW') 1356 | getrangeraw = command('GETRANGERAW') 1357 | items = command('ITEMS') 1358 | keys = command('KEYS') 1359 | prefix = command('PREFIX') 1360 | values = command('VALUES') 1361 | 1362 | # Client operations. 1363 | _sleep = command('SLEEP') 1364 | quit = command('QUIT', close_conn=True) 1365 | shutdown = command('SHUTDOWN', close_conn=True) 1366 | 1367 | def __setitem__(self, key, value): 1368 | self.set(key, value) 1369 | 1370 | def __getitem__(self, item): 1371 | if isinstance(item, slice): 1372 | return self.getrange(item.start, item.stop, item.step) 1373 | elif isinstance(item, (list, tuple)): 1374 | return self.mget(item) 1375 | return self.get(item) 1376 | 1377 | def __delitem__(self, item): 1378 | if isinstance(item, slice): 1379 | self.deleterange(item.start, item.stop, item.step) 1380 | elif isinstance(item, (list, tuple)): 1381 | self.mdelete(item) 1382 | else: 1383 | self.delete(item) 1384 | 1385 | __len__ = count 1386 | __contains__ = exists 1387 | 1388 | def update(self, __data=None, **kwargs): 1389 | if __data is not None: 1390 | params = __data 1391 | params.update(kwargs) 1392 | else: 1393 | params = kwargs 1394 | return self.mset(params) 1395 | 1396 | def __iter__(self): 1397 | return iter(self.keys()) 1398 | 1399 | 1400 | def get_option_parser(): 1401 | parser = optparse.OptionParser() 1402 | parser.add_option('-c', '--config', default='config.json', dest='config', 1403 | help='Config file (default="config.json")') 1404 | parser.add_option('-D', '--data-dir', default='data', dest='data_dir', 1405 | help='Directory to store db environment and data.') 1406 | parser.add_option('-d', '--debug', action='store_true', dest='debug', 1407 | help='Log debug messages.') 1408 | parser.add_option('-e', '--errors', action='store_true', dest='error', 1409 | help='Log error messages only.') 1410 | parser.add_option('-H', '--host', default='127.0.0.1', dest='host', 1411 | help='Host to listen on.') 1412 | parser.add_option('-l', '--log-file', dest='log_file', help='Log file.') 1413 | parser.add_option('-m', '--map-size', dest='map_size', help=( 1414 | 'Maximum size of memory-map used for database. The default value is ' 1415 | '256M and should be increased. Accepts value in bytes or file-size ' 1416 | 'using "M" or "G" suffix.')) 1417 | parser.add_option('--max-clients', default=1024, dest='max_clients', 1418 | help='Maximum number of clients.', type=int) 1419 | parser.add_option('-n', '--max-dbs', default=16, dest='max_dbs', 1420 | help='Number of databases in environment. Default=16.', 1421 | type='int') 1422 | parser.add_option('-p', '--port', default=31337, dest='port', 1423 | help='Port to listen on.', type=int) 1424 | parser.add_option('-r', '--reset', action='store_true', dest='reset', 1425 | help='Reset database and config. All data will be lost.') 1426 | parser.add_option('-s', '--sync', action='store_true', dest='sync', 1427 | help=('Flush system buffers to disk when committing a ' 1428 | 'transaction. Durable but much slower.')) 1429 | parser.add_option('-u', '--dupsort', action='append', dest='dupsort', 1430 | help='db index(es) to support dupsort', type='int'), 1431 | parser.add_option('-M', '--no-metasync', action='store_true', 1432 | dest='no_metasync', help=( 1433 | 'Flush system buffers to disk only once per ' 1434 | 'transaction, omit the metadata flush.')) 1435 | parser.add_option('-W', '--writemap', action='store_true', dest='writemap', 1436 | help='Use a writeable memory map.') 1437 | parser.add_option('-A', '--map-async', action='store_true', 1438 | dest='map_async', help=( 1439 | 'When used with "--writemap" (-W), use asynchronous ' 1440 | 'flushes to disk.')) 1441 | return parser 1442 | 1443 | 1444 | def configure_logger(options): 1445 | logger.addHandler(logging.StreamHandler()) 1446 | if options.log_file: 1447 | logger.addHandler(logging.FileHandler(options.log_file)) 1448 | if options.debug: 1449 | logger.setLevel(logging.DEBUG) 1450 | elif options.error: 1451 | logger.setLevel(logging.ERROR) 1452 | else: 1453 | logger.setLevel(logging.INFO) 1454 | 1455 | 1456 | def read_config(config_file): 1457 | if not os.path.exists(config_file): return {} 1458 | 1459 | with open(config_file) as fh: 1460 | data = fh.read() 1461 | 1462 | # Strip comments. 1463 | config = json.loads(re.sub('\s*\/\/.*', '', data)) 1464 | 1465 | if config.get('map_size'): 1466 | config['map_size'] = parse_map_size(config['map_size']) 1467 | return config 1468 | 1469 | def log_config(conf): 1470 | for key, value in sorted(conf.items()): 1471 | if value is not None: 1472 | logger.debug('%s=%s' % (key, value)) 1473 | 1474 | def parse_map_size(value): 1475 | mapsize = value.lower() if value else '256m' 1476 | 1477 | if mapsize.endswith('b'): 1478 | mapsize = mapsize[:-1] # Strip "b", as in "mb", "gb", etc. 1479 | n = 1 1480 | if mapsize.endswith('k'): 1481 | exp = 1024 1482 | elif mapsize.endswith('m'): 1483 | exp = 1024 * 1024 1484 | elif mapsize.endswith('g'): 1485 | exp = 1024 * 1024 * 1024 1486 | else: 1487 | exp = 1 1488 | n = 0 1489 | mapsize = mapsize[:-n] if n else mapsize # Strip trailing letter. 1490 | if not mapsize.isdigit(): 1491 | raise ValueError('cannot parse file-size "%s", use a valid file-' 1492 | 'size like "256m" or "1g"' % value) 1493 | return int(mapsize) * exp 1494 | 1495 | 1496 | if __name__ == '__main__': 1497 | options, args = get_option_parser().parse_args() 1498 | 1499 | configure_logger(options) 1500 | if options.reset and os.path.exists(options.data_dir): 1501 | shutil.rmtree(options.data_dir) 1502 | 1503 | config = read_config(options.config or 'config.json') 1504 | config.setdefault('path', options.data_dir) 1505 | config.setdefault('host', options.host) 1506 | config.setdefault('map_size', parse_map_size(options.map_size)) 1507 | config.setdefault('max_clients', options.max_clients) 1508 | config.setdefault('max_dbs', options.max_dbs) 1509 | config.setdefault('port', options.port) 1510 | config.setdefault('sync', bool(options.sync)) 1511 | config.setdefault('dupsort', options.dupsort) 1512 | config.setdefault('metasync', not options.no_metasync) 1513 | config.setdefault('writemap', options.writemap) 1514 | config.setdefault('map_async', options.map_async) 1515 | server = Server(**config) 1516 | 1517 | print('\x1b[32m .--.') 1518 | print(' /( \x1b[34m@\x1b[33m >\x1b[32m ,-. ' 1519 | '\x1b[1;32mgreendb ' 1520 | '\x1b[1;33m%s:%s\x1b[32m' % (server._host, server._port)) 1521 | print('/ \' .\'--._/ /') 1522 | print(': , , .\'') 1523 | print('\'. (___.\'_/') 1524 | print(' \x1b[33m((\x1b[32m-\x1b[33m((\x1b[32m-\'\'\x1b[0m') 1525 | log_config(config) 1526 | try: 1527 | server.run() 1528 | except KeyboardInterrupt: 1529 | print('\x1b[1;31mshutting down\x1b[0m') 1530 | --------------------------------------------------------------------------------