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())
|