90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
"""
|
|
Migration: Add TBD date support to job reservations
|
|
|
|
Adds columns:
|
|
- job_reservations.estimated_end_date: For planning when end is TBD
|
|
- job_reservations.end_date_tbd: Boolean flag for TBD end dates
|
|
- job_reservation_units.unit_start_date: Unit-specific start (for swaps)
|
|
- job_reservation_units.unit_end_date: Unit-specific end (for swaps)
|
|
- job_reservation_units.unit_end_tbd: Unit-specific TBD flag
|
|
- job_reservation_units.notes: Notes for the assignment
|
|
|
|
Also makes job_reservations.end_date nullable.
|
|
"""
|
|
|
|
import sqlite3
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def migrate(db_path: str):
|
|
"""Run the migration."""
|
|
print(f"Migrating database: {db_path}")
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
try:
|
|
# Check if job_reservations table exists
|
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='job_reservations'")
|
|
if not cursor.fetchone():
|
|
print("job_reservations table does not exist. Skipping migration.")
|
|
return
|
|
|
|
# Get existing columns in job_reservations
|
|
cursor.execute("PRAGMA table_info(job_reservations)")
|
|
existing_cols = {row[1] for row in cursor.fetchall()}
|
|
|
|
# Add new columns to job_reservations if they don't exist
|
|
if 'estimated_end_date' not in existing_cols:
|
|
print("Adding estimated_end_date column to job_reservations...")
|
|
cursor.execute("ALTER TABLE job_reservations ADD COLUMN estimated_end_date DATE")
|
|
|
|
if 'end_date_tbd' not in existing_cols:
|
|
print("Adding end_date_tbd column to job_reservations...")
|
|
cursor.execute("ALTER TABLE job_reservations ADD COLUMN end_date_tbd BOOLEAN DEFAULT 0")
|
|
|
|
# Get existing columns in job_reservation_units
|
|
cursor.execute("PRAGMA table_info(job_reservation_units)")
|
|
unit_cols = {row[1] for row in cursor.fetchall()}
|
|
|
|
# Add new columns to job_reservation_units if they don't exist
|
|
if 'unit_start_date' not in unit_cols:
|
|
print("Adding unit_start_date column to job_reservation_units...")
|
|
cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_start_date DATE")
|
|
|
|
if 'unit_end_date' not in unit_cols:
|
|
print("Adding unit_end_date column to job_reservation_units...")
|
|
cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_end_date DATE")
|
|
|
|
if 'unit_end_tbd' not in unit_cols:
|
|
print("Adding unit_end_tbd column to job_reservation_units...")
|
|
cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN unit_end_tbd BOOLEAN DEFAULT 0")
|
|
|
|
if 'notes' not in unit_cols:
|
|
print("Adding notes column to job_reservation_units...")
|
|
cursor.execute("ALTER TABLE job_reservation_units ADD COLUMN notes TEXT")
|
|
|
|
conn.commit()
|
|
print("Migration completed successfully!")
|
|
|
|
except Exception as e:
|
|
print(f"Migration failed: {e}")
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Default to dev database
|
|
db_path = "./data-dev/seismo_fleet.db"
|
|
|
|
if len(sys.argv) > 1:
|
|
db_path = sys.argv[1]
|
|
|
|
if not Path(db_path).exists():
|
|
print(f"Database not found: {db_path}")
|
|
sys.exit(1)
|
|
|
|
migrate(db_path)
|