st-ten-1/src/lib/db/__init__.py

66 lines
1.8 KiB
Python

import ast
import csv
import json
import logging
import re
from playhouse.sqlite_ext import JSONField
from .models import Archive, Log, Recipes, Session, Users, db
models_reference = {
"archive": Archive,
"log": Log,
"recipes": Recipes,
"users": Users,
}
db.connect()
db.create_tables(list(models_reference.values()))
log = logging.getLogger("db")
@db.atomic()
def init_db():
tables = db.get_tables()
tables.sort()
for table in tables:
count = 0
try:
with open(f"src/lib/db/imports/{table}.csv", "r") as f:
table = models_reference[table]
fields = list(table._meta.fields)
log.info(f"importing {table._meta.table_name}")
reader = csv.DictReader(f)
for row in reader:
obj = {}
for field in fields:
if type(table._meta.fields[field]) is JSONField:
obj[field] = json.loads(row[field])
else:
try:
obj[field] = ast.literal_eval(row[field])
except (SyntaxError, ValueError):
obj[field] = row[field]
table.insert(**obj).on_conflict_replace().execute()
count += 1
log.info(f"{table._meta.table_name}: imported {count} rows.")
except FileNotFoundError:
pass
init_db()
# register or reset admin
admin = Users.get_user("ADMIN")
if admin is None or not admin.is_admin:
Users.register(username="ADMIN", password="123123", roles=["admin"])
# register or reset user
Users.register(username="USER", password="user")
if True:
# crud_db must be imported after db and models_reference are available
from .crud_db import Crud_DB