#!/usr/bin/env python3
"""Update models.py to remove migrated functions and add re-exports.
Reads original models.py, strips out functions that moved to other modules,
adds import/re-export statements at top, preserves everything else intact.
"""
import re
import sys
MODEL_PATH = "app/models.py"
BACKUP_PATH = "app/models.py.bak"
MIGRATED_FUNCS = {
'migrate_add_field_config', 'migrate_add_submission_data',
'migrate_add_user_id_to_sites', 'migrate_add_usage',
'migrate_add_abuse_protection', 'migrate_add_rate_limit',
'migrate_add_analytics', 'migrate_add_email_verification',
'migrate_add_indexes', 'migrate_add_api_keys',
'migrate_encrypt_existing_data', 'migrate_rescope_to_user_keys',
'migrate_rotate_encryption_keys', 'migrate_normalize_field_configs',
'migrate_add_form_templates', 'migrate_add_teams',
'migrate_add_form_analytics', 'migrate_add_ai_builder',
'migrate_add_webhook_destinations', 'migrate_add_form_versions',
'migrate_add_form_sessions', 'migrate_add_form_variants',
'migrate_add_branding', 'migrate_add_webhook_retry_locking',
'migrate_add_document_templates', 'migrate_encrypt_user_pii',
'migrate_add_invite_expiration', 'migrate_add_user_id_to_doc_templates',
'migrate_add_user_id_to_invoice_schedules', 'migrate_add_email_phase3',
'migrate_add_email_phase4', 'migrate_add_email_tracking',
'migrate_add_email_campaigns', 'migrate_add_form_actions',
'migrate_add_user_referral_cols',
'resolve_geo', 'parse_user_agent', '_is_encrypted_value',
'_table_exists', '_current_month', '_decrypt_invoice_schedule',
'_decrypt_webhook_dest', '_decrypt_webhook_log', '_decrypt_team_member',
'_decrypt_user_row', '_encrypt_submission_field',
'_decrypt_submission_fields', '_get_site_owner_password_hash',
'_check_spam_content',
'site_session_stats', 'site_field_dropoff', 'site_hourly_trends',
'ab_test_stats', 'get_site_analytics', 'get_site_analytics_summary',
'site_geo_stats', 'site_device_stats', 'site_impression_stats',
'get_user_analytics', 'get_team_analytics', 'get_admin_stats',
'get_spam_submissions', 'dismiss_spam',
'create_document_template', 'get_document_template',
'list_document_templates', 'get_auto_generate_template',
'update_document_template', 'delete_document_template',
'increment_document_download', 'store_document', 'get_document',
'list_documents', 'update_document_status', 'update_document_due_date',
'create_invoice_schedule', 'get_invoice_schedule',
'list_invoice_schedules', 'update_invoice_schedule',
'delete_invoice_schedule', 'get_due_schedules',
'advance_schedule_next_run', 'generate_qr_code',
'generate_email_html', 'update_site_branding',
'generate_invoice_number', 'update_next_invoice_number',
'list_user_documents', 'count_user_documents',
'list_documents_by_email', 'update_document_stripe_payment',
'create_campaign', 'get_campaign', 'get_campaign_by_id',
'list_campaigns', 'add_campaign_recipients',
'update_campaign_status', 'delete_campaign',
'get_campaign_recipients', 'get_pending_recipients',
'update_recipient_status', 'get_recipient_by_message_id',
'get_email_settings', 'get_recipient_by_token',
'record_open', 'record_click', 'record_bounce',
'get_campaign_recipients_by_email', 'get_campaign_analytics',
'upsert_email_settings', 'create_reminder', 'get_reminder',
'get_reminders', 'cancel_reminder',
'schedule_reminders_for_campaign', 'get_pending_reminders',
'get_reminder_eligible_recipients', 'mark_reminder_sent',
'check_send_quota', 'record_send', 'cleanup_old_send_log',
'create_ab_test', 'get_ab_variant', 'get_ab_variants',
'get_ab_test_info', 'declare_ab_winner',
'assign_recipients_to_variants', 'send_ab_winner_to_remaining',
'check_pending_ab_tests', 'get_custom_templates',
'get_custom_template', 'create_custom_template',
'update_custom_template', 'delete_custom_template',
}
REEXPORT_IMPORTS = '''\
# ── Re-exports from split modules (models.py → domain modules) ──────────────
from app.migrations import (
migrate_add_abuse_protection, migrate_add_ai_builder, migrate_add_analytics,
migrate_add_api_keys, migrate_add_branding, migrate_add_email_campaigns,
migrate_add_email_phase3, migrate_add_email_phase4, migrate_add_email_tracking,
migrate_add_email_verification, migrate_add_field_config,
migrate_add_form_actions, migrate_add_form_analytics, migrate_add_form_sessions,
migrate_add_form_templates, migrate_add_form_versions, migrate_add_form_variants,
migrate_add_indexes, migrate_add_invoice_schedules, migrate_add_invite_expiration,
migrate_add_rate_limit, migrate_add_submission_data, migrate_add_teams,
migrate_add_usage, migrate_add_user_id_to_doc_templates,
migrate_add_user_id_to_invoice_schedules, migrate_add_user_id_to_sites,
migrate_add_user_referral_cols, migrate_add_webhook_destinations,
migrate_add_webhook_retry_locking, migrate_add_document_templates,
migrate_encrypt_existing_data, migrate_encrypt_user_pii,
migrate_normalize_field_configs, migrate_rescope_to_user_keys,
migrate_rotate_encryption_keys,
)
from app.helpers import (
_check_spam_content, _current_month, _decrypt_invoice_schedule,
_decrypt_submission_fields, _decrypt_team_member, _decrypt_user_row,
_decrypt_webhook_dest, _decrypt_webhook_log, _encrypt_submission_field,
_get_site_owner_password_hash, _is_encrypted_value, _table_exists,
parse_user_agent, resolve_geo,
)
from app.analytics import (
ab_test_stats, dismiss_spam, get_admin_stats, get_site_analytics,
get_site_analytics_summary, get_spam_submissions, get_team_analytics,
get_user_analytics, site_device_stats, site_field_dropoff,
site_geo_stats, site_hourly_trends, site_impression_stats,
site_session_stats,
)
from app.model_documents import (
advance_schedule_next_run, create_document_template,
create_invoice_schedule, delete_document_template, delete_invoice_schedule,
generate_email_html, generate_invoice_number, generate_qr_code,
get_auto_generate_template, get_document, get_document_template,
get_due_schedules, get_invoice_schedule, increment_document_download,
list_documents, list_document_templates, list_invoice_schedules,
list_user_documents, count_user_documents, list_documents_by_email,
store_document, update_document_due_date, update_document_status,
update_document_stripe_payment, update_document_template,
update_next_invoice_number, update_site_branding,
)
from app.model_email_campaigns import (
add_campaign_recipients, assign_recipients_to_variants,
cancel_reminder, check_pending_ab_tests, check_send_quota,
cleanup_old_send_log, create_ab_test, create_campaign, create_custom_template,
create_reminder, declare_ab_winner, delete_campaign, delete_custom_template,
get_ab_test_info, get_ab_variant, get_ab_variants, get_campaign,
get_campaign_analytics, get_campaign_by_id, get_campaign_recipients,
get_campaign_recipients_by_email, get_custom_template, get_email_settings,
get_pending_recipients, get_pending_reminders, get_recipient_by_message_id,
get_recipient_by_token, get_reminder, get_reminder_eligible_recipients,
get_reminders, list_campaigns, mark_reminder_sent, record_bounce,
record_click, record_open, record_send, schedule_reminders_for_campaign,
send_ab_winner_to_remaining, update_campaign_status, update_custom_template,
update_recipient_status, upsert_email_settings,
)
'''
def main():
with open(MODEL_PATH, 'r') as f:
lines = f.readlines()
# Backup original
with open(BACKUP_PATH, 'w') as f:
f.writelines(lines)
print(f"Backed up {MODEL_PATH} → {BACKUP_PATH}")
# Find function boundaries
func_pattern = re.compile(r'^(async\s+)?def\s+(\w+)\s*\(')
# Determine which line ranges to remove
keep_lines = []
i = 0
removed_count = 0
while i < len(lines):
m = func_pattern.match(lines[i])
if m and m.group(2) in MIGRATED_FUNCS:
# Skip this function entirely (until next function or end)
removed_count += 1
i += 1
while i < len(lines):
if lines[i].startswith('def ') or lines[i].startswith('class '):
break
i += 1
continue
keep_lines.append(lines[i])
i += 1
# Find where to insert re-exports: after the existing imports, before first function
# The first function is `def get_user_tier` at line 93 (0-indexed: 92)
insert_point = None
for idx, line in enumerate(keep_lines):
if line.startswith('def ') or line.startswith('class '):
insert_point = idx
break
if insert_point is None:
print("ERROR: Could not find insertion point in models.py")
return 1
# Insert re-exports before first function definition
# First, add a blank line separator before the imports
while insert_point > 0 and keep_lines[insert_point - 1].strip() == '':
insert_point -= 1
keep_lines.insert(insert_point, '\n')
keep_lines.insert(insert_point + 1, REEXPORT_IMPORTS + '\n')
with open(MODEL_PATH, 'w') as f:
f.writelines(keep_lines)
original_lines = sum(1 for _ in open(BACKUP_PATH))
new_lines = len(keep_lines)
print(f"\nResults:")
print(f" Original models.py: {original_lines} lines")
print(f" New models.py: {new_lines} lines")
print(f" Functions removed: {removed_count}")
print(f" Lines saved: {original_lines - new_lines}")
print(f"\nBackup: {BACKUP_PATH}")
print("Next: run tests to verify nothing broke")
return 0
if __name__ == '__main__':
sys.exit(main() or 0)