Merge pull request #134 from defog-ai/comext-sql-five-countries

Split EU trade SQL for more than 5 countries into several statements

Author: rishsriv

.claude-plugin/plugin.json

@@ -1,7 +1,7 @@
 {
   "name": "factiq",
   "description": "Answer economic and financial data questions with real data from FactIQ — official statistics, market data, earnings-call transcripts, executive media appearances, central-bank speeches and minutes, a curated business-news feed, and satellite-derived data (nighttime lights, shipping and port activity, reservoir levels, crop fires and wildfires with mappable footprints, NO2/SO2/CO activity, smoke and dust aerosol, monsoon rainfall) — return sourced answers, chart JSON, or report JSON.",
-  "version": "0.43.0",
+  "version": "0.44.0",
   "author": {
     "name": "FactIQ",
     "url": "https://www.factiq.com"

.codex-plugin/plugin.json

@@ -1,7 +1,7 @@
 {
   "name": "factiq",
   "description": "Answer economic and financial data questions with real FactIQ data — official data, markets, earnings calls, executive media appearances, central-bank documents, and satellite-derived signals including crop fires and wildfires — then return sourced answers, chart JSON, or report JSON.",
-  "version": "0.43.0",
+  "version": "0.44.0",
   "author": {
     "name": "FactIQ",
     "url": "https://www.factiq.com"

references/data/sql-guide.md

@@ -252,8 +252,9 @@ Exact series IDs are mandatory here, not just faster. Each country catalog
 holds millions of series, and the text collation is not byte-ordered, so
 `series_id LIKE 'prefix%'` cannot run on the index (it scans the table and
 hits the 30-second timeout) and `>=`/`<` string ranges over `series_id`
-return wrong or empty results. Equality joins on constructed IDs are index
-probes and stay fast at any scale.
+return wrong or empty results. Equality joins on constructed IDs avoid those
+scans. Their time still grows with the date range and with the number of
+reporter schemas in one statement, so keep both small.
 
 Use these rules:
 
@@ -262,8 +263,8 @@ Use these rules:
   `product_name`, `hs2_code`, `hs4_code`, or `hs6_code`.
   It is a complete product-code lookup, including historical codes and their
   first and last observed dates.
-- Build exact series IDs and join them to `data_points`, whose `series_id`
-  index makes the fetch fast. The detailed ID shape is
+- Build exact series IDs and join them to `data_points`; an exact
+  `series_id` match avoids the pattern scan. The detailed ID shape is
   `eu_comext_{M|X}_{reporter}_{partner}_{te|ti|tl}_p1_cn8_{product_code}_{eur|kg|su}`.
   `M` means imports and `X` exports. Reporter and partner are lowercase ISO2
   codes, except that `xi` is Eurostat's code for Northern Ireland. The stored
@@ -282,10 +283,13 @@ Use these rules:
   every CN8 product.
 - For several reporters, either run one already-aggregated query per reporter
   schema and combine the small results locally, or `UNION ALL` the
-  per-reporter aggregations in a single call — a union of exact-ID joins
-  across all 27 member schemas completes well within the 30-second limit.
-  What times out is pattern scanning, not combining: every branch must use
-  the exact-ID join, never a `series_id` pattern.
+  per-reporter aggregations in a single call. The API accepts at most 5
+  Comext country schemas in one statement that reads `data_points`, and
+  refuses a statement with more. For more reporters (or all 27), send
+  further groups of at most 5 one after another, not at the same time, and
+  combine the results locally. `comext_sql.py` splits a longer reporter list
+  into such groups for you. Every branch must use the exact-ID join, never a
+  `series_id` pattern.
 
 To inspect product mappings, qualify the shared table from any country query:
 
@@ -333,8 +337,8 @@ WHERE series_id = 'eu_comext_M_de_cn_te_p1_cn6_total_eur'
 ORDER BY time;
 ```
 
-For several reporters in one call, `UNION ALL` the same exact-ID aggregation
-per country schema. This tested example returns 2025 German- and
+For several reporters (at most 5) in one call, `UNION ALL` the same exact-ID
+aggregation per country schema. This tested example returns 2025 German- and
 French-reported electronics (HS chapter 85) imports from China; run it with
 `schema="eu_comext_de"` (any Comext schema works as the default):
 

references/report-patterns/bilateral-trade.md

@@ -64,7 +64,8 @@ monthly rows for short windows, annual or YTD summaries for long ones.
    only sets the default search path; qualified cross-schema references (for
    example `korea_trade.series` from an `india_trade` call) do execute, so a
    `UNION ALL` of per-reporter aggregations in one call is allowed when every
-   branch converts to a common unit. The reason to keep reporters separate is
+   branch converts to a common unit (for Eurostat Comext, at most 5 country
+   schemas per statement). The reason to keep reporters separate is
    the data, not the tool: mirror sources differ on units, valuation, timing,
    and definitions, so never mix raw values from two reporters in a single
    expression. Fetch comparable aggregates, normalize units, then compare.

scripts/comext_sql.py

@@ -4,8 +4,11 @@
 Comext series IDs must be constructed exactly (pattern scans time out; string
 ranges return wrong rows — see references/data/sql-guide.md). This emits the
 correct exact-ID SQL mechanically: the cn6_total series for all-goods totals,
-lookup-table joins for product breakdowns, and UNION ALL across any set of
-the 27 reporter schemas. Run the output with the run_sql tool using the
+lookup-table joins for product breakdowns, and UNION ALL across reporter
+schemas. One statement may read at most 5 reporter schemas, so a request for
+more reporters (or 'all') prints several statements, each with the schema to
+run it with; run them one after another and combine the results as the
+header comment says. Otherwise run the output with the run_sql tool using the
 schema printed in the header comment.
 
 ID grammar it encodes:
@@ -38,6 +41,7 @@
 import json
 import re
 import sys
+from collections.abc import Callable
 
 EU_MEMBERS = [
     "at",
@@ -90,6 +94,12 @@
 GB_EXTRA_EU_MONTH = "2020-02-01"
 NORTHERN_IRELAND_START = "2021-01-01"
 
+# The API refuses a statement that reads data_points in more than 5 Comext
+# reporter schemas, so larger reporter lists are split into statements.
+MAX_REPORTERS_PER_STATEMENT = 5
+# run_sql returns at most 50 rows.
+MAX_ROWS = 50
+
 METRICS = {"value": "eur", "weight": "kg", "supplementary": "su"}
 METRIC_UNITS = {"value": "EUR", "weight": "kg", "supplementary": "supplementary units"}
 
@@ -277,17 +287,66 @@ def product_id_expr(
     return f"'{prefix}' || lower(p.product_code) || '_{suffix}'"
 
 
+def generated_line(args: argparse.Namespace) -> str:
+    return (
+        f"-- generated by comext_sql.py {args.command}: "
+        f"{args.flow} vs partner '{args.partner}', {args.start}..{args.end}, metric={args.metric}"
+    )
+
+
 def header(reporters: list[str], args: argparse.Namespace, extra: str = "") -> str:
     lines = [
-        f"-- generated by comext_sql.py {args.command}: "
-        f"{args.flow} vs partner '{args.partner}', {args.start}..{args.end}, metric={args.metric}",
+        generated_line(args),
         f'-- run with the run_sql tool, schema="eu_comext_{reporters[0]}"',
     ]
     if extra:
         lines.append(f"-- {extra}")
     return "\n".join(lines)
 
 
+def reporter_groups(reporters: list[str]) -> list[list[str]]:
+    """Split reporters into the fewest groups of at most 5, sizes within one."""
+    count = -(-len(reporters) // MAX_REPORTERS_PER_STATEMENT)
+    size, extra = divmod(len(reporters), count)
+    groups, start = [], 0
+    for index in range(count):
+        end = start + size + (1 if index < extra else 0)
+        groups.append(reporters[start:end])
+        start = end
+    return groups
+
+
+def statements(
+    reporters: list[str],
+    args: argparse.Namespace,
+    note: str,
+    build: Callable[[list[str]], str],
+    combine: str,
+) -> str:
+    """One statement for up to 5 reporters; otherwise one per group of at most
+    5, with `combine` saying how to put the results together."""
+    if len(reporters) <= MAX_REPORTERS_PER_STATEMENT:
+        return header(reporters, args, note) + "\n" + build(reporters)
+    groups = reporter_groups(reporters)
+    lines = [generated_line(args)]
+    if note:
+        lines.append(f"-- {note}")
+    lines.append(
+        f"-- the API reads at most {MAX_REPORTERS_PER_STATEMENT} reporter schemas "
+        f"in one statement, so this request is {len(groups)} statements: run them "
+        f"one after another (not at the same time), each with the schema named "
+        f"above it, then {combine}"
+    )
+    for number, group in enumerate(groups, 1):
+        lines += [
+            "",
+            f"-- statement {number} of {len(groups)}: reporters {', '.join(group)}",
+            f'-- run with the run_sql tool, schema="eu_comext_{group[0]}"',
+            build(group),
+        ]
+    return "\n".join(lines)
+
+
 def sql_total(args: argparse.Namespace) -> str:
     reporters = resolve_reporters(args.reporters)
     partner = resolve_partner(args.partner)
@@ -302,39 +361,46 @@ def sql_total(args: argparse.Namespace) -> str:
     lo, hi, coverage_note = comext_date_bounds(args.start, args.end)
     periods, period_note = trade_periods(partner, lo, hi)
 
-    # Each reporter schema holds only its own series in data_points, so every
-    # reporter needs its own qualified branch — a flat IN list against one
-    # schema's table would silently drop the other 26.
-    branches = []
-    for r in reporters:
-        for period_partner, token, period_lo, period_hi in periods:
-            table = (
-                "data_points" if len(reporters) == 1 else f"eu_comext_{r}.data_points"
-            )
-            branches.append(
-                f"  SELECT '{r}' AS reporter, time AS t, value AS v\n"
-                f"  FROM {table}\n"
-                f"  WHERE series_id = "
-                f"'{total_series_id(flow, r, period_partner, token, suffix)}'\n"
-                f"    AND time >= DATE '{period_lo}' AND time < DATE '{period_hi}'"
-            )
-    union = "\nUNION ALL\n".join(branches)
     if args.monthly:
         select, group = (
             f"t::date AS month, SUM(v) AS total_{suffix}",
             "GROUP BY t\nORDER BY t",
         )
+        combine = f"add total_{suffix} across the statements for each month"
     else:
         select, group = (
             f"reporter, SUM(v) AS total_{suffix}",
             "GROUP BY reporter\nORDER BY 2 DESC",
         )
-    body = f"SELECT {select}\nFROM (\n{union}\n) u\n{group};"
+        combine = "put the reporter rows of all statements together"
+
+    def build(group_reporters: list[str]) -> str:
+        # Each reporter schema holds only its own series in data_points, so
+        # every reporter needs its own qualified branch — a flat IN list
+        # against one schema's table would silently drop the others.
+        branches = []
+        for r in group_reporters:
+            for period_partner, token, period_lo, period_hi in periods:
+                table = (
+                    "data_points"
+                    if len(group_reporters) == 1
+                    else f"eu_comext_{r}.data_points"
+                )
+                branches.append(
+                    f"  SELECT '{r}' AS reporter, time AS t, value AS v\n"
+                    f"  FROM {table}\n"
+                    f"  WHERE series_id = "
+                    f"'{total_series_id(flow, r, period_partner, token, suffix)}'\n"
+                    f"    AND time >= DATE '{period_lo}' AND time < DATE '{period_hi}'"
+                )
+        union = "\nUNION ALL\n".join(branches)
+        return f"SELECT {select}\nFROM (\n{union}\n) u\n{group};"
+
     note = f"all-goods total via the exact cn6_total series; values in {METRIC_UNITS[args.metric]}"
     note = add_note(note, domestic_note)
     note = add_note(note, coverage_note)
     note = add_note(note, period_note)
-    return header(reporters, args, note) + "\n" + body
+    return statements(reporters, args, note, build, combine)
 
 
 def product_branches(
@@ -395,13 +461,24 @@ def sql_products(args: argparse.Namespace) -> str:
         "cn8": "lower(p.product_code)",
     }[args.group_by]
     inner = f"{group_col} AS code, dp.value AS v"
-    union = product_branches(reporters, flow, suffix, periods, inner, "")
-    body = (
-        f"SELECT code, SUM(v) AS total_{suffix}\n"
-        f"FROM (\n{union}\n) t\n"
-        f"GROUP BY code\n"
-        f"ORDER BY 2 DESC\n"
-        f"LIMIT {args.top};"
+    # Split statements each return their top 50 so the combined ranking can
+    # be built from them; a code missing from a statement's list is not zero.
+    limit = args.top if len(reporters) <= MAX_REPORTERS_PER_STATEMENT else MAX_ROWS
+
+    def build(group_reporters: list[str]) -> str:
+        union = product_branches(group_reporters, flow, suffix, periods, inner, "")
+        return (
+            f"SELECT code, SUM(v) AS total_{suffix}\n"
+            f"FROM (\n{union}\n) t\n"
+            f"GROUP BY code\n"
+            f"ORDER BY 2 DESC\n"
+            f"LIMIT {limit};"
+        )
+
+    combine = (
+        f"add total_{suffix} across the statements for each code and keep the "
+        f"top {args.top}; each statement lists only its top {MAX_ROWS} codes, so "
+        f"a code missing from some statements has a total that is a lower bound"
     )
     label_hint = (
         "name the codes locally with the bundled hs_codes.py helper"
@@ -415,7 +492,7 @@ def sql_products(args: argparse.Namespace) -> str:
     note = add_note(note, domestic_note)
     note = add_note(note, coverage_note)
     note = add_note(note, period_note)
-    return header(reporters, args, note) + "\n" + body
+    return statements(reporters, args, note, build, combine)
 
 
 def sql_trend(args: argparse.Namespace) -> str:
@@ -452,31 +529,37 @@ def sql_trend(args: argparse.Namespace) -> str:
         if include_unit
         else "dp.time AS t, dp.value AS v"
     )
-    union = product_branches(
-        reporters,
-        flow,
-        suffix,
-        periods,
-        inner,
-        where_extra,
-        include_unit=include_unit,
-    )
-    if include_unit:
-        body = (
-            f"SELECT unit AS measurement_units, t::date AS month, SUM(v) AS total_{suffix}\n"
-            f"FROM (\n{union}\n) u\nGROUP BY unit, t\nORDER BY t, unit;"
+
+    def build(group_reporters: list[str]) -> str:
+        union = product_branches(
+            group_reporters,
+            flow,
+            suffix,
+            periods,
+            inner,
+            where_extra,
+            include_unit=include_unit,
         )
-        note = f"monthly trend for HS {code}; units returned with each row"
-    else:
-        body = (
+        if include_unit:
+            return (
+                f"SELECT unit AS measurement_units, t::date AS month, SUM(v) AS total_{suffix}\n"
+                f"FROM (\n{union}\n) u\nGROUP BY unit, t\nORDER BY t, unit;"
+            )
+        return (
             f"SELECT t::date AS month, SUM(v) AS total_{suffix}\n"
             f"FROM (\n{union}\n) u\nGROUP BY t\nORDER BY t;"
         )
+
+    if include_unit:
+        note = f"monthly trend for HS {code}; units returned with each row"
+        combine = f"add total_{suffix} across the statements for each month and unit"
+    else:
         note = f"monthly trend for HS {code}; values in {METRIC_UNITS[args.metric]}"
+        combine = f"add total_{suffix} across the statements for each month"
     note = add_note(note, domestic_note)
     note = add_note(note, coverage_note)
     note = add_note(note, period_note)
-    return header(reporters, args, note) + "\n" + body
+    return statements(reporters, args, note, build, combine)
 
 
 def main() -> None: