feat: added namespaced AddonStorage helper

Co-authored-by: f-tlm <f-tlm@users.noreply.github.com>
This commit is contained in:
əlemi 2022-08-13 14:14:57 +02:00
parent f5481ace65
commit 1499cab317
No known key found for this signature in database
GPG key ID: BBCBFE5D7244634E
3 changed files with 33 additions and 6 deletions

View file

@ -10,8 +10,8 @@ import inspect
from pathlib import Path from pathlib import Path
from importlib import import_module from importlib import import_module
import traceback import traceback
from typing import List, Type, Set, get_type_hints from typing import Type, Set, get_type_hints
from dataclasses import dataclass, MISSING, fields from dataclasses import MISSING, fields
from setproctitle import setproctitle from setproctitle import setproctitle

View file

@ -4,6 +4,8 @@ import logging
from typing import TYPE_CHECKING, Dict, Any, Optional, Union, List, Callable, get_type_hints, get_args, get_origin from typing import TYPE_CHECKING, Dict, Any, Optional, Union, List, Callable, get_type_hints, get_args, get_origin
from dataclasses import dataclass, MISSING, fields from dataclasses import dataclass, MISSING, fields
from treepuncher.storage import AddonStorage
from .scaffold import ConfigObject from .scaffold import ConfigObject
if TYPE_CHECKING: if TYPE_CHECKING:
@ -45,6 +47,7 @@ def parse_with_hint(val:str, hint:Any) -> Any:
class Addon: class Addon:
name: str name: str
config: ConfigObject config: ConfigObject
storage: AddonStorage
logger: logging.Logger logger: logging.Logger
_client: 'Treepuncher' _client: 'Treepuncher'
@ -81,6 +84,7 @@ class Addon:
else: # not really necessary since it's a dataclass but whatever else: # not really necessary since it's a dataclass but whatever
opts[field.name] = default opts[field.name] = default
self.config = self.Options(**opts) self.config = self.Options(**opts)
self.storage = client.storage.addon_storage(self.name)
self.logger = self._client.logger.getChild(self.name) self.logger = self._client.logger.getChild(self.name)
self.register() self.register()

View file

@ -20,11 +20,31 @@ class AuthenticatorState:
token : Dict[str, Any] token : Dict[str, Any]
legacy : bool = False legacy : bool = False
class AddonStorage:
db: sqlite3.Connection
name: str
def __init__(self, db:sqlite3.Connection, name:str):
self.db = db
self.name = name
self.db.cursor().execute('CREATE TABLE IF NOT EXISTS documents (name TEXT PRIMARY KEY, value TEXT)')
self.db.commit()
# fstrings in queries are evil but if you go to this length to fuck up you kinda deserve it :)
def get(self, key:str) -> Optional[Any]:
res = self.db.cursor().execute(f"SELECT * FROM documents_{self.name} WHERE name = ?", (key,)).fetchall()
return json.loads(res[0][1])
def put(self, key:str, val:Any) -> None:
cur = self.db.cursor()
cur.execute("DELETE FROM documents WHERE name = ?", (key,))
cur.execute(f"INSERT INTO documents_{self.name} VALUES (?, ?)", (key, json.dumps(val, default=str),))
self.db.commit()
class Storage: class Storage:
name : str name : str
db : sqlite3.Connection db : sqlite3.Connection
def __init__(self, name:str): def __init__(self, name:str):
self.name = name self.name = name
init = not os.path.isfile(name) init = not os.path.isfile(name)
@ -57,6 +77,9 @@ class Storage:
cur.execute('INSERT INTO authenticator VALUES (?, ?, ?)', (state.date.strftime(__DATE_FORMAT__), json.dumps(state.token), state.legacy)) cur.execute('INSERT INTO authenticator VALUES (?, ?, ?)', (state.date.strftime(__DATE_FORMAT__), json.dumps(state.token), state.legacy))
self.db.commit() self.db.commit()
def addon_storage(self, name:str) -> AddonStorage:
return AddonStorage(self.db, name)
def system(self) -> Optional[SystemState]: def system(self) -> Optional[SystemState]:
cur = self.db.cursor() cur = self.db.cursor()
val = cur.execute('SELECT * FROM system').fetchall() val = cur.execute('SELECT * FROM system').fetchall()
@ -78,10 +101,9 @@ class Storage:
token=json.loads(val[0][1]), token=json.loads(val[0][1]),
legacy=val[0][2] or False legacy=val[0][2] or False
) )
def get(self, key:str) -> Optional[Any]: def get(self, key:str) -> Optional[Any]:
cur = self.db.cursor() val = self.db.cursor().execute("SELECT * FROM documents WHERE name = ?", (key,)).fetchall()
val = cur.execute("SELECT * FROM documents WHERE name = ?", (key,)).fetchall()
return json.loads(val[0][1]) if val else None return json.loads(val[0][1]) if val else None
def put(self, key:str, val:Any) -> None: def put(self, key:str, val:Any) -> None:
@ -90,3 +112,4 @@ class Storage:
cur.execute("INSERT INTO documents VALUES (?, ?)", (key, json.dumps(val, default=str))) cur.execute("INSERT INTO documents VALUES (?, ?)", (key, json.dumps(val, default=str)))
self.db.commit() self.db.commit()