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

import ast
import csv
import json
import logging

from playhouse.sqlite_ext import JSONField

from .models import Archive, Autotests, Log, Recipes, Session, Users, db

models_reference = {
    "archive": Archive,
    "autotests": Autotests,
    "log": Log,
    "recipes": Recipes,
    "users": Users,
}

db.connect()
db.create_tables(list(models_reference.values()))

log = logging.getLogger("db")


def init_db():
    tables = db.get_tables()
    tables.sort()
    for table in tables:
        count = 0
        try:
            with open("src/lib/db/imports/{}.csv".format(table), "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")
# register test recipe
Recipes.replace(id=0, name="TEST", spec={
    # recipe
    "client": "TEST_CLIENT",
    "part_number": "TEST_PART_NUMBER",
    "station": "TEST_STATION",
    "description": "TEST_DESCRIPTION",
    # pre-filling
    "pre_filling_time": 5,
    "pre_filling_pressure": 3000,
    # filling
    "filling_time": 5,
    "settling_time": 10,
    "settling_pressure_min_percent": 10,
    "settling_pressure_max_percent": 10,
    # test
    "test_time": 10,
    "test_pressure_min_delta": 3.00,
    "test_pressure": 3000,
    "test_pressure_max_delta": 0.25,
    # flush
    "flush_time": 2,
    "flush_pressure": 5,
    # vision
    "vision_recipe": 0,
}, archived=False).execute()

if True:
    # crud_db must be imported after db and models_reference are available
    from .crud_db import Crud_DB