izzylib/IzzyLib/database_tiny.py

import base64
import json
import operator
import queue
import threading
import time

import tinydb
import tinydb_serialization

from . import misc


class AwaitingResult(object):
	pass


class DataBase(tinydb.TinyDB):
	def __init__(self, dbfile: misc.Path, queue_limit: int=64, serializers: list=[]):
		options = {
			'indent': 2,
			'separators': (',', ': '),
		}

		serialization = tinydb_serialization.SerializationMiddleware(ThreadSupport)
		serialization.register_serializer(tinydb_serialization.serializers.DateTimeSerializer(), 'TinyDate')
		serialization.register_serializer(ByteSerializer(), 'TinyBytes')
		serialization.register_serializer(PathSerialize(), 'IzzyLibPath')
		serialization.register_serializer(DotDictSerialize(), 'IzzyLibDotDict')

		for serializer in serializers:
			serialization.register_serializer(serializer(), serializer.__class__.__name__)

		options['storage'] = serialization

		super().__init__(dbfile, **options)

		self.dbfile = dbfile
		self.queue = queue.Queue(maxsize=queue_limit)
		self.sessions = {}


	@property
	def session(self):
		return TinySession(self)


	def fetch(self, table, single=True, orderby=None, reverse=False, **kwargs):
		query = tinydb.Query()

		if not kwargs:
			rows = TinyRows(self, self.table(table).all())
			single = False

		else:
			rows = TinyRows(self, self.table(table).search(query.fragment(kwargs)))

		if single:
			return rows[0] if rows else None

		return rows if not orderby else sorted(rows, key=operator.itemgetter(orderby), reverse=reverse)


	def insert(self, table, row=None, rowid=None, **kwargs):
		if row:
			rowid = row.doc_id

		elif not rowid:
			row = self.fetch(table, **kwargs)
			rowid = row.doc_id if row else None

		if rowid:
			return self.table(table).update(kwargs, doc_ids=[rowid])

		return self.table(table).insert(kwargs)


	def remove(self, table, row=None, rowid=None, **kwargs):
		query = tinydb.Query()

		if row or rowid:
			rowid = rowid or row.doc_id
			return self.table(table).remove(doc_ids=[row.doc_id])

		return self.table(table).remove(query.fragment(kwargs))


def TinyRows(db, rows):
	return [TinyRow(db, row) for row in rows]


class TinyRow(misc.DotDict):
	def __init__(self, db, row):
		super().__init({'id': row.doc_id})
		super().update({k: v for k,v in row.items()})


	def update(self, data={}):
		db.update(rowid=self.id, **data)
		super().update(data)


	def remove(self):
		db.remove(rowid=self.id)


class ThreadSupport(tinydb.storages.JSONStorage):
	def __init__(self, filename, *args, **kwargs):
		super().__init__(filename, *args, **kwargs)

		self._thread_event = threading.Event()
		self._shutdown = False
		self._results = {}
		self._queue = queue.Queue()

		self._lock = threading.Lock()
		self._thread = threading.Thread(target=self.process_queue)
		self._thread.daemon = True
		self._thread.start()

	## send all storage commands to the queue
	def read(self):
		data = self.queue_put('read')
		return data


	def write(self, data):
		self.queue_put('write', data)


	def close(self):
		self.queue_put('close')


	def get_action(self, action):
		return getattr(super(), action)


	def get_result(self, qid):
		with self._lock:
			return self._results[qid]


	def set_result(self, qid, data=AwaitingResult):
		with self._lock:
			self._results[qid] = data


	def pop_result(self, qid):
		with self._lock:
			return self._result.pop(qid)


	## queue
	def process_queue(self):
		while not self._thread_event.is_set():
			if not self._queue.empty():
				qid, action, args, kwargs = self._queue.get(block=False)

				if qid not in self._results:
					self.set_result(qid)
					if action == 'close':
						self._shutdown = True

					func = self.get_action(action)

					if action == 'read':
						self.set_result(qid, func(*args, **kwargs))

			else:
				time.sleep(0.1)

				if self._shutdown:
					self.get_action('close')()
					return


	def queue_put(self, func, *args, **kwargs):
		if self._shutdown:
			logging.error('Storage has been closed. Refusing to send more commands')
			return

		qid = misc.RandomGen()

		self._queue.put((qid, func, args, kwargs))

		if func != 'read':
			return

		sleep_time = 0.0

		while self.get_result(qid) == AwaitingResult:
			time.sleep(0.1)
			sleep_time += 0.1

			if sleep_time >= 5.0:
				raise TimeoutError(f'Timeout on "{func}" with args: {args}, {kwargs}')

		result = self.pop_result(qid)
		print(result)
		return result


class ByteSerializer(tinydb_serialization.Serializer):
	OBJ_CLASS = bytes

	def encode(self, obj):
		#print('encode', self.__class__.__name__, obj)
		return base64.b64encode(obj).decode('ascii')

	def decode(self, obj):
		#print('decode', self.__class__.__name__, obj)
		return base64.b64decode(obj)


class DictSerializer(tinydb_serialization.Serializer):
	OBJ_CLASS = dict

	def encode(self, obj):
		#print('encode', self.__class__.__name__, obj)
		return json.dumps(obj)

	def decode(self, obj):
		#print('decode', self.__class__.__name__, obj)
		return json.loads(obj)


class DotDictSerialize(tinydb_serialization.Serializer):
	OBJ_CLASS = misc.DotDict

	def encode(self, obj):
		#print('encode', self.__class__.__name__, obj)
		return obj.toJson()

	def decode(self, obj):
		#print('decode', self.__class__.__name__, obj)
		return misc.DotDict(obj)


class PathSerialize(tinydb_serialization.Serializer):
	OBJ_CLASS = misc.Path

	def encode(self, obj):
		#print('encode', self.__class__.__name__, obj)
		return obj.str()

	def decode(self, obj):
		#print('decode', self.__class__.__name__, obj)
		return misc.Path(obj)