File size: 765 Bytes
a101d9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import sqlite3
from pathlib import Path


class Database:
    DB_PATH = Path("data/")
    DB_FILE = DB_PATH / "models.db"

    def __init__(self):
        if not self.DB_FILE.exists():
            print("Creating database")
            print("DB_FILE", self.DB_FILE)
            db = sqlite3.connect(self.DB_FILE)
            with open(Path("schema.sql"), "r") as f:
                db.executescript(f.read())
            db.commit()
            db.close()

    def get_db(self):
        db = sqlite3.connect(self.DB_FILE, check_same_thread=False)
        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()