From 5da4e74f2d30721dedbfe6c9768d75119f5771a5 Mon Sep 17 00:00:00 2001 From: Ian Dees Date: Sun, 18 Oct 2026 09:43:50 -0500 Subject: [PATCH 1/4] fail fast when a data/show request expands past the geoid limit --- census_extractomatic/api.py | 22 ++++++---- census_extractomatic/test_expand_geoids.py | 49 ++++++++++++++++++++++ 2 files changed, 63 insertions(+), 8 deletions(-) create mode 100644 census_extractomatic/test_expand_geoids.py diff --git a/census_extractomatic/api.py b/census_extractomatic/api.py index 9c85076..f3ead15 100644 --- a/census_extractomatic/api.py +++ b/census_extractomatic/api.py @@ -738,12 +738,12 @@ def show_specified_geo_data(release): release = allowed_tiger[0] if release not in allowed_tiger: abort(404, "Unknown TIGER release") - geo_ids, child_parent_map = expand_geoids(request.qwargs.geo_ids, release_to_expand_with) + max_geoids = current_app.config.get('MAX_GEOIDS_TO_SHOW', 3000) + geo_ids, child_parent_map = expand_geoids(request.qwargs.geo_ids, release_to_expand_with, max_geoids) if not geo_ids: abort(404, 'None of the geo_ids specified were valid: %s' % ', '.join(geo_ids)) - max_geoids = current_app.config.get('MAX_GEOIDS_TO_SHOW', 3000) if len(geo_ids) > max_geoids: abort(400, 'You requested %s geoids. The maximum is %s. Please contact us for bulk data.' % (len(geo_ids), max_geoids)) @@ -1433,13 +1433,14 @@ def get_child_geoids_by_prefix(release, parent_geoid, child_summary_level): return result.mappings().fetchall() -def expand_geoids(geoid_list, release): +def expand_geoids(geoid_list, release, max_geoids=None): # Look for geoid "groups" of the form `child_sumlevel|parent_geoid`. # These will expand into a list of geoids like the old comparison endpoint used to expanded_geoids = set() explicit_geoids = set() child_parent_map = {} - for geoid_str in geoid_list: + # Repeated entries would just repeat the same expansion query + for geoid_str in dict.fromkeys(geoid_list): if not expandable_geoid_re.match(geoid_str): continue @@ -1453,6 +1454,11 @@ def expand_geoids(geoid_list, release): else: explicit_geoids.add(geoid_str) + # Fail fast: don't run the remaining expansions (or retry on the next release) + # for a request that's already over the limit + if max_geoids is not None and len(expanded_geoids | explicit_geoids) > max_geoids: + abort(400, 'You requested too many geoids. The maximum is %s. Please contact us for bulk data.' % max_geoids) + # Since the expanded geoids were sourced from the database they don't need to be checked valid_geo_ids = set(expanded_geoids) @@ -1501,9 +1507,10 @@ def show_specified_data(acs): expand_errors = [] valid_geo_ids = set() requested_geo_ids = request.qwargs.geo_ids + max_geoids = current_app.config.get('MAX_GEOIDS_TO_SHOW', 1000) for release in acs_to_try: try: - this_valid_geo_ids, child_parent_map = expand_geoids(requested_geo_ids, release) + this_valid_geo_ids, child_parent_map = expand_geoids(requested_geo_ids, release, max_geoids) if this_valid_geo_ids: releases_to_use.append(release) @@ -1518,7 +1525,6 @@ def show_specified_data(acs): if not valid_geo_ids: abort(404, 'None of the geo_ids specified were valid: %s' % ', '.join(requested_geo_ids)) - max_geoids = current_app.config.get('MAX_GEOIDS_TO_SHOW', 1000) if len(valid_geo_ids) > max_geoids: abort(400, 'You requested %s geoids. The maximum is %s. Please contact us for bulk data.' % (len(valid_geo_ids), max_geoids)) @@ -1892,9 +1898,10 @@ def download_specified_data(acs): expand_errors = [] valid_geo_ids = set() requested_geo_ids = request.qwargs.geo_ids + max_geoids = current_app.config.get('MAX_GEOIDS_TO_DOWNLOAD', 1000) for release in acs_to_try: try: - this_valid_geo_ids, child_parent_map = expand_geoids(requested_geo_ids, release) + this_valid_geo_ids, child_parent_map = expand_geoids(requested_geo_ids, release, max_geoids) if this_valid_geo_ids: releases_to_use.append(release) @@ -1909,7 +1916,6 @@ def download_specified_data(acs): if not valid_geo_ids: abort(404, 'None of the geo_ids specified were valid: %s' % ', '.join(valid_geo_ids)) - max_geoids = current_app.config.get('MAX_GEOIDS_TO_DOWNLOAD', 1000) if len(valid_geo_ids) > max_geoids: abort(400, 'You requested %s geoids. The maximum is %s. Please contact us for bulk data.' % (len(valid_geo_ids), max_geoids)) diff --git a/census_extractomatic/test_expand_geoids.py b/census_extractomatic/test_expand_geoids.py new file mode 100644 index 0000000..ec847a2 --- /dev/null +++ b/census_extractomatic/test_expand_geoids.py @@ -0,0 +1,49 @@ +"""expand_geoids should stop as soon as a request exceeds the geoid limit, instead of +running every remaining group expansion (and then repeating them for the next release).""" +import os +import sys +from unittest import mock + +os.environ.setdefault('DATABASE_URL', 'sqlite://') +os.environ.setdefault('REDIS_URL', 'redis://localhost:6379/0') +os.environ.setdefault('CACHE_TYPE', 'NullCache') + +# api.py imports GDAL (via user_geo), which isn't needed here and isn't always installed +try: + import osgeo # noqa: F401 +except ImportError: + sys.modules['osgeo'] = mock.MagicMock() + +import pytest +from werkzeug.exceptions import BadRequest + +from census_extractomatic import api + + +@pytest.fixture +def expansion_calls(monkeypatch): + calls = [] + + def fake_get_child_geoids(release, parent_geoid, child_summary_level): + calls.append(child_summary_level) + return [{'geoid': '%s00US%05d' % (child_summary_level, i)} for i in range(6)] + + monkeypatch.setattr(api, 'get_child_geoids', fake_get_child_geoids) + return calls + + +def test_stops_expanding_once_over_the_limit(expansion_calls): + groups = ['040|01000US', '050|01000US', '140|01000US', '150|01000US'] + with api.app.app_context(), pytest.raises(BadRequest): + api.expand_geoids(groups, 'acs2024_5yr', max_geoids=10) + + # 6 + 6 geoids passes the limit of 10 after the second group + assert expansion_calls == ['040', '050'] + + +def test_repeated_groups_are_only_expanded_once(expansion_calls): + with api.app.app_context(): + valid_geo_ids, child_parent_map = api.expand_geoids(['040|01000US'] * 5, 'acs2024_5yr') + + assert expansion_calls == ['040'] + assert len(valid_geo_ids) == 6 From e955c478f74c6bda227e2f403ae2a396c87da6e6 Mon Sep 17 00:00:00 2001 From: Ian Dees Date: Sun, 18 Oct 2026 09:44:22 -0500 Subject: [PATCH 2/4] add a 5s Postgres statement timeout for interactive endpoints and return 503 on timeout --- census_extractomatic/api.py | 33 +++++++++ .../test_statement_timeout.py | 73 +++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 census_extractomatic/test_statement_timeout.py diff --git a/census_extractomatic/api.py b/census_extractomatic/api.py index f3ead15..3ec1ee9 100644 --- a/census_extractomatic/api.py +++ b/census_extractomatic/api.py @@ -15,7 +15,9 @@ from flask_caching import Cache from flask_cors import CORS, cross_origin from flask_sqlalchemy import SQLAlchemy +from psycopg2.errors import QueryCanceled from sqlalchemy import text +from sqlalchemy.exc import DBAPIError from itertools import groupby from werkzeug.exceptions import HTTPException import math @@ -224,11 +226,42 @@ table_re = re.compile(r"^[BC]\d{5,6}(?:[A-Z]{1,3})?$") +# Interactive endpoints get a short Postgres statement timeout so one expensive query can't +# hold a DB backend (and a gunicorn worker) for minutes. Bulk endpoints keep the server default. +INTERACTIVE_STATEMENT_TIMEOUT_MS = 5000 +NO_STATEMENT_TIMEOUT_PREFIXES = ( + '/1.0/data/download/', + '/1.0/aggregate/', + '/1.0/user_geo/', +) + + +@app.before_request +def set_interactive_statement_timeout(): + timeout_ms = current_app.config.get('INTERACTIVE_STATEMENT_TIMEOUT_MS', INTERACTIVE_STATEMENT_TIMEOUT_MS) + path = request.path + if not timeout_ms or not (path.startswith('/1.') or path.startswith('/2.')): + return + if path.startswith(NO_STATEMENT_TIMEOUT_PREFIXES): + return + + # SET LOCAL only lasts for this request's transaction, so it can't leak to other + # requests through the connection pool + db.session.execute(text("SET LOCAL statement_timeout = %d" % int(timeout_ms))) + + @app.errorhandler(400) @app.errorhandler(404) @app.errorhandler(500) @cross_origin(origins='*') def jsonify_error_handler(error): + original_error = getattr(error, 'original_exception', None) + if isinstance(original_error, DBAPIError) and isinstance(original_error.orig, QueryCanceled): + app.logger.warning("Statement timeout on %s", request.full_path) + resp = jsonify(error="That query took too long. Try requesting fewer geographies or tables.") + resp.status_code = 503 + return resp + if isinstance(error, ClientRequestValidationException): resp = jsonify(error=error.description, errors=error.errors) resp.status_code = error.code diff --git a/census_extractomatic/test_statement_timeout.py b/census_extractomatic/test_statement_timeout.py new file mode 100644 index 0000000..7f8dccf --- /dev/null +++ b/census_extractomatic/test_statement_timeout.py @@ -0,0 +1,73 @@ +"""Interactive endpoints get a Postgres statement timeout; bulk endpoints don't, and a +timed-out query becomes a 503 rather than a generic 500.""" +import os +import sys +from unittest import mock + +os.environ.setdefault('DATABASE_URL', 'sqlite://') +os.environ.setdefault('REDIS_URL', 'redis://localhost:6379/0') +os.environ.setdefault('CACHE_TYPE', 'NullCache') + +# api.py imports GDAL (via user_geo), which isn't needed here and isn't always installed +try: + import osgeo # noqa: F401 +except ImportError: + sys.modules['osgeo'] = mock.MagicMock() + +import pytest +from psycopg2.errors import QueryCanceled +from sqlalchemy.exc import OperationalError + +from census_extractomatic import api + + +@pytest.fixture +def executed_sql(monkeypatch): + statements = [] + monkeypatch.setattr(api.db.session, 'execute', lambda stmt, *a, **kw: statements.append(str(stmt))) + return statements + + +@pytest.mark.parametrize('path', [ + '/1.0/data/show/latest', + '/1.0/geo/search', + '/2.1/full-text/search', +]) +def test_interactive_paths_get_a_statement_timeout(path, executed_sql): + with api.app.test_request_context(path): + api.set_interactive_statement_timeout() + + assert executed_sql == ['SET LOCAL statement_timeout = 5000'] + + +@pytest.mark.parametrize('path', [ + '/1.0/data/download/latest', + '/1.0/aggregate/acs/latest', + '/1.0/user_geo/import', + '/healthcheck', + '/robots.txt', +]) +def test_bulk_and_non_db_paths_are_left_alone(path, executed_sql): + with api.app.test_request_context(path): + api.set_interactive_statement_timeout() + + assert executed_sql == [] + + +def test_statement_timeout_can_be_disabled(executed_sql, monkeypatch): + monkeypatch.setitem(api.app.config, 'INTERACTIVE_STATEMENT_TIMEOUT_MS', 0) + with api.app.test_request_context('/1.0/data/show/latest'): + api.set_interactive_statement_timeout() + + assert executed_sql == [] + + +def test_timed_out_query_returns_503(): + error = mock.Mock() + error.original_exception = OperationalError('SELECT 1', {}, QueryCanceled('canceling statement due to statement timeout')) + + with api.app.test_request_context('/1.0/data/show/latest'): + resp = api.jsonify_error_handler(error) + + assert resp.status_code == 503 + assert 'took too long' in resp.get_json()['error'] From 6cef2e1fcdef924c0f3dcdec189c6b15b1476506 Mon Sep 17 00:00:00 2001 From: Ian Dees Date: Sun, 18 Oct 2026 09:57:40 -0500 Subject: [PATCH 3/4] reject known oversized geoid expansions without querying the database --- DATA_UPDATES.md | 1 + census_extractomatic/api.py | 51 ++++++++++++++++++++++ census_extractomatic/test_expand_geoids.py | 21 +++++++++ 3 files changed, 73 insertions(+) diff --git a/DATA_UPDATES.md b/DATA_UPDATES.md index 54cfba7..289c272 100644 --- a/DATA_UPDATES.md +++ b/DATA_UPDATES.md @@ -153,6 +153,7 @@ If this is a new release year, you'll want to set up the new TIGER geodata scrip - Add the new release to the `allowed_releases` list. Usually you want to replace the existing release with the newer one you just imported. You might want to move the older release of the same type to the bottom of the list so that people can use it for specific requests. - If you are updating a 5yr release, you probably want to update the `default_table_search_release` variables, too. - Add an entry for the `ACS_NAMES` dict for the new release. + - Optionally refresh `KNOWN_LARGE_EXPANSIONS` (the query is in the comment above it). Counts change very little between years, so this can wait. - Commit the changes - Push to the `dokku.censusreporter.org` remote - `git push dokku` diff --git a/census_extractomatic/api.py b/census_extractomatic/api.py index 3ec1ee9..42af79b 100644 --- a/census_extractomatic/api.py +++ b/census_extractomatic/api.py @@ -122,6 +122,50 @@ '150': ['101'], } +# Sizes of `sumlevel|parent_geoid` expansions in the 5-year releases that are too big to +# ever be useful, so requests for them can be rejected without querying the database. +# Only expansions over 5000 geoids are listed; the counts are only compared to the +# MAX_GEOIDS_* limits, so they don't need to track each release exactly. Measured on +# acs2024_5yr. To refresh (they change very little year to year), count geoheader rows: +# SELECT sumlevel, count(*) FROM geoheader WHERE component='00' GROUP BY 1 HAVING count(*) > 5000; +# and the same grouped by the state/county FIPS in the geoid for state/county parents. +KNOWN_LARGE_EXPANSIONS = { + # nation -> all + ('01000US', '150'): 242297, + ('01000US', '140'): 85382, + ('01000US', '070'): 73111, + ('01000US', '060'): 36421, + ('01000US', '155'): 33791, + ('01000US', '860'): 33772, + ('01000US', '160'): 32330, + ('01000US', '970'): 10901, + ('01000US', '622'): 8461, + ('01000US', '612'): 5557, + # states -> block groups / tracts + ('04000US06', '150'): 25607, + ('04000US48', '150'): 18638, + ('04000US36', '150'): 16031, + ('04000US12', '150'): 13388, + ('04000US42', '150'): 10173, + ('04000US17', '150'): 9898, + ('04000US39', '150'): 9472, + ('04000US26', '150'): 8386, + ('04000US13', '150'): 7446, + ('04000US37', '150'): 7111, + ('04000US34', '150'): 6599, + ('04000US51', '150'): 5963, + ('04000US53', '150'): 5311, + ('04000US18', '150'): 5290, + ('04000US25', '150'): 5116, + ('04000US29', '150'): 5031, + ('04000US06', '140'): 9129, + ('04000US48', '140'): 6896, + ('04000US36', '140'): 5397, + ('04000US12', '140'): 5160, + # counties -> block groups + ('05000US06037', '150'): 6591, +} + SUMLEV_NAMES = { "010": {"name": "nation", "plural": ""}, "020": {"name": "region", "plural": "regions"}, @@ -1480,6 +1524,13 @@ def expand_geoids(geoid_list, release, max_geoids=None): geoid_split = geoid_str.split('|') if len(geoid_split) == 2 and len(geoid_split[0]) == 3: (child_summary_level, parent_geoid) = geoid_split + + # Skip the database for expansions we already know are over the limit. + # The 1-year releases only cover places with 65k+ people, so they're always small. + known_size = KNOWN_LARGE_EXPANSIONS.get((parent_geoid.upper(), child_summary_level)) + if max_geoids is not None and release.endswith('_5yr') and known_size and known_size > max_geoids: + abort(400, 'You requested too many geoids. The maximum is %s. Please contact us for bulk data.' % max_geoids) + child_geoid_list = [child_geoid['geoid'] for child_geoid in get_child_geoids(release, parent_geoid, child_summary_level)] expanded_geoids.update(child_geoid_list) for child_geoid in child_geoid_list: diff --git a/census_extractomatic/test_expand_geoids.py b/census_extractomatic/test_expand_geoids.py index ec847a2..4259253 100644 --- a/census_extractomatic/test_expand_geoids.py +++ b/census_extractomatic/test_expand_geoids.py @@ -47,3 +47,24 @@ def test_repeated_groups_are_only_expanded_once(expansion_calls): assert expansion_calls == ['040'] assert len(valid_geo_ids) == 6 + + +def test_known_large_expansion_is_rejected_without_querying(expansion_calls): + with api.app.app_context(), pytest.raises(BadRequest): + api.expand_geoids(['150|01000US'], 'acs2024_5yr', max_geoids=10000) + + assert expansion_calls == [] + + +def test_known_large_expansion_is_still_expanded_for_1yr_releases(expansion_calls): + with api.app.app_context(): + api.expand_geoids(['150|01000US'], 'acs2024_1yr', max_geoids=10000) + + assert expansion_calls == ['150'] + + +def test_known_large_expansion_under_a_higher_limit_is_expanded(expansion_calls): + with api.app.app_context(): + api.expand_geoids(['150|04000US17'], 'acs2024_5yr', max_geoids=20000) + + assert expansion_calls == ['150'] From 8a6502aeb44a2bfedbc2072f035e2c11ca8b8175 Mon Sep 17 00:00:00 2001 From: Ian Dees Date: Sun, 18 Oct 2026 11:55:21 -0500 Subject: [PATCH 4/4] apply the known-large-expansion lookup to all releases --- census_extractomatic/api.py | 7 +++---- census_extractomatic/test_expand_geoids.py | 6 +++--- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/census_extractomatic/api.py b/census_extractomatic/api.py index 42af79b..60f9a6b 100644 --- a/census_extractomatic/api.py +++ b/census_extractomatic/api.py @@ -122,7 +122,7 @@ '150': ['101'], } -# Sizes of `sumlevel|parent_geoid` expansions in the 5-year releases that are too big to +# Sizes of `sumlevel|parent_geoid` expansions that are too big to # ever be useful, so requests for them can be rejected without querying the database. # Only expansions over 5000 geoids are listed; the counts are only compared to the # MAX_GEOIDS_* limits, so they don't need to track each release exactly. Measured on @@ -1525,10 +1525,9 @@ def expand_geoids(geoid_list, release, max_geoids=None): if len(geoid_split) == 2 and len(geoid_split[0]) == 3: (child_summary_level, parent_geoid) = geoid_split - # Skip the database for expansions we already know are over the limit. - # The 1-year releases only cover places with 65k+ people, so they're always small. + # Skip the database for expansions we already know are over the limit known_size = KNOWN_LARGE_EXPANSIONS.get((parent_geoid.upper(), child_summary_level)) - if max_geoids is not None and release.endswith('_5yr') and known_size and known_size > max_geoids: + if max_geoids is not None and known_size and known_size > max_geoids: abort(400, 'You requested too many geoids. The maximum is %s. Please contact us for bulk data.' % max_geoids) child_geoid_list = [child_geoid['geoid'] for child_geoid in get_child_geoids(release, parent_geoid, child_summary_level)] diff --git a/census_extractomatic/test_expand_geoids.py b/census_extractomatic/test_expand_geoids.py index 4259253..3599853 100644 --- a/census_extractomatic/test_expand_geoids.py +++ b/census_extractomatic/test_expand_geoids.py @@ -56,11 +56,11 @@ def test_known_large_expansion_is_rejected_without_querying(expansion_calls): assert expansion_calls == [] -def test_known_large_expansion_is_still_expanded_for_1yr_releases(expansion_calls): - with api.app.app_context(): +def test_known_large_expansion_is_rejected_for_1yr_releases_too(expansion_calls): + with api.app.app_context(), pytest.raises(BadRequest): api.expand_geoids(['150|01000US'], 'acs2024_1yr', max_geoids=10000) - assert expansion_calls == ['150'] + assert expansion_calls == [] def test_known_large_expansion_under_a_higher_limit_is_expanded(expansion_calls):