20 lines
610 B
Python
20 lines
610 B
Python
from typing import List
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import Session
|
|
|
|
from models.orm import Prompt
|
|
|
|
|
|
class DB:
|
|
def __init__(self, db: str = "sqlite:///example.db"):
|
|
self.engine = create_engine(db, connect_args={"check_same_thread": False})
|
|
|
|
def get_prompts(self) -> List[Prompt]:
|
|
with Session(self.engine) as session:
|
|
return session.query(Prompt).all()
|
|
|
|
def get_prompt_by_id(self, prompt_id: int) -> Prompt | None:
|
|
with Session(self.engine) as session:
|
|
return session.query(Prompt).filter(Prompt.id == prompt_id).first()
|