138 lines
4.0 KiB
Python
138 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Apply v9c Subscription System Migration
|
|
|
|
This script checks if v9c migration is needed and applies it.
|
|
Run automatically on container startup via main.py startup event.
|
|
"""
|
|
import os
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
|
|
def get_db_connection():
|
|
"""Get PostgreSQL connection."""
|
|
return psycopg2.connect(
|
|
host=os.getenv("DB_HOST", "postgres"),
|
|
port=int(os.getenv("DB_PORT", 5432)),
|
|
database=os.getenv("DB_NAME", "mitai_prod"),
|
|
user=os.getenv("DB_USER", "mitai_prod"),
|
|
password=os.getenv("DB_PASSWORD", ""),
|
|
cursor_factory=RealDictCursor
|
|
)
|
|
|
|
|
|
def migration_needed(conn):
|
|
"""Check if v9c migration is needed."""
|
|
cur = conn.cursor()
|
|
|
|
# Check if tiers table exists
|
|
cur.execute("""
|
|
SELECT EXISTS (
|
|
SELECT FROM information_schema.tables
|
|
WHERE table_name = 'tiers'
|
|
)
|
|
""")
|
|
tiers_exists = cur.fetchone()['exists']
|
|
|
|
# Check if features table exists
|
|
cur.execute("""
|
|
SELECT EXISTS (
|
|
SELECT FROM information_schema.tables
|
|
WHERE table_name = 'features'
|
|
)
|
|
""")
|
|
features_exists = cur.fetchone()['exists']
|
|
|
|
cur.close()
|
|
|
|
# Migration needed if either table is missing
|
|
return not (tiers_exists and features_exists)
|
|
|
|
|
|
def apply_migration():
|
|
"""Apply v9c migration if needed."""
|
|
print("[v9c Migration] Checking if migration is needed...")
|
|
|
|
try:
|
|
conn = get_db_connection()
|
|
|
|
if not migration_needed(conn):
|
|
print("[v9c Migration] Already applied, skipping.")
|
|
conn.close()
|
|
return
|
|
|
|
print("[v9c Migration] Applying subscription system migration...")
|
|
|
|
# Read migration SQL
|
|
migration_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"migrations",
|
|
"v9c_subscription_system.sql"
|
|
)
|
|
|
|
with open(migration_path, 'r', encoding='utf-8') as f:
|
|
migration_sql = f.read()
|
|
|
|
# Execute migration
|
|
cur = conn.cursor()
|
|
cur.execute(migration_sql)
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
|
|
print("[v9c Migration] ✅ Migration completed successfully!")
|
|
|
|
# Apply fix migration if exists
|
|
fix_migration_path = os.path.join(
|
|
os.path.dirname(__file__),
|
|
"migrations",
|
|
"v9c_fix_features.sql"
|
|
)
|
|
|
|
if os.path.exists(fix_migration_path):
|
|
print("[v9c Migration] Applying feature fixes...")
|
|
with open(fix_migration_path, 'r', encoding='utf-8') as f:
|
|
fix_sql = f.read()
|
|
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
cur.execute(fix_sql)
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print("[v9c Migration] ✅ Feature fixes applied!")
|
|
|
|
# Verify tables created
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
cur.execute("""
|
|
SELECT table_name FROM information_schema.tables
|
|
WHERE table_schema = 'public'
|
|
AND table_name IN ('tiers', 'features', 'tier_limits', 'access_grants', 'coupons')
|
|
ORDER BY table_name
|
|
""")
|
|
tables = [r['table_name'] for r in cur.fetchall()]
|
|
print(f"[v9c Migration] Created tables: {', '.join(tables)}")
|
|
|
|
# Verify initial data
|
|
cur.execute("SELECT COUNT(*) as count FROM tiers")
|
|
tier_count = cur.fetchone()['count']
|
|
cur.execute("SELECT COUNT(*) as count FROM features")
|
|
feature_count = cur.fetchone()['count']
|
|
cur.execute("SELECT COUNT(*) as count FROM tier_limits")
|
|
limit_count = cur.fetchone()['count']
|
|
|
|
print(f"[v9c Migration] Initial data: {tier_count} tiers, {feature_count} features, {limit_count} tier limits")
|
|
|
|
cur.close()
|
|
conn.close()
|
|
|
|
except Exception as e:
|
|
print(f"[v9c Migration] ❌ Error: {e}")
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
apply_migration()
|