""" Migration script to add job reservations for the Fleet Calendar feature. This creates two tables: - job_reservations: Track future unit assignments for jobs/projects - job_reservation_units: Link specific units to reservations Run this script once to migrate an existing database. """ import sqlite3 import os # Database path DB_PATH = "./data/seismo_fleet.db" def migrate_database(): """Create the job_reservations and job_reservation_units tables""" if not os.path.exists(DB_PATH): print(f"Database not found at {DB_PATH}") print("The database will be created automatically when you run the application.") return print(f"Migrating database: {DB_PATH}") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Check if job_reservations table already exists cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='job_reservations'") if cursor.fetchone(): print("Migration already applied - job_reservations table exists") conn.close() return print("Creating job_reservations table...") try: # Create job_reservations table cursor.execute(""" CREATE TABLE job_reservations ( id TEXT PRIMARY KEY, name TEXT NOT NULL, project_id TEXT, start_date DATE NOT NULL, end_date DATE NOT NULL, assignment_type TEXT NOT NULL DEFAULT 'quantity', device_type TEXT DEFAULT 'seismograph', quantity_needed INTEGER, notes TEXT, color TEXT DEFAULT '#3B82F6', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) print(" Created job_reservations table") # Create indexes for job_reservations cursor.execute("CREATE INDEX idx_job_reservations_project_id ON job_reservations(project_id)") print(" Created index on project_id") cursor.execute("CREATE INDEX idx_job_reservations_dates ON job_reservations(start_date, end_date)") print(" Created index on dates") # Create job_reservation_units table print("Creating job_reservation_units table...") cursor.execute(""" CREATE TABLE job_reservation_units ( id TEXT PRIMARY KEY, reservation_id TEXT NOT NULL, unit_id TEXT NOT NULL, assignment_source TEXT DEFAULT 'specific', assigned_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (reservation_id) REFERENCES job_reservations(id), FOREIGN KEY (unit_id) REFERENCES roster(id) ) """) print(" Created job_reservation_units table") # Create indexes for job_reservation_units cursor.execute("CREATE INDEX idx_job_reservation_units_reservation_id ON job_reservation_units(reservation_id)") print(" Created index on reservation_id") cursor.execute("CREATE INDEX idx_job_reservation_units_unit_id ON job_reservation_units(unit_id)") print(" Created index on unit_id") conn.commit() print("\nMigration completed successfully!") print("You can now use the Fleet Calendar to manage unit reservations.") except sqlite3.Error as e: print(f"\nError during migration: {e}") conn.rollback() raise finally: conn.close() if __name__ == "__main__": migrate_database()