68 lines
1.6 KiB
Python
68 lines
1.6 KiB
Python
"""Database related functions."""
|
|
|
|
import logging
|
|
import secrets
|
|
from pathlib import Path
|
|
from sqlalchemy import Engine
|
|
from sqlmodel import Session, create_engine, text
|
|
import bcrypt
|
|
|
|
from sqlalchemy.engine import URL
|
|
|
|
from .models import APIClient, init_db
|
|
|
|
from .settings import get_settings
|
|
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
|
|
def get_engine(filename: Path, echo: bool = False) -> Engine:
|
|
"""Initialize the engine."""
|
|
url = URL.create(drivername="sqlite", database=str(filename.absolute()))
|
|
engine = create_engine(url, echo=echo)
|
|
with engine.connect() as connection:
|
|
connection.execute(text("PRAGMA foreign_keys=ON")) # for SQLite only
|
|
|
|
return engine
|
|
|
|
|
|
def create_db_and_tables(filename: Path, echo: bool = False) -> bool:
|
|
"""Create database and tables.
|
|
|
|
Returns True if the database was created.
|
|
"""
|
|
created = False
|
|
if not filename.exists():
|
|
created = True
|
|
engine = get_engine(filename, echo)
|
|
|
|
init_db(engine)
|
|
return created
|
|
|
|
|
|
def create_api_token(session: Session, read_write: bool) -> str:
|
|
"""Create API token."""
|
|
token = secrets.token_urlsafe(32)
|
|
pwbytes = token.encode("utf-8")
|
|
salt = bcrypt.gensalt()
|
|
hashed_bytes = bcrypt.hashpw(password=pwbytes, salt=salt)
|
|
hashed = hashed_bytes.decode()
|
|
|
|
api_token = APIClient(token=hashed, read_write=read_write)
|
|
session.add(api_token)
|
|
session.commit()
|
|
|
|
return token
|
|
|
|
|
|
def generate_api_token() -> str:
|
|
"""Generate API token."""
|
|
settings = get_settings()
|
|
engine = get_engine(settings.db_file)
|
|
init_db(engine)
|
|
with Session(engine) as session:
|
|
token = create_api_token(session, True)
|
|
|
|
return token
|