database.py•826 B
from sqlmodel import SQLModel, create_engine, Session
from models import User, TrackedQuery, Run, RunItem, Aggregate
import os
# Database setup
DATABASE_URL = "sqlite:///./visibility.sqlite"
engine = create_engine(DATABASE_URL, echo=False)
def create_db_and_tables():
"""Create database and tables"""
SQLModel.metadata.create_all(engine)
def get_session():
"""Get database session"""
with Session(engine) as session:
yield session
def ensure_user_exists(phone: str) -> User:
"""Ensure user exists, create if not"""
with Session(engine) as session:
user = session.query(User).filter(User.phone == phone).first()
if not user:
user = User(phone=phone)
session.add(user)
session.commit()
session.refresh(user)
return user