cosmo/backend/scripts/add_predefined_task_columns.py

94 lines
3.0 KiB
Python

"""
Simple migration to add predefined task columns
"""
import asyncio
import sys
from pathlib import Path
# Add backend to path
sys.path.insert(0, str(Path(__file__).parent.parent))
from sqlalchemy import text
from app.database import engine
async def run_simple_migration():
"""Add the new columns to scheduled_jobs table"""
async with engine.begin() as conn:
print("🔄 Adding new columns to scheduled_jobs table...")
# Add job_type column
try:
await conn.execute(text("""
ALTER TABLE scheduled_jobs
ADD COLUMN job_type jobtype DEFAULT 'custom_code'::jobtype NOT NULL
"""))
print("✅ Added job_type column")
except Exception as e:
print(f"⚠️ job_type column: {e}")
# Add predefined_function column
try:
await conn.execute(text("""
ALTER TABLE scheduled_jobs
ADD COLUMN predefined_function VARCHAR(100)
"""))
print("✅ Added predefined_function column")
except Exception as e:
print(f"⚠️ predefined_function column: {e}")
# Add function_params column
try:
await conn.execute(text("""
ALTER TABLE scheduled_jobs
ADD COLUMN function_params JSONB DEFAULT '{}'::jsonb
"""))
print("✅ Added function_params column")
except Exception as e:
print(f"⚠️ function_params column: {e}")
# Set default for future records to 'predefined'
try:
await conn.execute(text("""
ALTER TABLE scheduled_jobs
ALTER COLUMN job_type SET DEFAULT 'predefined'::jobtype
"""))
print("✅ Set default job_type to 'predefined'")
except Exception as e:
print(f"⚠️ Setting default: {e}")
# Add check constraint
try:
await conn.execute(text("""
ALTER TABLE scheduled_jobs
DROP CONSTRAINT IF EXISTS chk_job_type_fields
"""))
await conn.execute(text("""
ALTER TABLE scheduled_jobs
ADD CONSTRAINT chk_job_type_fields
CHECK (
(job_type = 'predefined' AND predefined_function IS NOT NULL)
OR
(job_type = 'custom_code' AND python_code IS NOT NULL)
)
"""))
print("✅ Added check constraint")
except Exception as e:
print(f"⚠️ Check constraint: {e}")
print("\n📋 Final table structure:")
result = await conn.execute(text("""
SELECT column_name, data_type, is_nullable
FROM information_schema.columns
WHERE table_name = 'scheduled_jobs'
ORDER BY ordinal_position
"""))
rows = result.fetchall()
for row in rows:
print(f" - {row[0]}: {row[1]} (nullable: {row[2]})")
if __name__ == "__main__":
asyncio.run(run_simple_migration())