Coverage for backend/tests/utils/seed_database.py: 82%

1""" 

2Database seeding script for development. 

3This script will drop all data and repopulate with hard-coded sample data. 

4""" 

5 

6import os 

7import sys 

8 

9from sqlalchemy import text, inspect 

10 

11from app.database import engine, session_local, Base 

12from tests.utils.create_data import ( 

13 create_users, 

14 create_settings, 

15 create_companies, 

16 create_locations, 

17 create_aggregators, 

18 create_keywords, 

19 create_people, 

20 create_jobs, 

21 create_files, 

22 create_interviews, 

23 create_job_alert_emails, 

24 create_scraped_jobs, 

25 create_service_logs, 

26 create_job_application_updates, 

27) 

28 

29sys.path.append(os.path.dirname(os.path.abspath(__file__))) 

30 

31 

32def reset_database(db_engine) -> None: 

33 """Drop ALL tables in the database (including orphaned ones) and recreate from models""" 

34 print("Dropping all tables in the database...") 

35 

36 with db_engine.connect() as conn: 

37 # Get the database inspector to find all existing tables 

38 inspector = inspect(db_engine) 

39 table_names = inspector.get_table_names() 

40 

41 # For PostgreSQL: Drop tables with CASCADE to handle foreign key constraints 

42 for table_name in table_names: 

43 conn.execute(text(f'DROP TABLE IF EXISTS "{table_name}" CASCADE')) 

44 

45 conn.commit() 

46 

47 print("Creating all tables from models...") 

48 Base.metadata.create_all(bind=db_engine) 

49 

50 

51def seed_database() -> None: 

52 """Main function to seed the database""" 

53 print("Starting database seeding...") 

54 

55 # Reset the database 

56 reset_database(engine) 

57 

58 # Create a database session 

59 db = session_local() 

60 # 

61 # try: 

62 # # Create data in order of dependencies 

63 # users = create_users(db) 

64 # settings = create_settings(db) 

65 # keywords = create_keywords(db, users) 

66 # aggregators = create_aggregators(db, users) 

67 # locations = create_locations(db, users) 

68 # companies = create_companies(db, users) 

69 # people = create_people(db, users, companies) 

70 # files = create_files(db, users) 

71 # jobs = create_jobs(db, keywords, people, users, companies, locations, aggregators, files) 

72 # interviews = create_interviews(db, people, users, locations, jobs) 

73 # job_application_updates = create_job_application_updates(db, users, jobs) 

74 # service_logs = create_service_logs(db) 

75 # alert_emails = create_job_alert_emails(db, users, service_logs) 

76 # scraped_jobs = create_scraped_jobs(db, alert_emails, users) 

77 # 

78 # print("\n" + "=" * 50) 

79 # print("DATABASE SEEDING COMPLETED SUCCESSFULLY!") 

80 # print("=" * 50) 

81 # print(f"Users: {len(users)}") 

82 # print(f"Settings: {len(settings)}") 

83 # print(f"Companies: {len(companies)}") 

84 # print(f"Locations: {len(locations)}") 

85 # print(f"Aggregators: {len(aggregators)}") 

86 # print(f"Keywords: {len(keywords)}") 

87 # print(f"People: {len(people)}") 

88 # print(f"Jobs: {len(jobs)}") 

89 # print(f"Files: {len(files)}") 

90 # print(f"Interviews: {len(interviews)}") 

91 # print(f"Service Logs: {len(service_logs)}") 

92 # print(f"Job Alert Emails: {len(alert_emails)}") 

93 # print(f"Scraped Jobs: {len(scraped_jobs)}") 

94 # print(f"Job Application Updates: {len(job_application_updates)}") 

95 # print("=" * 50) 

96 # 

97 # except Exception as e: 

98 # print(f"Error during seeding: {e}") 

99 # db.rollback() 

100 # raise 

101 # finally: 

102 # db.close() 

103 

104 

105if __name__ == "__main__": 

106 seed_database()