06/03 Optimize codes, fix Schwab Account

app/routes/accounts.py

@@ -60,19 +60,28 @@ def index():
         tab = 'bank'
 
     from app.models.teller_enrollment import TellerAccount
+    from app.models.schwab_connection import SchwabAccount
     all_accounts = Account.query.filter_by(is_active=True).order_by(Account.name).all()
+    pfm_ids = [a.id for a in all_accounts]
 
-    # Build Teller map first so we can skip calc_balance for Teller-linked accounts.
-    # Their balance comes from Teller (live refresh or transaction sync) and must not
-    # be overwritten by the transaction-computed sum on every page load.
+    # Build sync-provider maps before the calc_balance loop.
+    # Accounts linked to Teller or Schwab get their balance from the provider,
+    # not from transaction summation, so we skip calc_balance for them.
     teller_accounts = TellerAccount.query.filter(
-        TellerAccount.pfm_account_id.in_([a.id for a in all_accounts]),
+        TellerAccount.pfm_account_id.in_(pfm_ids),
         TellerAccount.is_active == True,
     ).all()
     teller_map = {ta.pfm_account_id: ta for ta in teller_accounts}
 
+    schwab_accounts = SchwabAccount.query.filter(
+        SchwabAccount.pfm_account_id.in_(pfm_ids),
+        SchwabAccount.is_active == True,
+    ).all()
+    schwab_map = {sa.pfm_account_id: sa for sa in schwab_accounts}
+
+    provider_ids = set(teller_map) | set(schwab_map)
     for a in all_accounts:
-        if a.id not in teller_map:
+        if a.id not in provider_ids:
             calc_balance(a.id)
     all_accounts = Account.query.filter_by(is_active=True).order_by(Account.name).all()
 
@@ -101,7 +110,8 @@ def index():
                            bank_count=len(bank_accounts),
                            credit_count=len(credit_accounts),
                            monthly_charges=monthly_charges,
-                           teller_map=teller_map)
+                           teller_map=teller_map,
+                           schwab_map=schwab_map)
 
 
 @accounts_bp.route('/new', methods=['GET', 'POST'])

app/routes/schwab.py

@@ -292,6 +292,9 @@ def sync_snapshot(schwab_account_id):
     except Exception as e:
         log.error('[schwab] sync_snapshot failed for id=%s: %s', schwab_account_id, e, exc_info=True)
         flash(f'Snapshot sync failed: {e}', 'danger')
+    next_url = request.form.get('next', '')
+    if next_url and next_url.startswith('/'):
+        return redirect(next_url)
     return redirect(url_for('schwab.index'))