67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
import abc
|
|
import json
|
|
import sqlite3
|
|
|
|
|
|
class Beaker(abc.ABC):
|
|
def __init__(self, name: str, recipe):
|
|
self.name = name
|
|
self.recipe = recipe
|
|
|
|
def __repr__(self):
|
|
return f"Beaker({self.name})"
|
|
|
|
@abc.abstractmethod
|
|
def items(self):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def __len__(self):
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def add_item(self, item: dict, from_table=None, from_id=None) -> None:
|
|
pass
|
|
|
|
|
|
class TempBeaker(Beaker):
|
|
def __init__(self, name: str, recipe):
|
|
super().__init__(name, recipe)
|
|
self._items = []
|
|
|
|
def __len__(self):
|
|
return len(self._items)
|
|
|
|
def add_item(self, item: dict, from_table=None, from_id=None) -> None:
|
|
self._items.append((from_id, item))
|
|
|
|
def items(self):
|
|
yield from self._items
|
|
|
|
|
|
class SqliteBeaker(Beaker):
|
|
def __init__(self, name: str, recipe):
|
|
super().__init__(name, recipe)
|
|
# create table if it doesn't exist
|
|
self.cursor = self.recipe.db.cursor()
|
|
self.cursor.row_factory = sqlite3.Row
|
|
self.cursor.execute(
|
|
f"CREATE TABLE IF NOT EXISTS {self.name} (id INTEGER PRIMARY KEY, data JSON, from_table TEXT NULL, from_id INTEGER NULL)"
|
|
)
|
|
|
|
def items(self):
|
|
self.cursor.execute(f"SELECT id, data FROM {self.name}")
|
|
data = self.cursor.fetchall()
|
|
for item in data:
|
|
yield item["id"], json.loads(item["data"])
|
|
|
|
def __len__(self):
|
|
self.cursor.execute(f"SELECT COUNT(*) FROM {self.name}")
|
|
return self.cursor.fetchone()[0]
|
|
|
|
def add_item(self, item: dict, from_table=None, from_id=None) -> None:
|
|
self.cursor.execute(
|
|
f"INSERT INTO {self.name} (data) VALUES (?)", (json.dumps(item),)
|
|
)
|
|
self.recipe.db.commit()
|