Initial Checkin

backend/flex_query.py

@@ -0,0 +1,141 @@
+"""Utilities for Interactive Brokers Flex Web Service integration.
+
+This module provides the first layer of infrastructure for pulling
+historical trade data from a saved Flex Query. It intentionally keeps the
+surface area small so the app can validate configuration, trigger a sync,
+and persist the raw XML responses locally before the data is wired into the
+lot engine.
+"""
+from __future__ import annotations
+
+import os
+import time
+import urllib.parse
+import urllib.request
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass
+from pathlib import Path
+
+
+FLEX_BASE_URL = "https://ndcdyn.interactivebrokers.com/AccountManagement/FlexWebService"
+FLEX_SEND_REQUEST_URL = f"{FLEX_BASE_URL}/SendRequest"
+FLEX_GET_STATEMENT_URL = f"{FLEX_BASE_URL}/GetStatement"
+
+
+@dataclass
+class FlexConfig:
+    token: str
+    query_id: str
+    data_dir: Path
+    poll_interval_seconds: float = 2.0
+    max_polls: int = 30
+
+
+class FlexQueryError(RuntimeError):
+    pass
+
+
+def load_flex_config() -> FlexConfig | None:
+    token = os.environ.get("IB_FLEX_TOKEN", "").strip()
+    query_id = os.environ.get("IB_FLEX_QUERY_ID", "").strip()
+    if not token or not query_id:
+        return None
+    data_dir = Path(os.environ.get("IB_FLEX_DATA_DIR", Path(__file__).resolve().parent.parent / "data" / "flex"))
+    return FlexConfig(token=token, query_id=query_id, data_dir=data_dir)
+
+
+def get_flex_status() -> dict:
+    config = load_flex_config()
+    if config is None:
+        return {
+            "configured": False,
+            "missing": [
+                name
+                for name in ("IB_FLEX_TOKEN", "IB_FLEX_QUERY_ID")
+                if not os.environ.get(name, "").strip()
+            ],
+        }
+    latest_file = None
+    latest_mtime = None
+    if config.data_dir.exists():
+        xml_files = sorted(config.data_dir.glob("flex_*.xml"))
+        if xml_files:
+            latest = xml_files[-1]
+            latest_file = str(latest)
+            latest_mtime = latest.stat().st_mtime
+    return {
+        "configured": True,
+        "query_id": config.query_id,
+        "data_dir": str(config.data_dir),
+        "latest_file": latest_file,
+        "latest_timestamp": latest_mtime,
+    }
+
+
+def _http_get(url: str, params: dict[str, str]) -> bytes:
+    full_url = f"{url}?{urllib.parse.urlencode(params)}"
+    request = urllib.request.Request(
+        full_url,
+        headers={"User-Agent": "IB-Dashboard/1.0"},
+    )
+    with urllib.request.urlopen(request, timeout=30) as response:
+        return response.read()
+
+
+def _parse_send_request(xml_bytes: bytes) -> tuple[str, str]:
+    root = ET.fromstring(xml_bytes)
+    status = (root.findtext("Status") or "").strip()
+    if status.lower() != "success":
+        raise FlexQueryError(root.findtext("ErrorMessage") or "Flex send-request failed")
+    reference_code = (root.findtext("ReferenceCode") or "").strip()
+    if not reference_code:
+        raise FlexQueryError("Flex send-request returned no ReferenceCode")
+    return status, reference_code
+
+
+def _parse_statement_status(xml_bytes: bytes) -> tuple[bool, str]:
+    root = ET.fromstring(xml_bytes)
+    if root.tag == "FlexStatementResponse":
+        error_message = (root.findtext("ErrorMessage") or "").strip()
+        if error_message:
+            raise FlexQueryError(error_message)
+        return False, ""
+    return True, xml_bytes.decode("utf-8", errors="replace")
+
+
+def fetch_flex_statement() -> dict:
+    config = load_flex_config()
+    if config is None:
+        raise FlexQueryError("Flex Query is not configured")
+
+    config.data_dir.mkdir(parents=True, exist_ok=True)
+
+    _, reference_code = _parse_send_request(_http_get(
+        FLEX_SEND_REQUEST_URL,
+        {"t": config.token, "q": config.query_id, "v": "3"},
+    ))
+
+    statement_xml = None
+    for _ in range(config.max_polls):
+        xml_bytes = _http_get(
+            FLEX_GET_STATEMENT_URL,
+            {"t": config.token, "q": reference_code, "v": "3"},
+        )
+        done, payload = _parse_statement_status(xml_bytes)
+        if done:
+            statement_xml = payload
+            break
+        time.sleep(config.poll_interval_seconds)
+
+    if statement_xml is None:
+        raise FlexQueryError("Timed out waiting for Flex statement generation")
+
+    timestamp = time.strftime("%Y%m%d_%H%M%S")
+    destination = config.data_dir / f"flex_{timestamp}.xml"
+    destination.write_text(statement_xml, encoding="utf-8")
+
+    return {
+        "saved_to": str(destination),
+        "reference_code": reference_code,
+        "bytes": len(statement_xml.encode("utf-8")),
+    }