Skip to content
This repository
Browse code

added sqlite3 support

  • Loading branch information...
commit 212f56ddedb96ffb170f49884a6deab4903c1e45 1 parent a669eab
Andrey Usov authored

Showing 1 changed file with 268 additions and 51 deletions. Show diff stats Hide diff stats

  1. 319  kvlite.py
319  kvlite.py
@@ -37,11 +37,13 @@
37 37
 ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
38 38
 POSSIBILITY OF SUCH DAMAGE."""
39 39
 
  40
+import os
40 41
 import cmd
41 42
 import sys
42 43
 import zlib
43 44
 import uuid
44 45
 import pprint
  46
+import sqlite3
45 47
 import binascii
46 48
 
47 49
 __all__ = ['open', 'remove',]
@@ -52,23 +54,82 @@
52 54
     print >> sys.stderr, 'Error! MySQLdb package is not installed, please install python-mysqldb'
53 55
     sys.exit()
54 56
 
  57
+
55 58
 try:
56 59
     import cPickle as pickle
57 60
 except ImportError:
58 61
     import pickle
59  
-
60  
-# TODO add JSON support for serialization
  62
+# TODO add test cases for serializers
61 63
 # TODO add support user specific serializators
62  
-from json import loads as json_decode
63  
-from json import dumps as json_encode
64 64
 
65 65
 SUPPORTED_BACKENDS = ['mysql', 'sqlite', ]
66 66
 
  67
+'''
  68
+A collection is a group of documents stored in kvlite2, 
  69
+and can be thought of as roughly the equivalent of a 
  70
+table in a relational database.
  71
+
  72
+
  73
+
  74
+For using JSON as serialization
  75
+
  76
+>>> import json
  77
+>>> collection = open('sqlite://test.sqlite:test', serializer=json)
  78
+>>>
  79
+
  80
+'''
  81
+# -----------------------------------------------------------------
  82
+# cPickleSerializer class
  83
+# -----------------------------------------------------------------
  84
+
  85
+class cPickleSerializer(object):
  86
+    ''' cPickleSerializer '''
  87
+
  88
+    @staticmethod
  89
+    def dumps(v):
  90
+        ''' dumps value '''
  91
+        if isinstance(v, unicode):
  92
+            v = str(v)
  93
+        return pickle.dumps(v)
  94
+
  95
+    @staticmethod
  96
+    def loads(v):
  97
+        ''' loads value  '''
  98
+        if isinstance(v, unicode):
  99
+            v = str(v)
  100
+        return pickle.loads(v)
  101
+
  102
+# -----------------------------------------------------------------
  103
+# CompressedJsonSerializer class
  104
+# -----------------------------------------------------------------
  105
+
  106
+class CompressedJsonSerializer(object):
  107
+    ''' CompressedJsonSerializer '''
  108
+
  109
+    @staticmethod
  110
+    def dumps(v):
  111
+        ''' dumps value '''
  112
+        return zlib.compress(json_encode(v))
  113
+
  114
+    @staticmethod
  115
+    def loads(v):
  116
+        ''' loads value  '''
  117
+        return json_decode(zlib.decompress(v))
  118
+
  119
+# -----------------------------------------------------------------
  120
+# SERIALIZERS 
  121
+# -----------------------------------------------------------------
  122
+
  123
+SERIALIZERS = {
  124
+    'pickle': cPickleSerializer,
  125
+    'completed_json': CompressedJsonSerializer,
  126
+}
  127
+
67 128
 
68 129
 # -----------------------------------------------------------------
69 130
 # KVLite utils
70 131
 # -----------------------------------------------------------------
71  
-def open(uri, serializer=pickle):
  132
+def open(uri, serializer=cPickleSerializer):
72 133
     ''' 
73 134
     open collection by URI, 
74 135
     if collection does not exist kvlite will try to create it
@@ -88,7 +149,7 @@ def open(uri, serializer=pickle):
88 149
     params = manager.parse_uri(uri)
89 150
     if params['collection'] not in manager.collections():
90 151
         manager.create(params['collection'])
91  
-    return manager.get_collection(params['collection'])(uri, serializer)
  152
+    return manager.collection_class(manager.connection, params['collection'], serializer)
92 153
 
93 154
 def remove(uri):
94 155
     '''
@@ -127,29 +188,39 @@ def __init__(self, uri):
127 188
         backend, rest_uri = uri.split('://')
128 189
         if backend in SUPPORTED_BACKENDS:
129 190
             if backend == 'mysql':
130  
-                self.manager = MysqlCollectionManager
  191
+                self.backend_manager = MysqlCollectionManager(uri)
131 192
             elif backend == 'sqlite':
132  
-                self.manager = SqliteCollectionManager
  193
+                self.backend_manager = SqliteCollectionManager(uri)
133 194
             else:
134 195
                 raise NotImplementedError()
135 196
         else:
136 197
             raise RuntimeError('Unknown backend: {}'.format(backend))
137 198
 
  199
+    def parse_uri(self, uri):
  200
+        ''' parse_uri '''
  201
+        return self.backend_manager.parse_uri(uri)
  202
+
138 203
     def create(self, name):
139 204
         ''' create collection '''
140  
-        self.manager.create(name)
  205
+        self.backend_manager.create(name)
141 206
     
142  
-    def get_collection(self, name):
  207
+    @property
  208
+    def collection_class(self):
143 209
         ''' return object MysqlCollection or SqliteCollection '''
144  
-        return self.manager.get_collection(name)
  210
+        return self.backend_manager.collection_class
  211
+    
  212
+    @property
  213
+    def connection(self):
  214
+        ''' return reference to backend connection '''
  215
+        return self.backend_manager.connection
145 216
     
146 217
     def collections(self):
147 218
         ''' return list of collections '''
148  
-        return self.manager.collections()
  219
+        return self.backend_manager.collections()
149 220
     
150 221
     def remove(self, name):
151 222
         ''' remove collection '''
152  
-        self.manager.remove(name)
  223
+        self.backend_manager.remove(name)
153 224
 
154 225
 # -----------------------------------------------------------------
155 226
 # MysqlCollectionManager class
@@ -160,6 +231,7 @@ class MysqlCollectionManager(object):
160 231
     def __init__(self, uri):
161 232
         
162 233
         params = self.parse_uri(uri) 
  234
+        
163 235
         try:
164 236
             self.__conn = MySQLdb.connect(
165 237
                                 host=params['host'], port = params['port'], 
@@ -194,7 +266,7 @@ def parse_uri(uri):
194 266
             parsed_uri['db'] = rest_uri
195 267
             parsed_uri['collection'] = None
196 268
         return parsed_uri
197  
-    
  269
+        
198 270
     def create(self, name):
199 271
         ''' create collection '''
200 272
 
@@ -207,9 +279,15 @@ def create(self, name):
207 279
         self.__cursor.execute(SQL_CREATE_TABLE % name)
208 280
         self.__conn.commit()
209 281
 
210  
-    def get_collection(self, name):
  282
+    @property
  283
+    def connection(self):
  284
+        ''' return connection '''
  285
+        return self.__conn
  286
+
  287
+    @property
  288
+    def collection_class(self):
211 289
         ''' return MysqlCollection object'''
212  
-        return MysqlCollection(self.__conn, name)
  290
+        return MysqlCollection
213 291
     
214 292
     def collections(self):
215 293
         ''' return collection list'''
@@ -235,11 +313,13 @@ def close(self):
235 313
 class MysqlCollection(object):
236 314
     ''' Mysql Connection '''
237 315
 
238  
-    def __init__(self, connection, collection_name):
  316
+    def __init__(self, connection, collection_name, serializer=cPickleSerializer):
239 317
 
240 318
         self.__conn = connection
241 319
         self.__cursor = self.__conn.cursor()
242 320
         self.__collection = collection_name
  321
+        self.__serializer = serializer
  322
+
243 323
         self.__uuid_cache = list()
244 324
 
245 325
     def get_uuid(self):
@@ -271,9 +351,9 @@ def __get_many(self):
271 351
                 rowid = r[0]
272 352
                 k = binascii.b2a_hex(r[1])
273 353
                 try:
274  
-                    v = self.unpack(r[2])
  354
+                    v = self.__serializer.loads(r[2])
275 355
                 except Exception, err:
276  
-                    raise ValueUnpackError('key %s, %s' % (k, err))
  356
+                    raise RuntimeError('key %s, %s' % (k, err))
277 357
                 yield (k, v)
278 358
 
279 359
     def get(self, k=None):
@@ -292,9 +372,9 @@ def get(self, k=None):
292 372
             result = self.__cursor.fetchone()
293 373
             if result:
294 374
                 try:
295  
-                    v = self.unpack(result[1])
  375
+                    v = self.__serializer.loads(result[1])
296 376
                 except Exception, err:
297  
-                    raise ValueUnpackError('key %s, %s' % (k, err))
  377
+                    raise RuntimeError('key %s, %s' % (k, err))
298 378
                 return (binascii.b2a_hex(result[0]), v)
299 379
             else:
300 380
                 return (None, None)
@@ -303,11 +383,14 @@ def get(self, k=None):
303 383
 
304 384
     def put(self, k, v):
305 385
         ''' put document in collection '''
  386
+        
  387
+        # TODO many k/v in put()
  388
+        
306 389
         if len(k) > 40:
307 390
             raise RuntimeError('The length of key is more than 40 bytes')
308 391
         SQL_INSERT = 'INSERT INTO %s (k,v) ' % self.__collection
309 392
         SQL_INSERT += 'VALUES (%s,%s) ON DUPLICATE KEY UPDATE v=%s;;'
310  
-        v = self.pack(v)
  393
+        v = self.__serializer.dumps(v)
311 394
         try:
312 395
             self.__cursor.execute(SQL_INSERT, (binascii.a2b_hex(k), v, v))
313 396
         except TypeError, err:
@@ -337,6 +420,7 @@ def keys(self):
337 420
                 k = binascii.b2a_hex(r[1])
338 421
                 yield k
339 422
 
  423
+    @property
340 424
     def count(self):
341 425
         ''' return amount of documents in collection'''
342 426
         self.__cursor.execute('SELECT count(*) FROM %s;' % self.__collection)
@@ -354,50 +438,183 @@ def close(self):
354 438
 # -----------------------------------------------------------------
355 439
 class SqliteCollectionManager(object):
356 440
     ''' Sqlite Collection Manager '''
357  
-    pass
  441
+    def __init__(self, uri):
  442
+        
  443
+        params = self.parse_uri(uri) 
  444
+        
  445
+        self.__conn = sqlite3.connect(params['db'])       
  446
+        self.__cursor = self.__conn.cursor()
  447
+
  448
+    @staticmethod
  449
+    def parse_uri(uri):
  450
+        '''parse URI 
  451
+        
  452
+        return driver, database, collection
  453
+        '''
  454
+        parsed_uri = dict()
  455
+        parsed_uri['backend'], rest_uri = uri.split('://', 1)
  456
+        if ':' in rest_uri:
  457
+            parsed_uri['db'], parsed_uri['collection'] = rest_uri.split(':',1)
  458
+        else:
  459
+            parsed_uri['db'] = rest_uri
  460
+            parsed_uri['collection'] = None
  461
+        if parsed_uri['db'] == 'memory':
  462
+            parsed_uri['db'] = ':memory:'
  463
+        return parsed_uri
  464
+
  465
+    @property
  466
+    def connection(self):
  467
+        ''' return connection '''
  468
+        return self.__conn
  469
+
  470
+    @property
  471
+    def collection_class(self):
  472
+        ''' return SqliteCollection object'''
  473
+        return SqliteCollection
  474
+
  475
+    def collections(self):
  476
+        ''' return collection list'''
  477
+
  478
+        self.__cursor.execute('SELECT name FROM sqlite_master WHERE type="table";')
  479
+        return [t[0] for t in self.__cursor.fetchall()]
358 480
 
  481
+    def create(self, name):
  482
+        ''' create collection '''
  483
+
  484
+        SQL_CREATE_TABLE = '''CREATE TABLE IF NOT EXISTS %s (
  485
+                                k NOT NULL, v, UNIQUE (k) );'''
  486
+
  487
+        self.__cursor.execute(SQL_CREATE_TABLE % name)
  488
+        self.__conn.commit()
  489
+
  490
+    def remove(self, name):
  491
+        ''' remove collection '''
  492
+        if name in self.collections():
  493
+            self.__cursor.execute('DROP TABLE %s;' % name)
  494
+            self.__conn.commit()
  495
+        else:
  496
+            raise RuntimeError('No collection with name: {}'.format(name))
  497
+
  498
+    def close(self):
  499
+        ''' close connection to database '''
  500
+        self.__conn.close()
  501
+    
359 502
 # -----------------------------------------------------------------
360 503
 # SqliteCollection class
361 504
 # -----------------------------------------------------------------
362 505
 class SqliteCollection(object):
363 506
     ''' Sqlite Collection'''
364 507
     
365  
-    def __init__(self, uri):
366  
-        raise NotImplementedError('SqliteCollection is not implemented yet')
  508
+    def __init__(self, connection, collection_name, serializer=cPickleSerializer):
367 509
 
368  
-class Collection(object):
369  
-    ''' 
370  
-    kvlite2 collection
  510
+        self.__conn = connection
  511
+        self.__cursor = self.__conn.cursor()
  512
+        self.__collection = collection_name
  513
+        self.__serializer = serializer
371 514
 
372  
-    A collection is a group of documents stored in kvlite2, 
373  
-    and can be thought of as roughly the equivalent of a 
374  
-    table in a relational database.
  515
+        self.__uuid_cache = list()
375 516
 
376  
-    '''
377  
-    def __init__(self, db_uri):
378  
-        '''
379  
-        db_uri - URI to databases, 
380  
-        URI format: driver://username:passwd@host[:port]/database.collection
381  
-        '''
382  
-        params = parse_uri(db_uri)
383  
-        self.__conn = MySQLdb.connect(host=params['host'], port = params['port'], 
384  
-                        user=params['usr'], passwd=params['pwd'], db=params['db'])
385  
-        self.__collection = params['coll']
386  
-        self.__cursor = self.__conn.cursor()
387  
-        self.__uuids = []
  517
+    def get_uuid(self):
  518
+        """ return id based on uuid """
388 519
 
389  
-    def pack(self, v):
390  
-        ''' pack value 
  520
+        if not self.__uuid_cache:
  521
+            for uuid in get_uuid():
  522
+                self.__uuid_cache.append(uuid)
  523
+        return self.__uuid_cache.pop()
  524
+
  525
+    def put(self, k, v):
  526
+        ''' put document in collection '''
391 527
         
392  
-        Note: before pack the value it's better to encode it by base64
393  
-        '''
394  
-        return zlib.compress(json_encode(v))
  528
+        # TODO many k/v in put()
  529
+        
  530
+        if len(k) > 40:
  531
+            raise RuntimeError('The length of key is more than 40 bytes')
  532
+        SQL_INSERT = 'INSERT OR REPLACE INTO %s (k,v) ' % self.__collection
  533
+        SQL_INSERT += 'VALUES (?,?)'
  534
+        v = self.__serializer.dumps(v)
  535
+        try:
  536
+            self.__cursor.execute(SQL_INSERT, (k, v))
  537
+        except TypeError, err:
  538
+            raise RuntimeError(err)
395 539
 
396  
-    def unpack(self, v):
397  
-        ''' unpack value 
  540
+    def __get_many(self):
  541
+        ''' return all docs '''
  542
+        rowid = 0
  543
+        while True:
  544
+            SQL_SELECT_MANY = 'SELECT rowid, k,v FROM %s WHERE rowid > %d LIMIT 1000 ;' % (self.__collection, rowid)
  545
+            self.__cursor.execute(SQL_SELECT_MANY)
  546
+            result = self.__cursor.fetchall()
  547
+            if not result:
  548
+                break
  549
+            for r in result:
  550
+                rowid = r[0]
  551
+                k = r[1]
  552
+                try:
  553
+                    v = self.__serializer.loads(r[2])
  554
+                except Exception, err:
  555
+                    raise RuntimeError('key %s, %s' % (k, err))
  556
+                yield (k, v)
  557
+
  558
+    def get(self, k=None):
  559
+        ''' 
  560
+        return document by key from collection 
  561
+        return documents if key is not defined
398 562
         '''
399  
-        return json_decode(zlib.decompress(v))
  563
+        if k:
  564
+            if len(k) > 40:
  565
+                raise RuntimeError('The key length is more than 40 bytes')
  566
+            SQL = 'SELECT k,v FROM %s WHERE k = ?;' % self.__collection
  567
+            try:
  568
+                self.__cursor.execute(SQL, (k,))
  569
+            except TypeError, err:
  570
+                raise WronKeyValue(err)
  571
+            result = self.__cursor.fetchone()
  572
+            if result:
  573
+                try:
  574
+                    v = self.__serializer.loads(result[1])
  575
+                except Exception, err:
  576
+                    raise RuntimeError('key %s, %s' % (k, err))
  577
+                return (result[0], v)
  578
+            else:
  579
+                return (None, None)
  580
+        else:
  581
+            return self.__get_many()            
  582
+
  583
+    def keys(self):
  584
+        ''' return document keys in collection'''
  585
+        rowid = 0
  586
+        while True:
  587
+            SQL_SELECT_MANY = 'SELECT rowid, k FROM %s WHERE rowid > %d LIMIT 1000 ;' % (self.__collection, rowid)
  588
+            self.__cursor.execute(SQL_SELECT_MANY)
  589
+            result = self.__cursor.fetchall()
  590
+            if not result:
  591
+                break
  592
+            for r in result:
  593
+                rowid = r[0]
  594
+                yield r[1]
400 595
 
  596
+    def delete(self, k):
  597
+        ''' delete document by k '''
  598
+        if len(k) > 40:
  599
+            raise RuntimeError('The key length is more than 40 bytes')
  600
+        SQL_DELETE = '''DELETE FROM %s WHERE k = ?;''' % self.__collection
  601
+        try:
  602
+            self.__cursor.execute(SQL_DELETE, (k,))
  603
+        except TypeError, err:
  604
+            raise WronKeyValue(err)
  605
+
  606
+    @property
  607
+    def count(self):
  608
+        ''' return amount of documents in collection'''
  609
+        self.__cursor.execute('SELECT count(*) FROM %s;' % self.__collection)
  610
+        return int(self.__cursor.fetchone()[0])
  611
+
  612
+    def commit(self):
  613
+        self.__conn.commit()
  614
+
  615
+    def close(self):
  616
+        ''' close connection to database '''
  617
+        self.__conn.close()
401 618
                     
402 619
 # -----------------------------------------------------------------
403 620
 # Console class

0 notes on commit 212f56d

Please sign in to comment.
Something went wrong with that request. Please try again.