test_500_tables.py

#!/usr/bin/env python3
"""
Test Schema Ranker with 500 tables - realistic enterprise schema
"""
import sys
sys.path.insert(0, 'src')

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
import time

# Load model
print("Loading Schema Ranker model...")
model = SentenceTransformer('models/schema_ranker/schema_ranker_model')

# Generate 500 realistic enterprise tables
TABLES = [
    # Core business (relevant for many queries)
    "customers (id, name, email, phone, address, created_at)",
    "orders (id, customer_id, order_date, total, status)",
    "products (id, name, sku, price, category_id, stock)",
    "payments (id, order_id, amount, method, status, date)",
    "invoices (id, order_id, amount, due_date, paid_date)",
    
    # E-commerce
    "shopping_carts (id, customer_id, created_at)",
    "cart_items (id, cart_id, product_id, quantity)",
    "wishlists (id, customer_id, product_id)",
    "reviews (id, product_id, customer_id, rating, text)",
    "categories (id, name, parent_id, description)",
    "brands (id, name, logo_url, country)",
    "suppliers (id, name, contact, address)",
    "inventory (id, product_id, warehouse_id, quantity)",
    "warehouses (id, name, address, capacity)",
    "shipments (id, order_id, carrier, tracking, status)",
    "returns (id, order_id, reason, status, refund_amount)",
    "coupons (id, code, discount, valid_from, valid_to)",
    "promotions (id, name, start_date, end_date, discount_pct)",
    
    # Users & Auth
    "users (id, email, password_hash, role, active)",
    "user_profiles (id, user_id, avatar, bio, settings)",
    "user_sessions (id, user_id, token, expires_at)",
    "user_roles (id, name, permissions)",
    "user_addresses (id, user_id, type, street, city, zip)",
    "password_resets (id, user_id, token, expires_at)",
    "email_verifications (id, user_id, token, verified_at)",
    "oauth_tokens (id, user_id, provider, access_token)",
    "login_history (id, user_id, ip, device, timestamp)",
    "failed_logins (id, email, ip, attempts, locked_until)",
    
    # HR & Employees
    "employees (id, name, email, department_id, salary)",
    "departments (id, name, manager_id, budget)",
    "positions (id, title, department_id, salary_range)",
    "employee_benefits (id, employee_id, type, value)",
    "payroll (id, employee_id, period, gross, net, taxes)",
    "timesheets (id, employee_id, date, hours, project_id)",
    "leave_requests (id, employee_id, type, start, end, status)",
    "performance_reviews (id, employee_id, reviewer_id, score)",
    "training_records (id, employee_id, course, completed_at)",
    "expense_reports (id, employee_id, amount, status, date)",
    
    # Finance & Accounting
    "accounts (id, name, type, balance)",
    "transactions (id, account_id, amount, type, date)",
    "ledger_entries (id, transaction_id, debit, credit)",
    "budgets (id, department_id, year, amount)",
    "tax_records (id, year, type, amount, filed_date)",
    "bank_accounts (id, company_id, account_number, balance)",
    "bank_transfers (id, from_account, to_account, amount)",
    "credit_cards (id, customer_id, number, limit, balance)",
    "billing_cycles (id, customer_id, start, end, amount)",
    "payment_methods (id, customer_id, type, details)",
    
    # Support & Tickets
    "support_tickets (id, customer_id, subject, status, priority)",
    "ticket_comments (id, ticket_id, user_id, text, created_at)",
    "ticket_attachments (id, ticket_id, file_url)",
    "knowledge_base (id, title, content, category)",
    "faq (id, question, answer, category)",
    "chat_sessions (id, customer_id, agent_id, started_at)",
    "chat_messages (id, session_id, sender, text, timestamp)",
    "escalations (id, ticket_id, level, assigned_to)",
    "resolution_times (id, ticket_id, opened_at, resolved_at)",
    "customer_satisfaction (id, ticket_id, rating, feedback)",
    
    # Marketing
    "campaigns (id, name, type, start_date, budget)",
    "email_campaigns (id, campaign_id, subject, sent_count)",
    "email_subscribers (id, email, subscribed_at, status)",
    "email_opens (id, campaign_id, subscriber_id, opened_at)",
    "email_clicks (id, campaign_id, subscriber_id, link, clicked_at)",
    "landing_pages (id, campaign_id, url, conversions)",
    "ad_campaigns (id, platform, budget, impressions, clicks)",
    "social_posts (id, platform, content, posted_at, engagement)",
    "influencers (id, name, platform, followers, rate)",
    "affiliate_partners (id, name, commission_rate, earnings)",
    
    # Analytics & Logs
    "page_views (id, session_id, url, timestamp)",
    "user_events (id, user_id, event_type, metadata, timestamp)",
    "conversion_funnels (id, name, steps, conversion_rate)",
    "ab_tests (id, name, variant_a, variant_b, winner)",
    "error_logs (id, level, message, stack_trace, timestamp)",
    "audit_logs (id, user_id, action, resource, timestamp)",
    "api_logs (id, endpoint, method, status, duration)",
    "system_metrics (id, metric_name, value, timestamp)",
    "server_health (id, server_id, cpu, memory, disk, timestamp)",
    "job_queues (id, job_type, status, attempts, scheduled_at)",
    
    # Content Management
    "pages (id, title, slug, content, published_at)",
    "blog_posts (id, title, author_id, content, published_at)",
    "blog_categories (id, name, slug, description)",
    "comments (id, post_id, user_id, text, approved)",
    "media_files (id, filename, type, size, url)",
    "tags (id, name, slug)",
    "post_tags (post_id, tag_id)",
    "revisions (id, page_id, content, created_by, created_at)",
    "templates (id, name, html, category)",
    "widgets (id, name, type, config, position)",
    
    # Notifications
    "notifications (id, user_id, type, message, read_at)",
    "notification_settings (id, user_id, type, enabled)",
    "push_tokens (id, user_id, device, token)",
    "sms_logs (id, phone, message, status, sent_at)",
    "email_queue (id, to, subject, body, status)",
    "webhooks (id, url, events, secret, active)",
    "webhook_logs (id, webhook_id, payload, response, status)",
    "alerts (id, type, severity, message, acknowledged)",
    "alert_rules (id, condition, threshold, action)",
    "subscriptions (id, user_id, plan_id, status, expires_at)",
    
    # Settings & Config
    "settings (id, key, value, type)",
    "feature_flags (id, name, enabled, rollout_pct)",
    "configurations (id, environment, key, value)",
    "locales (id, code, name, active)",
    "translations (id, locale_id, key, value)",
    "currencies (id, code, name, exchange_rate)",
    "countries (id, code, name, region)",
    "timezones (id, name, offset)",
    "themes (id, name, colors, fonts)",
    "branding (id, logo_url, primary_color, tagline)",
]

# Add more generic tables to reach 500
PREFIXES = ["archive_", "backup_", "temp_", "staging_", "legacy_", "v2_", "test_", "dev_"]
DOMAINS = [
    "vendor", "contract", "project", "task", "milestone", "resource", "asset",
    "license", "certificate", "document", "version", "approval", "workflow",
    "rule", "policy", "compliance", "report", "dashboard", "metric", "kpi",
    "goal", "objective", "strategy", "initiative", "program", "portfolio",
    "risk", "issue", "incident", "change_request", "release", "deployment",
    "environment", "cluster", "node", "service", "endpoint", "route",
    "schema", "table_metadata", "column_metadata", "index", "constraint",
    "backup_schedule", "restore_point", "snapshot", "replica", "shard",
    "partition", "cache", "session_store", "rate_limit", "quota",
    "usage_stats", "billing_detail", "cost_center", "budget_allocation"
]

for prefix in PREFIXES:
    for domain in DOMAINS[:15]:
        TABLES.append(f"{prefix}{domain} (id, data, created_at, updated_at)")

for i in range(500 - len(TABLES)):
    TABLES.append(f"misc_table_{i:03d} (id, field1, field2, field3)")

print(f"Total tables: {len(TABLES)}")

# Test queries
QUERIES = [
    ("Which customer paid the most?", ["customers", "orders", "payments"]),
    ("Show me products that are out of stock", ["products", "inventory"]),
    ("Find employees with highest salary in engineering", ["employees", "departments", "positions"]),
    ("List all support tickets that are critical and unresolved", ["support_tickets"]),
    ("Get monthly revenue by category", ["orders", "products", "categories", "payments"]),
    ("Show user login history from last week", ["users", "login_history"]),
    ("Find all invoices that are overdue", ["invoices"]),
    ("List blog posts with most comments", ["blog_posts", "comments"]),
    ("Show campaign performance and conversions", ["campaigns", "landing_pages", "conversion_funnels"]),
    ("Get employee expense reports pending approval", ["employees", "expense_reports"]),
]

print("\n" + "="*80)
print("SCHEMA RANKER TEST: 500 Tables")
print("="*80)

# Pre-encode all tables
print("\nPre-encoding 500 tables...")
start = time.time()
table_embeddings = model.encode(TABLES)
encode_time = time.time() - start
print(f"Encoding time: {encode_time:.2f}s ({encode_time/len(TABLES)*1000:.1f}ms per table)")

for query, expected in QUERIES:
    print(f"\n{'─'*80}")
    print(f"Query: \"{query}\"")
    print(f"Expected tables: {expected}")
    
    # Rank
    start = time.time()
    query_emb = model.encode(query)
    scores = cos_sim(query_emb, table_embeddings)[0].tolist()
    rank_time = time.time() - start
    
    ranked = sorted(zip(TABLES, scores), key=lambda x: -x[1])
    
    print(f"Ranking time: {rank_time*1000:.1f}ms")
    print(f"\nTop 10 ranked tables:")
    
    found_expected = []
    for i, (table, score) in enumerate(ranked[:10]):
        table_name = table.split()[0]
        is_expected = any(exp in table_name for exp in expected)
        marker = "✓" if is_expected else " "
        if is_expected:
            found_expected.append(table_name)
        print(f"  {i+1:2d}. {marker} {table[:55]:55s} {score:.3f}")
    
    # Check how many expected were in top 10
    found = len(found_expected)
    total = len(expected)
    print(f"\n  Found {found}/{total} expected tables in top 10: {found_expected}")

print("\n" + "="*80)
print("SUMMARY")
print("="*80)
print(f"Total tables: {len(TABLES)}")
print(f"Encoding time (one-time): {encode_time:.2f}s")
print(f"Average ranking time: ~{rank_time*1000:.1f}ms per query")
print("\nConclusion: Schema Ranker is good for filtering 500→10 tables quickly!")