cache.py

# -*- coding: utf-8 -*-
import json

from peewee import (
    BooleanField,
    CharField,
    Field,
    FloatField,
    ForeignKeyField,
    Model,
    SqliteDatabase,
    TextField,
    UUIDField,
)

db = SqliteDatabase(None)


class JSONField(Field):
    field_type = "text"

    def db_value(self, value):
        if value is None:
            return
        return json.dumps(value)

    def python_value(self, value):
        if value is None:
            return
        return json.loads(value)


class CachedElement(Model):
    id = UUIDField(primary_key=True)
    parent_id = UUIDField(null=True)
    type = CharField(max_length=50)
    polygon = JSONField(null=True)
    initial = BooleanField(default=False)
    worker_version_id = UUIDField(null=True)

    class Meta:
        database = db
        table_name = "elements"


class CachedTranscription(Model):
    id = UUIDField(primary_key=True)
    element_id = ForeignKeyField(CachedElement, backref="transcriptions")
    text = TextField()
    confidence = FloatField()
    worker_version_id = UUIDField()

    class Meta:
        database = db
        table_name = "transcriptions"


def init_cache_db(path):
    db.init(
        path,
        pragmas={
            # SQLite ignores foreign keys and check constraints by default!
            "foreign_keys": 1,
            "ignore_check_constraints": 0,
        },
    )
    db.connect()


def create_tables():
    """
    Creates the tables in the cache DB only if they do not already exist.
    """
    db.create_tables([CachedElement, CachedTranscription])