123 lines
4.9 KiB
Python
123 lines
4.9 KiB
Python
from typing import Set
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import text
|
|
from loguru import logger
|
|
|
|
|
|
_checked_subscription_plan_columns: bool = False
|
|
_checked_usage_summaries_columns: bool = False
|
|
|
|
|
|
def ensure_subscription_plan_columns(db: Session) -> None:
|
|
"""Ensure required columns exist on subscription_plans for runtime safety.
|
|
|
|
This is a defensive guard for environments where migrations have not yet
|
|
been applied. If columns are missing (e.g., exa_calls_limit), we add them
|
|
with a safe default so ORM queries do not fail.
|
|
"""
|
|
global _checked_subscription_plan_columns
|
|
if _checked_subscription_plan_columns:
|
|
return
|
|
|
|
try:
|
|
# Discover existing columns using PRAGMA
|
|
result = db.execute(text("PRAGMA table_info(subscription_plans)"))
|
|
cols: Set[str] = {row[1] for row in result}
|
|
|
|
logger.debug(f"Schema check: Found {len(cols)} columns in subscription_plans table")
|
|
|
|
# Columns we may reference in models but might be missing in older DBs
|
|
required_columns = {
|
|
"exa_calls_limit": "INTEGER DEFAULT 0",
|
|
"video_calls_limit": "INTEGER DEFAULT 0",
|
|
"image_edit_calls_limit": "INTEGER DEFAULT 0",
|
|
"audio_calls_limit": "INTEGER DEFAULT 0",
|
|
}
|
|
|
|
for col_name, ddl in required_columns.items():
|
|
if col_name not in cols:
|
|
logger.info(f"Adding missing column {col_name} to subscription_plans table")
|
|
try:
|
|
db.execute(text(f"ALTER TABLE subscription_plans ADD COLUMN {col_name} {ddl}"))
|
|
db.commit()
|
|
logger.info(f"Successfully added column {col_name}")
|
|
except Exception as alter_err:
|
|
logger.error(f"Failed to add column {col_name}: {alter_err}")
|
|
db.rollback()
|
|
# Don't set flag on error - allow retry
|
|
raise
|
|
else:
|
|
logger.debug(f"Column {col_name} already exists")
|
|
|
|
# Only set flag if we successfully completed the check
|
|
_checked_subscription_plan_columns = True
|
|
except Exception as e:
|
|
logger.error(f"Error ensuring subscription_plan columns: {e}", exc_info=True)
|
|
db.rollback()
|
|
# Don't set the flag if there was an error, so we retry next time
|
|
_checked_subscription_plan_columns = False
|
|
raise
|
|
|
|
|
|
def ensure_usage_summaries_columns(db: Session) -> None:
|
|
"""Ensure required columns exist on usage_summaries for runtime safety.
|
|
|
|
This is a defensive guard for environments where migrations have not yet
|
|
been applied. If columns are missing (e.g., exa_calls, exa_cost), we add them
|
|
with a safe default so ORM queries do not fail.
|
|
"""
|
|
global _checked_usage_summaries_columns
|
|
if _checked_usage_summaries_columns:
|
|
return
|
|
|
|
try:
|
|
# Discover existing columns using PRAGMA
|
|
result = db.execute(text("PRAGMA table_info(usage_summaries)"))
|
|
cols: Set[str] = {row[1] for row in result}
|
|
|
|
logger.debug(f"Schema check: Found {len(cols)} columns in usage_summaries table")
|
|
|
|
# Columns we may reference in models but might be missing in older DBs
|
|
required_columns = {
|
|
"exa_calls": "INTEGER DEFAULT 0",
|
|
"exa_cost": "REAL DEFAULT 0.0",
|
|
"video_calls": "INTEGER DEFAULT 0",
|
|
"video_cost": "REAL DEFAULT 0.0",
|
|
"image_edit_calls": "INTEGER DEFAULT 0",
|
|
"image_edit_cost": "REAL DEFAULT 0.0",
|
|
"audio_calls": "INTEGER DEFAULT 0",
|
|
"audio_cost": "REAL DEFAULT 0.0",
|
|
}
|
|
|
|
for col_name, ddl in required_columns.items():
|
|
if col_name not in cols:
|
|
logger.info(f"Adding missing column {col_name} to usage_summaries table")
|
|
try:
|
|
db.execute(text(f"ALTER TABLE usage_summaries ADD COLUMN {col_name} {ddl}"))
|
|
db.commit()
|
|
logger.info(f"Successfully added column {col_name}")
|
|
except Exception as alter_err:
|
|
logger.error(f"Failed to add column {col_name}: {alter_err}")
|
|
db.rollback()
|
|
# Don't set flag on error - allow retry
|
|
raise
|
|
else:
|
|
logger.debug(f"Column {col_name} already exists")
|
|
|
|
# Only set flag if we successfully completed the check
|
|
_checked_usage_summaries_columns = True
|
|
except Exception as e:
|
|
logger.error(f"Error ensuring usage_summaries columns: {e}", exc_info=True)
|
|
db.rollback()
|
|
# Don't set the flag if there was an error, so we retry next time
|
|
_checked_usage_summaries_columns = False
|
|
raise
|
|
|
|
|
|
def ensure_all_schema_columns(db: Session) -> None:
|
|
"""Ensure all required columns exist in subscription-related tables."""
|
|
ensure_subscription_plan_columns(db)
|
|
ensure_usage_summaries_columns(db)
|
|
|
|
|