Coverage for backend / app / demo / setup.py: 97%

1"""Demo schema setup and initialization.""" 

2 

3import datetime as dt 

4 

5from sqlalchemy import text 

6from sqlalchemy.orm import Session 

7 

8from tests.utils.test_data import data_tables 

9from tests.utils.create_data.utils import create_db_entries 

10from app import models 

11from app.database import Base, demo_engine, demo_session_local, engine 

12from app.job_rating.prompts import seed_ai_prompts 

13 

14 

15_DEMO_SETUP_LOCK_ID = 987_654_321 

16 

17 

18def setup_demo_schema() -> None: 

19 """Create the demo schema and all tables, then seed required data. 

20 Called at application startup.""" 

21 

22 with engine.connect() as conn: 

23 # Serialize setup across processes (e.g. multi-port dev servers) 

24 conn.execute(text("SELECT pg_advisory_lock(:lock_id)"), {"lock_id": _DEMO_SETUP_LOCK_ID}) 

25 try: 

26 conn.execute(text("DROP SCHEMA IF EXISTS demo CASCADE")) 

27 conn.execute(text("CREATE SCHEMA demo")) 

28 conn.commit() 

29 

30 # Create all tables in the demo schema 

31 Base.metadata.create_all(bind=demo_engine) 

32 

33 # Seed required data in the demo schema 

34 db = demo_session_local() 

35 try: 

36 seed_ai_prompts(db) 

37 create_db_entries(db, models.Geolocation, data_tables.GEOLOCATION_DATA) 

38 cleanup_stale_demo_users(db) 

39 finally: 

40 db.close() 

41 finally: 

42 conn.execute(text("SELECT pg_advisory_unlock(:lock_id)"), {"lock_id": _DEMO_SETUP_LOCK_ID}) 

43 

44 

45def seed_demo_ai_prompts(db: Session) -> None: 

46 """Seed AI prompts in the demo schema if they don't exist.""" 

47 

48 existing = db.query(models.AiSystemPrompt).first() 

49 if not existing: 

50 seed_ai_prompts(db) 

51 

52 

53def cleanup_stale_demo_users(db: Session) -> None: 

54 """Delete demo users older than 24 hours from the demo schema.""" 

55 

56 cutoff = dt.datetime.now(dt.timezone.utc) - dt.timedelta(hours=24) 

57 stale_users = ( 

58 db.query(models.User).filter(models.User.is_demo.is_(True)).filter(models.User.created_at < cutoff).all() 

59 ) 

60 for user in stale_users: 

61 db.delete(user) 

62 if stale_users: 

63 db.commit()