149 lines
5.5 KiB
Python
149 lines
5.5 KiB
Python
"""
|
|
Migration script for ScriptsManager Proxy System
|
|
Adds support for project-based script execution and internal proxy
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, '/app')
|
|
|
|
from app.config.database import db
|
|
from app.models import ScriptProxyExecution, ExecutionLog
|
|
from sqlalchemy import text
|
|
|
|
def run_migration():
|
|
"""Run the migration to add proxy support"""
|
|
|
|
print("=== ScriptsManager Proxy Migration ===")
|
|
|
|
try:
|
|
# Add new columns to execution_logs table
|
|
print("Adding project_id, proxy_port, proxy_url, execution_mode to execution_logs...")
|
|
|
|
# Check if columns already exist
|
|
result = db.session.execute(text("""
|
|
SELECT column_name
|
|
FROM information_schema.columns
|
|
WHERE table_name = 'execution_logs'
|
|
AND column_name IN ('project_id', 'proxy_port', 'proxy_url', 'execution_mode')
|
|
""")).fetchall()
|
|
|
|
existing_columns = [row[0] for row in result]
|
|
|
|
if 'project_id' not in existing_columns:
|
|
db.session.execute(text("""
|
|
ALTER TABLE execution_logs
|
|
ADD COLUMN project_id INTEGER REFERENCES user_projects(id)
|
|
"""))
|
|
print("✓ Added project_id column")
|
|
|
|
if 'proxy_port' not in existing_columns:
|
|
db.session.execute(text("""
|
|
ALTER TABLE execution_logs
|
|
ADD COLUMN proxy_port INTEGER
|
|
"""))
|
|
print("✓ Added proxy_port column")
|
|
|
|
if 'proxy_url' not in existing_columns:
|
|
db.session.execute(text("""
|
|
ALTER TABLE execution_logs
|
|
ADD COLUMN proxy_url VARCHAR(255)
|
|
"""))
|
|
print("✓ Added proxy_url column")
|
|
|
|
if 'execution_mode' not in existing_columns:
|
|
db.session.execute(text("""
|
|
ALTER TABLE execution_logs
|
|
ADD COLUMN execution_mode VARCHAR(20) DEFAULT 'direct'
|
|
"""))
|
|
print("✓ Added execution_mode column")
|
|
|
|
# Create script_proxy_executions table
|
|
print("Creating script_proxy_executions table...")
|
|
|
|
# Check if table exists
|
|
result = db.session.execute(text("""
|
|
SELECT table_name
|
|
FROM information_schema.tables
|
|
WHERE table_name = 'script_proxy_executions'
|
|
""")).fetchall()
|
|
|
|
if not result:
|
|
db.session.execute(text("""
|
|
CREATE TABLE script_proxy_executions (
|
|
id SERIAL PRIMARY KEY,
|
|
script_id INTEGER NOT NULL REFERENCES scripts(id),
|
|
user_id INTEGER NOT NULL REFERENCES users(id),
|
|
project_id INTEGER NOT NULL REFERENCES user_projects(id),
|
|
internal_port INTEGER NOT NULL UNIQUE,
|
|
proxy_path VARCHAR(255) NOT NULL,
|
|
workspace_path VARCHAR(500) NOT NULL,
|
|
process_id INTEGER,
|
|
status VARCHAR(20) DEFAULT 'starting',
|
|
conda_environment VARCHAR(100),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
started_at TIMESTAMP,
|
|
last_activity TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
stopped_at TIMESTAMP,
|
|
script_name VARCHAR(100),
|
|
parameters TEXT,
|
|
environment_vars TEXT,
|
|
CONSTRAINT unique_project_script_user UNIQUE (project_id, script_id, user_id)
|
|
)
|
|
"""))
|
|
print("✓ Created script_proxy_executions table")
|
|
else:
|
|
print("✓ script_proxy_executions table already exists")
|
|
|
|
# Create indexes for better performance
|
|
print("Creating indexes...")
|
|
|
|
try:
|
|
db.session.execute(text("""
|
|
CREATE INDEX IF NOT EXISTS idx_proxy_executions_status
|
|
ON script_proxy_executions(status)
|
|
"""))
|
|
|
|
db.session.execute(text("""
|
|
CREATE INDEX IF NOT EXISTS idx_proxy_executions_user_project
|
|
ON script_proxy_executions(user_id, project_id)
|
|
"""))
|
|
|
|
db.session.execute(text("""
|
|
CREATE INDEX IF NOT EXISTS idx_proxy_executions_last_activity
|
|
ON script_proxy_executions(last_activity)
|
|
"""))
|
|
|
|
print("✓ Created performance indexes")
|
|
|
|
except Exception as e:
|
|
print(f"⚠ Warning creating indexes: {e}")
|
|
|
|
# Commit all changes
|
|
db.session.commit()
|
|
print("✓ Migration completed successfully!")
|
|
|
|
# Show final status
|
|
print("\n=== Migration Summary ===")
|
|
|
|
# Count existing records
|
|
execution_logs_count = db.session.execute(text("SELECT COUNT(*) FROM execution_logs")).scalar()
|
|
proxy_executions_count = db.session.execute(text("SELECT COUNT(*) FROM script_proxy_executions")).scalar()
|
|
|
|
print(f"Execution logs: {execution_logs_count}")
|
|
print(f"Proxy executions: {proxy_executions_count}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"✗ Migration failed: {e}")
|
|
db.session.rollback()
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
from app.app import create_app
|
|
|
|
app = create_app()
|
|
with app.app_context():
|
|
success = run_migration()
|
|
sys.exit(0 if success else 1) |