06/03 Optimize codes, fix Schwab Account sync error 2

app/routes/schwab.py

@@ -10,7 +10,8 @@ from app.models.account import Account
 from app.models.schwab_connection import SchwabConnection, SchwabAccount
 from app.services.schwab_service import (
     get_auth_url, exchange_code, _apply_token_data,
-    get_account_number_hashes, get_accounts, sync_preview, import_transactions,
+    get_account_number_hashes, get_accounts,
+    sync_preview, import_transactions, sync_account_snapshot,
     ACCOUNT_TYPE_MAP,
 )
 
@@ -87,12 +88,22 @@ def callback():
     for ra in raw_accounts:
         sec       = ra.get('securitiesAccount', {})
         acct_num  = sec.get('accountNumber', '')
-        acct_hash  = hash_map.get(acct_num, acct_num)   # use hashValue, fall back to raw number
+        acct_hash = hash_map.get(acct_num, acct_num)
         if not acct_hash:
             continue
-        existing = SchwabAccount.query.filter_by(account_hash=acct_hash).first()
-        if not existing:
         masked = '…' + acct_num[-4:] if len(acct_num) >= 4 else acct_num
+
+        # Try to find an existing record (by new hash or old raw number) so we
+        # can update it in-place and preserve the pfm_account_id mapping.
+        existing = (SchwabAccount.query.filter_by(account_hash=acct_hash).first() or
+                    SchwabAccount.query.filter_by(account_hash=acct_num).first() or
+                    SchwabAccount.query.filter_by(account_number_display=masked).first())
+
+        if existing:
+            existing.connection  = conn
+            existing.account_hash = acct_hash
+            existing.is_active   = True
+        else:
             db.session.add(SchwabAccount(
                 connection=conn,
                 account_hash=acct_hash,
@@ -261,6 +272,29 @@ def full_resync(schwab_account_id):
     return redirect(url_for('schwab.index'))
 
 
+# ── Balance + position snapshot ───────────────────────────────────────────────
+
+@schwab_bp.route('/snapshot/<int:schwab_account_id>', methods=['POST'])
+@login_required
+def sync_snapshot(schwab_account_id):
+    """Pull live balance and investment positions from Schwab and write to PFM."""
+    sa = db.get_or_404(SchwabAccount, schwab_account_id)
+    if not sa.pfm_account_id:
+        flash('Map this account to a PFM account first.', 'warning')
+        return redirect(url_for('schwab.index'))
+    try:
+        bal_updated, pos_synced = sync_account_snapshot(sa)
+        flash(
+            f'{sa.account_name}: balance updated'
+            f'{f", {pos_synced} position(s) synced" if pos_synced else " (no positions found)"}.',
+            'success',
+        )
+    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')
+    return redirect(url_for('schwab.index'))
+
+
 # ── Disconnect ────────────────────────────────────────────────────────────────
 
 @schwab_bp.route('/disconnect', methods=['POST'])

app/services/schwab_service.py

@@ -8,7 +8,9 @@ Auth flow:
   4. Access token expires in 30 min — auto-refresh via refresh_token (valid 7 days)
 
 Endpoints used:
-  GET /trader/v1/accounts                                         → list accounts
+  GET /trader/v1/accounts/accountNumbers                           → {accountNumber: hashValue} map
+  GET /trader/v1/accounts?fields=positions                         → list accounts with balances + positions
+  GET /trader/v1/accounts/{hash}?fields=positions                  → single account with balance + positions
   GET /trader/v1/accounts/{hash}/transactions?startDate&endDate    → transactions
 """
 
@@ -55,6 +57,21 @@ ACCOUNT_TYPE_MAP = {
     'MARGIN': 'investment',
 }
 
+# Schwab instrument asset type → PFM investment asset type
+ASSET_TYPE_MAP = {
+    'EQUITY':                'stock',
+    'ETF':                   'etf',
+    'MUTUAL_FUND':           'etf',
+    'COLLECTIVE_INVESTMENT':  'etf',
+    'INDEX':                 'etf',
+    'FIXED_INCOME':          'bond',
+    'BOND':                  'bond',
+    'CASH_EQUIVALENT':       'cash',
+    'CURRENCY':              'cash',
+    'OPTION':                'other',
+    'FUTURE':                'other',
+}
+
 
 def get_auth_url():
     cfg = current_app.config
@@ -181,6 +198,90 @@ def get_accounts(connection):
     return data
 
 
+def get_account(connection, account_hash):
+    """Fetch a single account's balance and positions by its hash."""
+    data = _authed_get(connection, f'/trader/v1/accounts/{account_hash}',
+                       params={'fields': 'positions'})
+    log.info('[schwab] get_account: hash=%s…', account_hash[:8])
+    return data
+
+
+def sync_account_snapshot(schwab_account):
+    """
+    Pull live balance and equity positions for one Schwab account and write to PFM.
+
+    Balance: sets the linked PFM account balance to Schwab's liquidationValue
+             (total portfolio value = cash + market value of all holdings).
+    Positions: upserts Investment records for every long equity/ETF/fund/bond
+               position; updates shares, avg cost, and current price.
+
+    Returns (balance_updated: bool, positions_synced: int).
+    """
+    from app.extensions import db
+    from app.models.investment import Investment
+
+    connection = schwab_account.connection
+    data       = get_account(connection, schwab_account.account_hash)
+    sec        = data.get('securitiesAccount', {})
+
+    # ── 1. Balance ────────────────────────────────────────────────────────────
+    balance_updated = False
+    if schwab_account.pfm_account:
+        balances  = sec.get('currentBalances', {})
+        liq_value = float(
+            balances.get('liquidationValue') or
+            balances.get('cashBalance') or 0
+        )
+        schwab_account.pfm_account.balance = liq_value
+        balance_updated = True
+        log.info('[schwab] balance set to %.2f for %s',
+                 liq_value, schwab_account.account_name)
+
+    # ── 2. Positions ──────────────────────────────────────────────────────────
+    positions_synced = 0
+    for pos in sec.get('positions', []):
+        instrument  = pos.get('instrument', {})
+        asset_key   = instrument.get('assetType', '')
+        symbol      = (instrument.get('symbol') or '').upper().strip()
+        long_qty    = float(pos.get('longQuantity') or 0)
+
+        pfm_type = ASSET_TYPE_MAP.get(asset_key)
+        if not pfm_type or not symbol or long_qty <= 0:
+            continue
+
+        avg_price    = float(pos.get('averagePrice') or pos.get('averageLongPrice') or 0)
+        market_value = float(pos.get('marketValue') or 0)
+        cur_price    = round(market_value / long_qty, 4) if long_qty > 0 else avg_price
+
+        inv = Investment.query.filter_by(ticker=symbol, is_active=True).first()
+        if inv:
+            inv.shares            = long_qty
+            if avg_price > 0:
+                inv.avg_cost_basis = avg_price
+            inv.current_price     = cur_price
+            inv.last_price_update = datetime.utcnow()
+        else:
+            name = (instrument.get('description') or symbol).strip()
+            inv  = Investment(
+                asset_name        = name,
+                ticker            = symbol,
+                asset_type        = pfm_type,
+                shares            = long_qty,
+                avg_cost_basis    = avg_price,
+                current_price     = cur_price,
+                last_price_update = datetime.utcnow(),
+                notes             = 'Imported from Schwab',
+            )
+            db.session.add(inv)
+
+        positions_synced += 1
+
+    db.session.commit()
+    log.info('[schwab] snapshot done for %s: positions=%d',
+             schwab_account.account_name, positions_synced)
+    return balance_updated, positions_synced
+
+
 def get_transactions(connection, account_hash, start_date, end_date):
     """
     Fetch transactions for one account.
@@ -232,10 +333,15 @@ def parse_transaction(schwab_txn, pfm_account_id, cat_id_map):
     cat_name     = CATEGORY_MAP.get(schwab_type, 'Other')
     category_id  = cat_id_map.get(cat_name)
 
-    # Parse date — Schwab uses ISO-8601 with timezone offset
+    # Parse date — Schwab uses ISO-8601 with various timezone offset forms
-    raw_time = schwab_txn.get('time', '')
+    # e.g. "2024-01-05T18:45:45+0000" or "2024-01-05T18:45:45Z"
+    raw_time = schwab_txn.get('time', '') or ''
     try:
-        txn_date = datetime.fromisoformat(raw_time.replace('Z', '+00:00')).date()
+        normalized = raw_time.replace('Z', '+00:00')
+        # Normalise +0000 → +00:00 so fromisoformat accepts it on all Python versions
+        import re as _re
+        normalized = _re.sub(r'([+-]\d{2})(\d{2})$', r'\1:\2', normalized)
+        txn_date = datetime.fromisoformat(normalized).date()
     except (ValueError, AttributeError):
         txn_date = date.today()
 

app/templates/schwab/index.html

@@ -89,10 +89,22 @@ SCHWAB_REDIRECT_URI=https://pfm.ngodanguyen.tech/schwab/callback</pre>
         </div>
         <div class="d-flex align-items-center gap-2 flex-shrink-0 ms-2">
             {% if sa.pfm_account %}
+            <!-- Sync balance + investment positions -->
+            <form method="POST" action="{{ url_for('schwab.sync_snapshot', schwab_account_id=sa.id) }}"
+                  style="display:inline;">
+                <input type="hidden" name="csrf_token" value="{{ csrf_token() }}">
+                <button type="submit" class="btn btn-sm btn-outline-success" style="font-size:11px;"
+                        title="Pull live balance and investment positions from Schwab">
+                    <i class="bi bi-arrow-clockwise me-1"></i>Balance &amp; Positions
+                </button>
+            </form>
+            <!-- Sync transactions -->
             <a href="{{ url_for('schwab.sync_preview_view', schwab_account_id=sa.id) }}"
-               class="btn btn-sm btn-outline-primary" style="font-size:11px;">
+               class="btn btn-sm btn-outline-primary" style="font-size:11px;"
-                <i class="bi bi-cloud-download me-1"></i>Sync
+               title="Import new transactions">
+                <i class="bi bi-cloud-download me-1"></i>Transactions
             </a>
+            <!-- Reset sync cursor -->
             <form method="POST" action="{{ url_for('schwab.full_resync', schwab_account_id=sa.id) }}"
                   style="display:inline;"
                   onsubmit="return confirm('Reset sync cursor for {{ sa.account_name }}?\nNext sync re-fetches 90 days. Duplicates are skipped.')">