120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
"""
|
||
Fix enum type and add 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 fix_enum_and_migrate():
|
||
"""Fix enum type and add columns"""
|
||
async with engine.begin() as conn:
|
||
# First check enum values
|
||
result = await conn.execute(text("""
|
||
SELECT enumlabel
|
||
FROM pg_enum
|
||
WHERE enumtypid = 'jobtype'::regtype
|
||
ORDER BY enumsortorder
|
||
"""))
|
||
enum_values = [row[0] for row in result.fetchall()]
|
||
print(f"Current enum values: {enum_values}")
|
||
|
||
# Add missing enum values if needed
|
||
if 'predefined' not in enum_values:
|
||
await conn.execute(text("ALTER TYPE jobtype ADD VALUE 'predefined'"))
|
||
print("✅ Added 'predefined' to enum")
|
||
|
||
if 'custom_code' not in enum_values:
|
||
await conn.execute(text("ALTER TYPE jobtype ADD VALUE 'custom_code'"))
|
||
print("✅ Added 'custom_code' to enum")
|
||
|
||
# Now add columns in separate transaction
|
||
async with engine.begin() as conn:
|
||
print("\n🔄 Adding 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:
|
||
if "already exists" in str(e):
|
||
print("ℹ️ job_type column already exists")
|
||
else:
|
||
raise
|
||
|
||
# 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:
|
||
if "already exists" in str(e):
|
||
print("ℹ️ predefined_function column already exists")
|
||
else:
|
||
raise
|
||
|
||
# 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:
|
||
if "already exists" in str(e):
|
||
print("ℹ️ function_params column already exists")
|
||
else:
|
||
raise
|
||
|
||
# Set defaults and constraints in separate transaction
|
||
async with engine.begin() as conn:
|
||
# Set default for future records
|
||
await conn.execute(text("""
|
||
ALTER TABLE scheduled_jobs
|
||
ALTER COLUMN job_type SET DEFAULT 'predefined'::jobtype
|
||
"""))
|
||
print("✅ Set default job_type to 'predefined'")
|
||
|
||
# Drop and recreate check constraint
|
||
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")
|
||
|
||
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(fix_enum_and_migrate())
|