# -*- coding: utf-8 -*- from collections import deque import inspect import os import uuid import sys from peewee import CharField, BlobField, IntegerField, \ ForeignKeyField, Model, BooleanField, TextField, Field, Database from solar.dblayer.model import clear_cache from threading import RLock # msgpack is way faster but less readable # using json for easier debug import json encoder = json.dumps def wrapped_loads(data, *args, **kwargs): if not isinstance(data, basestring): data = str(data) return json.loads(data, *args, **kwargs) decoder = wrapped_loads class _DataField(BlobField): def db_value(self, value): return super(_DataField, self).db_value(encoder(value)) def python_value(self, value): return decoder(super(_DataField, self).python_value(value)) class _LinksField(_DataField): def db_value(self, value): return super(_LinksField, self).db_value(list(value)) def python_value(self, value): ret = super(_LinksField, self).python_value(value) return [tuple(e) for e in ret] class _SqlBucket(Model): def __init__(self, *args, **kwargs): self._new = kwargs.pop('_new', False) ed = kwargs.pop('encoded_data', None) if ed: self.encoded_data = ed if 'data' not in kwargs: kwargs['data'] = {} super(_SqlBucket, self).__init__(*args, **kwargs) key = CharField(primary_key=True, null=False) data = _DataField(null=False) vclock = CharField(max_length=32, null=False) links = _LinksField(null=False, default=list) @property def encoded_data(self): return self.data.get('_encoded_data') @encoded_data.setter def encoded_data(self, value): self.data['_encoded_data'] = value def save(self, force_insert=False, only=None): if self._new: force_insert = True self._new = False ret = super(_SqlBucket, self).save(force_insert, only) return ret @property def sql_session(self): return self.bucket.sql_session class FieldWrp(object): def __init__(self, name): self.name = name def __get__(self, instance, owner): return getattr(instance._sql_bucket_obj, self.name) def __set__(self, instance, value): setattr(instance._sql_bucket_obj, self.name, value) class _SqlIdx(Model): name = CharField(null=False, index=True) value = CharField(null=False, index=True) class RiakObj(object): key = FieldWrp('key') data = FieldWrp('data') vclock = FieldWrp('vclock') links = FieldWrp('links') encoded_data = FieldWrp('encoded_data') def __init__(self, sql_bucket_obj, new=False): self._sql_bucket_obj = sql_bucket_obj self.new = sql_bucket_obj._new self.fetch_indexes() @property def sql_session(self): return self._sql_bucket_obj.sql_session @property def bucket(self): return self._sql_bucket_obj.bucket @property def indexes(self): self.fetch_indexes() return self._indexes def fetch_indexes(self): if not hasattr(self, '_indexes'): idxes = self.bucket._sql_idx.select().where( self.bucket._sql_idx.key == self.key) self._indexes = set((idx.name, idx.value) for idx in idxes) @indexes.setter def indexes(self, value): assert isinstance(value, set) self._indexes = value def _save_indexes(self): # TODO: possible optimization # update only what's needed # don't delete all at first q = self.bucket._sql_idx.delete().where( self.bucket._sql_idx.key == self.key) q.execute() for iname, ival in self.indexes: idx = self.bucket._sql_idx(key=self.key, name=iname, value=ival) idx.save() def add_index(self, field, value): self.indexes.add((field, value)) return self def set_index(self, field, value): to_rem = set((x for x in self.indexes if x[0] == field)) self.indexes.difference_update(to_rem) return self.add_index(field, value) def remove_index(self, field=None, value=None): if field is None and value is None: # q = self.bucket._sql_idx.delete().where( # self.bucket._sql_idx.key == self.key) # q.execute() self.indexes.clear() elif field is not None and value is None: # q = self.bucket._sql_idx.delete().where( # (self.bucket._sql_idx.key == self.key) & # (self.bucket._sql_idx.name == field)) # q.execute() to_rem = set((x for x in self.indexes if x[0] == field)) self.indexes.difference_update(to_rem) elif field is not None and value is not None: # q = self.bucket._sql_idx.delete().where( # (self.bucket._sql_idx.key == self.key) & # (self.bucket._sql_idx.name == field) & # (self.bucket._sql_idx.value == value)) # q.execute() to_rem = set((x for x in self.indexes if x[0] == field and x[1] == value)) self.indexes.difference_update(to_rem) return self def store(self, return_body=True): self.vclock = uuid.uuid4().hex assert self._sql_bucket_obj is not None self._sql_bucket_obj.save() self._save_indexes() return self def delete(self): self._sql_bucket_obj.delete() return self @property def exists(self): return not self.new def get_link(self, tag): return next(x[1] for x in self.links if x[2] == tag) def set_link(self, obj, tag=None): if isinstance(obj, tuple): newlink = obj else: newlink = (obj.bucket.name, obj.key, tag) multi = [x for x in self.links if x[0:1] == newlink[0:1]] for item in multi: self.links.remove(item) self.links.append(newlink) return self def del_link(self, obj=None, tag=None): assert obj is not None or tag is not None if tag is not None: links = [x for x in self.links if x[2] != tag] else: links = self.links if obj is not None: if not isinstance(obj, tuple): obj = (obj.bucket.name, obj.key, tag) links = [x for x in links if x[0:1] == obj[0:1]] self.links = links return self class IndexPage(object): def __init__(self, index, results, return_terms, max_results, continuation): self.max_results = max_results self.index = index if not return_terms: self.results = tuple(x[0] for x in results) else: self.results = tuple(results) if not max_results or not self.results: self.continuation = None else: self.continuation = str(continuation + len(self.results)) self.return_terms = return_terms def __len__(self): return len(self.results) def __getitem__(self, item): return self.results[item] class Bucket(object): def __init__(self, name, client): self.client = client table_name = "bucket_%s" % name.lower() self.name = table_name idx_table_name = 'idx_%s' % name.lower() class ModelMeta: db_table = table_name database = self.client.sql_session self._sql_model = type(table_name, (_SqlBucket,), {'Meta': ModelMeta, 'bucket': self}) _idx_key = ForeignKeyField(self._sql_model, null=False, index=True) class IdxMeta: db_table = idx_table_name database = self.client.sql_session self._sql_idx = type(idx_table_name, (_SqlIdx,), {'Meta': IdxMeta, 'bucket': self, 'key': _idx_key}) def search(self, q, rows=10, start=0, sort=''): raise NotImplementedError() def create_search(self, index): raise NotImplementedError() def set_property(self, name, value): return def get_properties(self): return {'search_index': False} def get(self, key): try: ret = self._sql_model.get(self._sql_model.key == key) except self._sql_model.DoesNotExist: ret = None new = ret is None if new: ret = self._sql_model(key=key, _new=new) return RiakObj(ret, new) def delete(self, data, *args, **kwargs): if isinstance(data, basestring): key = data else: key = data.key self._sql_model.delete().where(self._sql_model.key == key).execute() self._sql_idx.delete().where(self._sql_idx.key == key).execute() return self def new(self, key, data=None, encoded_data=None, **kwargs): if key is not None: try: ret = self._sql_model.get(self._sql_model.key == key) except self._sql_model.DoesNotExist: ret = None new = ret is None else: key = uuid.uuid4().hex new = True if new: ret = self._sql_model(key=key, _new=new) ret.key = key ret.data = data if data is not None else {} if encoded_data: ret.encoded_data = encoded_data ret.links = [] ret.vclock = "new" return RiakObj(ret, new) def get_index(self, index, startkey, endkey=None, return_terms=None, max_results=None, continuation=None, timeout=None, fmt=None, term_regex=None): if startkey and endkey is None: endkey = startkey if startkey > endkey: startkey, endkey = endkey, startkey if index == '$key': if return_terms: q = self._sql_model.select( self._sql_model.value, self._sql_model.key) else: q = self._sql_model.select(self._sql_model.key) q = q.where( self._sql_model.key >= startkey, self._sql_model.key <= endkey ).order_by(self._sql_model.key) elif index == '$bucket': if return_terms: q = self._sql_model.select( self._sql_model.value, self._sql_model.key) else: q = self._sql_model.select(self._sql_model.key) if not startkey == '_' and endkey == '_': q = q.where( self._sql_model.key >= startkey, self._sql_model.key <= endkey ) else: if return_terms: q = self._sql_idx.select( self._sql_idx.value, self._sql_idx.key) else: q = self._sql_idx.select(self._sql_idx.key) q = q.where( self._sql_idx.name == index, self._sql_idx.value >= startkey, self._sql_idx.value <= endkey ).order_by(self._sql_idx.value) max_results = int(max_results or 0) continuation = int(continuation or 0) if max_results: q = q.limit(max_results) if continuation: q = q.offset(continuation) q = q.tuples() return IndexPage(index, q, return_terms, max_results, continuation) def multiget(self, keys): if not keys: return [] else: q = self._sql_model.select().where(self._sql_model.key << list(keys)) print q return map(RiakObj, list(q)) @property def sql_session(self): return self.client.sql_session class SqlClient(object): block = RLock() search_dir = None def __init__(self, *args, **kwargs): db_class_str = kwargs.pop("db_class", 'SqliteDatabase') try: mod, fromlist = db_class_str.split('.') except ValueError: mod = 'peewee' fromlist = db_class_str __import__(mod, fromlist=[fromlist]) db_class = getattr(sys.modules[mod], fromlist) session = db_class(*args, **kwargs) self._sql_session = session self.buckets = {} def bucket(self, name): with self.block: if name not in self.buckets: b = Bucket(name, self) b._sql_model.create_table(fail_silently=True) b._sql_idx.create_table(fail_silently=True) self.buckets[name] = b return self.buckets[name] @property def sql_session(self): return self._sql_session def session_start(self): clear_cache() sess = self._sql_session sess.begin() def session_end(self, result=True): sess = self._sql_session if result: sess.commit() else: sess.rollback() clear_cache() def delete_all(self, cls): # naive way for SQL, we could delete whole table contents rst = cls.bucket.get_index('$bucket', startkey='_', max_results=100000).results for key in rst: cls.bucket.delete(key)