basis/components/database/__init__.py
apeters 519dbc73c6 Cat.
Signed-off-by: apeters <apeters@korves.net>
2025-06-03 11:40:56 +00:00

93 lines
2.5 KiB
Python

import contextvars
import os
import shutil
import json
from aiotinydb import AIOTinyDB as TinyDB
from aiotinydb.storage import AIOStorage
from copy import copy
from tinydb import Query
from tinydb.table import Document
from components.utils import is_path_within_cwd
__all__ = [
"TinyDB",
"Query",
"Document",
"TINYDB_PARAMS",
"IN_MEMORY_DB",
"CTX_TICKET",
"evaluate_db_params",
"dbcommit",
]
TinyDB.DEFAULT_TABLE_KWARGS = {"cache_size": 0}
TINYDB_PARAMS = {
"filename": "database/main",
"indent": 2,
"sort_keys": True,
}
SYSTEM_CACHE_ID = "00000000-0000-0000-0000-000000000000"
ANONYMOUS_CACHE_ID = "99999999-9999-9999-9999-999999999999"
IN_MEMORY_DB = dict()
IN_MEMORY_DB["SESSION_VALIDATED"] = dict()
IN_MEMORY_DB["WS_CONNECTIONS"] = dict()
IN_MEMORY_DB["CACHE"] = {
"MODELS": {
SYSTEM_CACHE_ID: dict(),
ANONYMOUS_CACHE_ID: dict(),
},
"FORMS": dict(),
}
IN_MEMORY_DB["APP_LOGS_FULL_PULL"] = dict()
IN_MEMORY_DB["PROMOTE_USERS"] = set()
IN_MEMORY_DB["TOKENS"] = {
"REGISTER": dict(),
"LOGIN": dict(),
}
CTX_TICKET = contextvars.ContextVar("CTX_TICKET", default=None)
if not os.path.exists("database/main") or os.path.getsize("database/main") == 0:
os.makedirs(os.path.dirname("database/main"), exist_ok=True)
with open("database/main", "w") as f:
f.write("{}")
os.chmod("database/main", 0o600)
def evaluate_db_params(ticket: str | None = None):
db_params = copy(TINYDB_PARAMS)
transaction_file = (
f"database/main.{ticket}"
if ticket
else f"database/main.{CTX_TICKET.get() or ''}".rstrip(".")
)
if transaction_file != "database/main" and is_path_within_cwd(transaction_file):
if not os.path.exists(transaction_file):
shutil.copy("database/main", transaction_file)
db_params["filename"] = transaction_file
return db_params
async def dbcommit(commit_tables: set, ticket: str | None = None) -> None:
assert commit_tables
db_params = evaluate_db_params(ticket)
with open(db_params["filename"], "r") as f:
modified_db = json.load(f)
async with TinyDB(**TINYDB_PARAMS) as db:
current_db = json.load(db._storage._handle)
for t in commit_tables:
current_db[t] = modified_db[t]
db._storage._handle.seek(0)
serialized = json.dumps(current_db, **db._storage.kwargs)
db._storage._handle.write(serialized)
db._storage._handle.flush()
db._storage._handle.truncate()
os.unlink(db_params["filename"])