File size: 916 Bytes
9a5a85e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import sqlite3
from pathlib import Path
import math


def power(x, y):
    return math.pow(x, y)


class Database:
    def __init__(self, db_path=None):
        if db_path is None:
            raise ValueError("db_path must be provided")
        self.db_path = db_path
        self.db_file = self.db_path / "models.db"
        with self.get_db() as db:
            try:
                db.execute("ALTER TABLE models ADD COLUMN flags TEXT")
            except sqlite3.OperationalError:
                print("Column Flags already exists")
                pass

    def get_db(self):
        db = sqlite3.connect(self.db_file, check_same_thread=False)
        db.create_function("MYPOWER", 2, power)
        db.row_factory = sqlite3.Row
        return db

    def __enter__(self):
        self.db = self.get_db()
        return self.db

    def __exit__(self, exc_type, exc_value, traceback):
        self.db.close()