From 1f296fc8b9b4fbd85bab2fe35b5a2751aea89cb2 Mon Sep 17 00:00:00 2001 From: Santiago Garces Escobar Date: Thu, 24 Sep 2026 23:01:51 -0400 Subject: [PATCH] feat(arcgis): aggregate_data tool (outStatistics + groupBy) Totals by area needed every row pulled and summed client-side: five-county voucher sums took 23 paged query_data calls. aggregate_data sends one outStatistics query with groupByFieldsForStatistics instead. - statistics: [{type: count|sum|avg|min|max|stddev, field, as}]; a count without a field counts records by the object-id field. - group_by, where, having, order_by, limit, layer, format (text|json|csv). - Field and group names must be plain identifiers and exist in the layer schema (case-insensitive; sent with the schema's spelling). where and having use WhereValidator; order_by uses validate_order_by. All input checks run before any request. - Layers reporting supportsStatistics: false are refused. - The guidance line says nulls are excluded, so sums over suppressed values are lower bounds. get_schema, query_data and aggregate_data now share _layer_url_for (item -> trusted layer URL) and _layer_fields (fields + metadata, with the one-row-query fallback). Closes #31 Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/BUILT_IN_PLUGINS.md | 2 + plugins/arcgis/plugin.py | 433 ++++++++++++++++-- .../security/test_arcgis_aggregate_guards.py | 81 ++++ .../plugins/arcgis/test_aggregate_data.py | 246 ++++++++++ .../unit/plugins/arcgis/test_arcgis_plugin.py | 5 +- 5 files changed, 721 insertions(+), 46 deletions(-) create mode 100644 tests/security/test_arcgis_aggregate_guards.py create mode 100644 tests/unit/plugins/arcgis/test_aggregate_data.py diff --git a/docs/BUILT_IN_PLUGINS.md b/docs/BUILT_IN_PLUGINS.md index 6726eeb..8ff549d 100644 --- a/docs/BUILT_IN_PLUGINS.md +++ b/docs/BUILT_IN_PLUGINS.md @@ -88,6 +88,7 @@ plugins: | `arcgis__get_aggregations(field, q)` | Facet counts for type, tags, categories, or access | | `arcgis__get_schema(dataset_id, layer)` | Field names, types and aliases for a layer or table | | `arcgis__query_data(dataset_id, layer, where, out_fields, limit, offset, order_by, format)` | Query a Feature Service, with paging and a total match count | +| `arcgis__aggregate_data(dataset_id, statistics, group_by, where, having, order_by, layer, limit, format)` | Server-side count/sum/avg/min/max/stddev, grouped by fields | ### Usage Notes @@ -95,6 +96,7 @@ plugins: - `get_aggregations` accepts `field` values: `"type"`, `"tags"`, `"categories"`, `"access"`. This is a catalog-level tool, not a DataPlugin method — it has no equivalent in other plugins. - `query_data` uses the ArcGIS Feature Service query interface. The `where` parameter is a SQL WHERE clause (e.g., `"population > 10000"`). Only Feature Layer, Feature Service, Map Service, and Table types are queryable. - `layer`: a Feature Service can hold several layers and tables. `get_dataset` lists them (`0: SAFMR_Zip_Code_Tab_Areas (polygon layer) [default]`, `1: SAFMR_table (table)`), and `get_schema`/`query_data` take `layer` to reach one other than the default. The default is the layer the item's URL names, otherwise the service's first layer (then first table). An id the service does not list is refused with the list of valid ids. +- `aggregate_data` runs one `outStatistics` query instead of paging rows to sum them client-side. `statistics` is a list of `{"type": "sum", "field": "HCV_PUBLIC", "as": "vouchers"}` (types: count, sum, avg, min, max, stddev; a `count` without a field counts records by the object-id field), `group_by` a list of fields. Field names must exist in the layer's schema (matched case-insensitively); `where` and `having` go through the WHERE validator and `order_by` through the sort-order validator. Nulls are excluded from every statistic, so a sum over suppressed values is a lower bound. Layers that report `supportsStatistics: false` are refused with a pointer to `query_data`. - `query_data` paging: the reply starts with `Returned N of M matching record(s) (offset: O, limit: L)` and, when more records match, ends with `Next page: offset=…`. `offset` maps to `resultOffset` and `order_by` to `orderByFields` (field names, each optionally `ASC`/`DESC`; anything else is refused). Pass `order_by` for stable paging. The total comes from a `returnCountOnly` request, which is skipped when a first page comes back short of `limit`. - `query_data` output: `format` is `text` (default, `Record N:` blocks), `json` (an array with one object per line) or `csv` (a header row plus one line per record). Every returned record is shown; if the reply would pass the response size limit, whole records are dropped from the end, a `Showing N of M record(s)` notice is added, and the next offset accounts for it. diff --git a/plugins/arcgis/plugin.py b/plugins/arcgis/plugin.py index fd363af..a9cd33d 100644 --- a/plugins/arcgis/plugin.py +++ b/plugins/arcgis/plugin.py @@ -5,6 +5,7 @@ """ import ipaddress +import json import logging import re from typing import Any @@ -31,6 +32,13 @@ r"^(?P.+/(?:FeatureServer|MapServer))(?:/(?P\d+))?/?$", re.IGNORECASE, ) +# Statistic types accepted by aggregate_data (ArcGIS ``statisticType``). +STATISTIC_TYPES = ("count", "sum", "avg", "min", "max", "stddev") +# A plain field name or output alias. +_FIELD_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]{0,63}$") +# Upper bounds on aggregate_data inputs. +_MAX_STATISTICS = 20 +_MAX_GROUP_BY = 10 # Hostname suffixes that are never valid public service hosts. _BLOCKED_HOST_SUFFIXES = (".local", ".localhost", ".internal", ".localdomain") @@ -340,6 +348,98 @@ def get_tools(self) -> list[ToolDefinition]: "required": ["dataset_id"], }, ), + ToolDefinition( + name="aggregate_data", + description=( + "Compute statistics (count, sum, avg, min, max, stddev) on an " + "ArcGIS Feature Service layer, optionally grouped by fields, in " + "one server-side call. Use it for totals and counts by area " + "instead of paging through query_data. Null values are left out " + "of every statistic, so a sum over a field with suppressed " + "values is a lower bound. Use get_schema first for field names." + ), + input_schema={ + "type": "object", + "properties": { + "dataset_id": { + "type": "string", + "description": "Hub item ID (same as get_dataset)", + }, + "layer": { + "type": "integer", + "description": ( + "Layer or table id within the Feature Service, as " + "listed by get_dataset. Default: the layer the item " + "points to, or the service's first layer." + ), + "minimum": 0, + }, + "statistics": { + "type": "array", + "description": ( + "Statistics to compute, e.g. " + '[{"type": "sum", "field": "HCV_PUBLIC", "as": ' + '"vouchers"}, {"type": "count", "field": "GEOID"}]. ' + "count may omit field to count records." + ), + "items": { + "type": "object", + "properties": { + "type": { + "type": "string", + "enum": list(STATISTIC_TYPES), + }, + "field": {"type": "string"}, + "as": {"type": "string"}, + }, + "required": ["type"], + }, + "minItems": 1, + "maxItems": _MAX_STATISTICS, + }, + "group_by": { + "type": "array", + "items": {"type": "string"}, + "description": "Fields to group by (one row per group)", + "maxItems": _MAX_GROUP_BY, + }, + "where": { + "type": "string", + "description": "SQL WHERE clause applied before grouping", + "default": "1=1", + }, + "having": { + "type": "string", + "description": ( + "Filter on the computed statistics, e.g. " + '"SUM(HCV_PUBLIC) > 100" (only on services that ' + "support it)" + ), + }, + "order_by": { + "type": "string", + "description": ( + "Sort by group fields or statistic names, each " + 'optionally ASC or DESC (e.g. "vouchers DESC")' + ), + }, + "limit": { + "type": "integer", + "description": "Maximum number of groups (default: 1000)", + "default": 1000, + "minimum": 1, + "maximum": 5000, + }, + "format": { + "type": "string", + "enum": list(ROW_FORMATS), + "description": "Output format: text, json or csv", + "default": "text", + }, + }, + "required": ["dataset_id", "statistics"], + }, + ), ] def tool_handlers(self) -> dict[str, ToolHandler]: @@ -375,6 +475,15 @@ def tool_handlers(self) -> dict[str, ToolHandler]: "query_data": ToolHandler( handler=self._tool_query_data, required_args=("dataset_id",) ), + "aggregate_data": ToolHandler( + handler=self._tool_aggregate_data, + required_args=("dataset_id", "statistics"), + guidance=( + "Statistics ignore null values: where a source suppresses " + "small counts as null, a sum is a lower bound and a count of " + "a field counts only non-null values." + ), + ), } async def _tool_search_datasets(self, arguments: dict[str, Any]) -> ToolResult: @@ -452,6 +561,27 @@ async def _tool_query_data(self, arguments: dict[str, Any]) -> ToolResult: success=True, ) + async def _tool_aggregate_data(self, arguments: dict[str, Any]) -> ToolResult: + fmt = arguments.get("format") or "text" + if fmt not in ROW_FORMATS: + raise ValueError( + f"format must be one of {', '.join(ROW_FORMATS)} (got {fmt!r})" + ) + result = await self.aggregate_features( + arguments["dataset_id"], + arguments["statistics"], + group_by=arguments.get("group_by"), + where=arguments.get("where", "1=1"), + having=arguments.get("having"), + order_by=arguments.get("order_by"), + limit=arguments.get("limit", 1000), + layer=arguments.get("layer"), + ) + return ToolResult( + content=[{"type": "text", "text": self._format_aggregate(result, fmt=fmt)}], + success=True, + ) + # ── DataPlugin abstract method implementations ────────────────────── async def search_datasets( @@ -545,21 +675,62 @@ async def get_schema( Raises: ValueError: If the dataset has no queryable service URL. """ + layer_url = await self._layer_url_for(dataset_id, layer) + fields, _ = await self._layer_fields(layer_url) + return [ + { + "name": f.get("name", ""), + "type": f.get("type", ""), + "alias": f.get("alias", ""), + } + for f in fields + ] + + async def _layer_url_for( + self, dataset_id: str, layer: int | None, *, queryable_only: bool = False + ) -> str: + """Resolve a Hub item to a trusted layer URL (two-hop). + + Args: + dataset_id: Hub item ID. + layer: Layer or table id, or ``None`` for the default layer. + queryable_only: Refuse item types outside :attr:`QUERYABLE_TYPES`. + + Returns: + The validated layer URL (``.../FeatureServer/``). + + Raises: + ValueError: If the item has no service URL, is not a queryable + type, its host is not trusted, or ``layer`` is invalid. + """ dataset = await self.get_dataset(dataset_id) service_url = dataset.get("service_url") if not service_url: raise ValueError( f"Dataset {dataset_id} does not have a queryable Feature Service URL" ) - + ds_type = dataset.get("type", "") + if queryable_only and ds_type and ds_type not in self.QUERYABLE_TYPES: + raise ValueError( + f"Dataset type '{ds_type}' is not queryable. " + f"query_data only supports: {', '.join(sorted(self.QUERYABLE_TYPES))}." + ) service_url = self._validate_feature_url( service_url, self.plugin_config.portal_url, self.plugin_config.trusted_service_hosts, auto_trust=self.plugin_config.auto_trust_hub_services, ) - layer_url = await self._resolve_layer_url(service_url, layer) + return await self._resolve_layer_url(service_url, layer) + + async def _layer_fields( + self, layer_url: str + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Return a layer's ``fields`` and its raw metadata. + Reads ``{layer_url}?f=json``; when that body is unusable, the fields + come from a one-row query and the metadata is ``{}``. + """ # The format must travel as a query parameter: httpx replaces the # URL's own query string with ``params``, so ``{url}?f=json`` plus # ``params={}`` sent a bare request and ArcGIS answered with an HTML @@ -567,25 +738,18 @@ async def get_schema( response = await self._call_feature_service(layer_url, {"f": "json"}) fields = self._fields_from_layer_metadata(response) - if fields is None: - # Some self-hosted services return HTML or a malformed body on - # the layer metadata endpoint while the query endpoint works - # (seen on Columbus and Indianapolis). Derive the schema from a - # one-row query instead of failing. - logger.warning( - "Layer metadata at %s was not usable; deriving schema from a " - "1-row query", - layer_url, - ) - fields = await self._schema_from_query(layer_url) - return [ - { - "name": f.get("name", ""), - "type": f.get("type", ""), - "alias": f.get("alias", ""), - } - for f in fields - ] + if fields is not None: + metadata = response.json() + return fields, metadata if isinstance(metadata, dict) else {} + # Some self-hosted services return HTML or a malformed body on the + # layer metadata endpoint while the query endpoint works (seen on + # Columbus and Indianapolis). Derive the schema from a one-row query + # instead of failing. + logger.warning( + "Layer metadata at %s was not usable; deriving schema from a 1-row query", + layer_url, + ) + return await self._schema_from_query(layer_url), {} @staticmethod def _fields_from_layer_metadata(response: httpx.Response) -> list | None: @@ -745,30 +909,10 @@ async def _query_page( if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0: raise ValueError(f"offset must be a non-negative integer (got {offset})") sort = WhereValidator.validate_order_by(order_by) - - dataset = await self.get_dataset(dataset_id) - service_url = dataset.get("service_url") - ds_type = dataset.get("type", "") - if not service_url: - raise ValueError( - f"Dataset {dataset_id} does not have a queryable Feature Service URL" - ) - - if ds_type and ds_type not in self.QUERYABLE_TYPES: - raise ValueError( - f"Dataset type '{ds_type}' is not queryable. " - f"query_data only supports: {', '.join(sorted(self.QUERYABLE_TYPES))}." - ) - where_clause = WhereValidator.validate(where) - service_url = self._validate_feature_url( - service_url, - self.plugin_config.portal_url, - self.plugin_config.trusted_service_hosts, - auto_trust=self.plugin_config.auto_trust_hub_services, - ) - service_url = await self._resolve_layer_url(service_url, layer) - query_url = f"{service_url}/query" + + layer_url = await self._layer_url_for(dataset_id, layer, queryable_only=True) + query_url = f"{layer_url}/query" record_count = min(limit, 1000) params: dict[str, Any] = { "where": where_clause, @@ -848,6 +992,188 @@ async def _count_features(self, query_url: str, where: str) -> int | None: return None return count + async def aggregate_features( + self, + dataset_id: str, + statistics: list[dict[str, Any]], + *, + group_by: list[str] | None = None, + where: str = "1=1", + having: str | None = None, + order_by: str | None = None, + limit: int = 1000, + layer: int | None = None, + ) -> dict[str, Any]: + """Run a server-side statistics query (``outStatistics``). + + Every field named in ``statistics`` and ``group_by`` must exist in the + layer's schema (matched case-insensitively and sent with the schema's + spelling). ``where`` and ``having`` go through + :class:`WhereValidator`, ``order_by`` through + :meth:`WhereValidator.validate_order_by`. A ``count`` without a field + counts records using the layer's object-id field. + + Args: + dataset_id: Hub item ID. + statistics: ``[{"type", "field", "as"}]``; ``type`` is one of + :data:`STATISTIC_TYPES`, ``as`` names the output column + (default ``_``). + group_by: Fields to group by. + where: WHERE clause applied before grouping. + having: Filter on the computed statistics. + order_by: Sort order over group fields or output names. + limit: Maximum number of groups returned. + layer: Layer or table id, or ``None`` for the default layer. + + Returns: + ``{"rows", "total_groups", "layer_url", "exceeded"}``. + + Raises: + ValueError: On invalid arguments, unknown fields, or a layer that + does not support statistics. + RuntimeError: If the service reports an error. + """ + if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1: + raise ValueError(f"limit must be at least 1 (got {limit})") + if not isinstance(statistics, list) or not statistics: + raise ValueError("statistics must be a non-empty list") + if len(statistics) > _MAX_STATISTICS: + raise ValueError(f"statistics accepts at most {_MAX_STATISTICS} entries") + group_by = group_by or [] + if not isinstance(group_by, list) or not all( + isinstance(g, str) for g in group_by + ): + raise ValueError("group_by must be a list of field names") + if len(group_by) > _MAX_GROUP_BY: + raise ValueError(f"group_by accepts at most {_MAX_GROUP_BY} fields") + for name in group_by: + self._check_field_name(name, "group_by field") + requested = [self._parse_statistic(stat) for stat in statistics] + where_clause = WhereValidator.validate(where) + having_clause = ( + WhereValidator.validate(having) + if isinstance(having, str) and having.strip() + else None + ) + if having is not None and not isinstance(having, str): + raise ValueError("having must be a string") + sort = WhereValidator.validate_order_by(order_by) + + layer_url = await self._layer_url_for(dataset_id, layer, queryable_only=True) + fields, metadata = await self._layer_fields(layer_url) + capabilities = metadata.get("advancedQueryCapabilities") or {} + if capabilities.get("supportsStatistics") is False: + raise ValueError( + "This layer does not support server-side statistics; use " + "query_data and aggregate the rows instead." + ) + + by_name = { + str(f.get("name", "")).lower(): f for f in fields if isinstance(f, dict) + } + object_id = next( + (f.get("name") for f in fields if f.get("type") == "esriFieldTypeOID"), + None, + ) + + def schema_name(name: str, role: str) -> str: + field = by_name.get(name.lower()) + if field is None: + raise ValueError( + f"Unknown {role} {name!r}; use get_schema to list this " + "layer's fields" + ) + return str(field["name"]) + + groups = [schema_name(g, "group_by field") for g in group_by] + out_statistics = [] + out_names: set[str] = {g.lower() for g in groups} + for stat_type, field, alias in requested: + if field is None: + if stat_type != "count": + raise ValueError(f"{stat_type} needs a field") + if not object_id: + raise ValueError( + "count without a field needs an object-id field; name " + "a field to count" + ) + field = object_id + else: + field = schema_name(field, "statistics field") + alias = alias or f"{stat_type}_{field}" + if alias.lower() in out_names: + raise ValueError( + f"Output name {alias!r} is used twice; give each statistic " + "a distinct 'as'" + ) + out_names.add(alias.lower()) + out_statistics.append( + { + "statisticType": stat_type, + "onStatisticField": field, + "outStatisticFieldName": alias, + } + ) + + params: dict[str, Any] = { + "where": where_clause, + "outStatistics": json.dumps(out_statistics), + "returnGeometry": "false", + "f": "json", + } + if groups: + params["groupByFieldsForStatistics"] = ",".join(groups) + if having_clause: + params["havingClause"] = having_clause + if sort: + params["orderByFields"] = sort + + data = self._query_json( + await self._call_feature_service(f"{layer_url}/query", params) + ) + rows = [f.get("attributes", {}) for f in data.get("features") or []] + return { + "rows": rows[:limit], + "total_groups": len(rows), + "layer_url": layer_url, + "exceeded": bool(data.get("exceededTransferLimit")), + } + + @staticmethod + def _check_field_name(name: Any, role: str) -> str: + """Require a plain field name or output alias.""" + if not isinstance(name, str) or not _FIELD_NAME.match(name): + raise ValueError( + f"Invalid {role} {str(name)[:80]!r}: use a plain field name " + "(letters, digits, underscores)" + ) + return name + + def _parse_statistic(self, stat: Any) -> tuple[str, str | None, str | None]: + """Validate one ``statistics`` entry; return ``(type, field, as)``.""" + if not isinstance(stat, dict): + raise ValueError( + 'Each statistic must be an object like {"type": "sum", ' + '"field": "UNITS"}' + ) + stat_type = str(stat.get("type", "")).lower() + if stat_type not in STATISTIC_TYPES: + raise ValueError( + f"Unknown statistic type {str(stat.get('type'))[:40]!r}; use one " + f"of {', '.join(STATISTIC_TYPES)}" + ) + field = stat.get("field") + if field in (None, "", "*"): + field = None + else: + self._check_field_name(field, "statistics field") + alias = stat.get("as") + if alias in (None, ""): + alias = None + else: + self._check_field_name(alias, "statistic name") + return stat_type, field, alias + # ── Aggregations (standalone helper, not a DataPlugin method) ─────── async def get_aggregations( @@ -1399,6 +1725,25 @@ def _format_query_results( ) return text + def _format_aggregate(self, result: dict[str, Any], *, fmt: str = "text") -> str: + """Format an :meth:`aggregate_features` result.""" + rows = result.get("rows") or [] + if not rows: + return "No groups matched." + total = result.get("total_groups", len(rows)) + header = ( + f"Aggregated {total} group(s):" + if total == len(rows) + else f"Aggregated {total} group(s); showing the first {len(rows)} (limit):" + ) + text, _ = self.render_rows(rows, fmt, header=header, skip_keys=frozenset()) + if result.get("exceeded"): + text += ( + "\n\nThe service hit its transfer limit, so some groups are " + "missing. Narrow where or group by fewer fields." + ) + return text + def _format_aggregations(self, field: str, buckets: list[dict[str, Any]]) -> str: if not buckets: return f"No aggregation results for '{field}'." diff --git a/tests/security/test_arcgis_aggregate_guards.py b/tests/security/test_arcgis_aggregate_guards.py new file mode 100644 index 0000000..dc65fda --- /dev/null +++ b/tests/security/test_arcgis_aggregate_guards.py @@ -0,0 +1,81 @@ +"""Security tests for ArcGIS aggregate_data inputs (outStatistics, having).""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from plugins.arcgis.plugin import ArcGISPlugin + +pytestmark = pytest.mark.security + +INJECTIONS = [ + "x; DROP TABLE t", + "x) OR (1=1", + "x -- c", + "UPPER(x)", + "'x'", + "a" * 100, + "1x", +] + + +@pytest.fixture +def plugin(): + p = ArcGISPlugin( + {"portal_url": "https://hub.arcgis.com", "city_name": "T", "timeout": 30} + ) + p.feature_client = AsyncMock() + return p + + +async def _run(plugin, arguments): + with patch.object(plugin, "get_dataset", new_callable=AsyncMock) as get_ds: + result = await plugin.execute_tool( + "aggregate_data", {"dataset_id": "abc", **arguments} + ) + return result, get_ds + + +class TestAggregateInputsBlockedBeforeRequest: + @pytest.mark.parametrize("bad", INJECTIONS) + async def test_statistic_field(self, plugin, bad): + result, get_ds = await _run( + plugin, {"statistics": [{"type": "sum", "field": bad}]} + ) + assert result.success is False + get_ds.assert_not_called() + plugin.feature_client.get.assert_not_called() + + @pytest.mark.parametrize("bad", INJECTIONS) + async def test_statistic_alias(self, plugin, bad): + result, get_ds = await _run( + plugin, {"statistics": [{"type": "count", "field": "A", "as": bad}]} + ) + assert result.success is False + get_ds.assert_not_called() + + @pytest.mark.parametrize("bad", INJECTIONS) + async def test_group_by_field(self, plugin, bad): + result, get_ds = await _run( + plugin, {"statistics": [{"type": "count"}], "group_by": [bad]} + ) + assert result.success is False + get_ds.assert_not_called() + + @pytest.mark.parametrize( + "bad", + ["COUNT(*) > 1; DROP TABLE t", "COUNT(*) > 0 OR DELETE FROM t"], + ) + async def test_having_clause(self, plugin, bad): + result, get_ds = await _run( + plugin, {"statistics": [{"type": "count"}], "having": bad} + ) + assert result.success is False + get_ds.assert_not_called() + + async def test_order_by(self, plugin): + result, get_ds = await _run( + plugin, {"statistics": [{"type": "count"}], "order_by": "x; DROP"} + ) + assert result.success is False + get_ds.assert_not_called() diff --git a/tests/unit/plugins/arcgis/test_aggregate_data.py b/tests/unit/plugins/arcgis/test_aggregate_data.py new file mode 100644 index 0000000..731506f --- /dev/null +++ b/tests/unit/plugins/arcgis/test_aggregate_data.py @@ -0,0 +1,246 @@ +"""Tests for the ArcGIS aggregate_data tool (#31). + +Modelled on HUD's Housing Choice Vouchers by Tract layer: sum HCV_PUBLIC and +count tracts, grouped by state and county. +""" + +import json +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from plugins.arcgis.plugin import ArcGISPlugin + +LAYER = "https://services.arcgis.com/VTyQ9soqVukalItT/arcgis/rest/services/HCV/FeatureServer/0" +FIELDS = [ + {"name": "OBJECTID", "type": "esriFieldTypeOID"}, + {"name": "GEOID", "type": "esriFieldTypeString"}, + {"name": "STATE", "type": "esriFieldTypeString"}, + {"name": "COUNTY", "type": "esriFieldTypeString"}, + {"name": "HCV_PUBLIC", "type": "esriFieldTypeInteger"}, +] +METADATA = { + "fields": FIELDS, + "advancedQueryCapabilities": {"supportsStatistics": True}, +} +GROUPS = { + "features": [ + { + "attributes": { + "STATE": "18", + "COUNTY": "141", + "vouchers": 2609, + "tracts": 82, + } + }, + {"attributes": {"STATE": "18", "COUNTY": "039", "vouchers": 767, "tracts": 45}}, + ] +} +FIVE = "STATE||COUNTY IN ('18141','18039','18091','18099','26021')" + + +def _response(data): + mock = Mock() + mock.json.return_value = data + mock.raise_for_status = Mock() + mock.headers = Mock() + mock.headers.get = Mock(return_value="application/json") + return mock + + +@pytest.fixture +def plugin(): + p = ArcGISPlugin( + {"portal_url": "https://hub.arcgis.com", "city_name": "HUD", "timeout": 30} + ) + p.feature_client = AsyncMock() + return p + + +async def _aggregate(plugin, arguments, *payloads): + plugin.feature_client.get = AsyncMock(side_effect=[_response(p) for p in payloads]) + dataset = {"id": "abc", "type": "Feature Layer", "service_url": LAYER} + with patch.object( + plugin, "get_dataset", new_callable=AsyncMock, return_value=dataset + ): + return await plugin.execute_tool( + "aggregate_data", {"dataset_id": "abc", **arguments} + ) + + +def _query_params(plugin): + call = plugin.feature_client.get.call_args_list[-1] + assert call.args[0] == f"{LAYER}/query" + return call.kwargs["params"] + + +VOUCHERS = { + "statistics": [ + {"type": "sum", "field": "HCV_PUBLIC", "as": "vouchers"}, + {"type": "count", "field": "GEOID", "as": "tracts"}, + ], + "group_by": ["STATE", "COUNTY"], + "where": FIVE, +} + + +class TestRequest: + async def test_builds_out_statistics_and_group_by(self, plugin): + result = await _aggregate(plugin, VOUCHERS, METADATA, GROUPS) + assert result.success is True + params = _query_params(plugin) + assert json.loads(params["outStatistics"]) == [ + { + "statisticType": "sum", + "onStatisticField": "HCV_PUBLIC", + "outStatisticFieldName": "vouchers", + }, + { + "statisticType": "count", + "onStatisticField": "GEOID", + "outStatisticFieldName": "tracts", + }, + ] + assert params["groupByFieldsForStatistics"] == "STATE,COUNTY" + assert params["where"] == FIVE + assert "havingClause" not in params + assert "orderByFields" not in params + + async def test_fields_match_schema_case_insensitively(self, plugin): + args = { + "statistics": [{"type": "SUM", "field": "hcv_public"}], + "group_by": ["county"], + } + await _aggregate(plugin, args, METADATA, GROUPS) + params = _query_params(plugin) + stat = json.loads(params["outStatistics"])[0] + assert stat["onStatisticField"] == "HCV_PUBLIC" + assert stat["outStatisticFieldName"] == "sum_HCV_PUBLIC" + assert params["groupByFieldsForStatistics"] == "COUNTY" + + async def test_count_without_field_uses_object_id(self, plugin): + await _aggregate(plugin, {"statistics": [{"type": "count"}]}, METADATA, GROUPS) + stat = json.loads(_query_params(plugin)["outStatistics"])[0] + assert stat["onStatisticField"] == "OBJECTID" + + async def test_having_and_order_by_are_sent(self, plugin): + args = { + **VOUCHERS, + "having": "SUM(HCV_PUBLIC) > 100", + "order_by": "vouchers desc", + } + await _aggregate(plugin, args, METADATA, GROUPS) + params = _query_params(plugin) + assert params["havingClause"] == "SUM(HCV_PUBLIC) > 100" + assert params["orderByFields"] == "vouchers DESC" + + async def test_layer_argument_selects_a_table(self, plugin): + service = LAYER.rsplit("/", 1)[0] + info = {"layers": [{"id": 0, "name": "Tracts"}], "tables": [{"id": 1}]} + plugin.feature_client.get = AsyncMock( + side_effect=[_response(info), _response(METADATA), _response(GROUPS)] + ) + dataset = {"id": "abc", "type": "Feature Service", "service_url": service} + with patch.object( + plugin, "get_dataset", new_callable=AsyncMock, return_value=dataset + ): + result = await plugin.execute_tool( + "aggregate_data", + {"dataset_id": "abc", "layer": 1, "statistics": [{"type": "count"}]}, + ) + assert result.success is True + urls = [c.args[0] for c in plugin.feature_client.get.call_args_list] + assert urls == [service, f"{service}/1", f"{service}/1/query"] + + +class TestOutput: + async def test_text_output_lists_each_group(self, plugin): + result = await _aggregate(plugin, VOUCHERS, METADATA, GROUPS) + text = result.content[0]["text"] + assert "Aggregated 2 group(s):" in text + assert "vouchers: 2609" in text + assert "tracts: 45" in text + assert "Statistics ignore null values" in text + + async def test_csv_output(self, plugin): + result = await _aggregate( + plugin, {**VOUCHERS, "format": "csv"}, METADATA, GROUPS + ) + text = result.content[0]["text"] + assert "STATE,COUNTY,vouchers,tracts\n18,141,2609,82\n18,039,767,45" in text + + async def test_limit_trims_groups_with_notice(self, plugin): + result = await _aggregate(plugin, {**VOUCHERS, "limit": 1}, METADATA, GROUPS) + text = result.content[0]["text"] + assert "Aggregated 2 group(s); showing the first 1 (limit):" in text + assert "COUNTY: 039" not in text + + async def test_transfer_limit_is_reported(self, plugin): + payload = {**GROUPS, "exceededTransferLimit": True} + result = await _aggregate(plugin, VOUCHERS, METADATA, payload) + assert "hit its transfer limit" in result.content[0]["text"] + + async def test_no_groups(self, plugin): + result = await _aggregate(plugin, VOUCHERS, METADATA, {"features": []}) + assert result.success is True + assert "No groups matched." in result.content[0]["text"] + + +class TestRefusals: + async def test_unknown_field_names_the_problem(self, plugin): + args = {"statistics": [{"type": "sum", "field": "NOPE"}]} + result = await _aggregate(plugin, args, METADATA) + assert result.success is False + assert "Unknown statistics field 'NOPE'" in result.error_message + + async def test_unknown_group_by_field(self, plugin): + args = {"statistics": [{"type": "count"}], "group_by": ["TRACT"]} + result = await _aggregate(plugin, args, METADATA) + assert "Unknown group_by field 'TRACT'" in result.error_message + + async def test_layer_without_statistics_support(self, plugin): + metadata = { + **METADATA, + "advancedQueryCapabilities": {"supportsStatistics": False}, + } + result = await _aggregate(plugin, {"statistics": [{"type": "count"}]}, metadata) + assert result.success is False + assert "does not support server-side statistics" in result.error_message + + async def test_duplicate_output_names(self, plugin): + args = { + "statistics": [ + {"type": "sum", "field": "HCV_PUBLIC", "as": "n"}, + {"type": "count", "field": "GEOID", "as": "N"}, + ] + } + result = await _aggregate(plugin, args, METADATA) + assert "used twice" in result.error_message + + async def test_alias_clashing_with_group_field(self, plugin): + args = { + "statistics": [{"type": "count", "as": "county"}], + "group_by": ["COUNTY"], + } + result = await _aggregate(plugin, args, METADATA) + assert "used twice" in result.error_message + + async def test_sum_needs_a_field(self, plugin): + result = await _aggregate(plugin, {"statistics": [{"type": "sum"}]}, METADATA) + assert "sum needs a field" in result.error_message + + @pytest.mark.parametrize( + "statistics", + [[], "sum", [{"type": "median", "field": "X"}], ["sum"], [{"field": "X"}]], + ) + async def test_malformed_statistics_never_reach_the_service( + self, plugin, statistics + ): + result = await _aggregate(plugin, {"statistics": statistics}) + assert result.success is False + plugin.feature_client.get.assert_not_called() + + @pytest.mark.parametrize("fmt", ["xml", "TEXT"]) + async def test_unknown_format(self, plugin, fmt): + result = await _aggregate(plugin, {**VOUCHERS, "format": fmt}) + assert "format must be one of" in result.error_message diff --git a/tests/unit/plugins/arcgis/test_arcgis_plugin.py b/tests/unit/plugins/arcgis/test_arcgis_plugin.py index b9e7534..9ad229c 100644 --- a/tests/unit/plugins/arcgis/test_arcgis_plugin.py +++ b/tests/unit/plugins/arcgis/test_arcgis_plugin.py @@ -159,17 +159,18 @@ async def test_shutdown_closes_tracked_clients(self, arcgis_config): class TestGetTools: - def test_get_tools_returns_five_tools(self, arcgis_config): + def test_get_tools_returns_six_tools(self, arcgis_config): plugin = ArcGISPlugin(arcgis_config) tools = plugin.get_tools() - assert len(tools) == 5 + assert len(tools) == 6 tool_names = [t.name for t in tools] assert "search_datasets" in tool_names assert "get_dataset" in tool_names assert "get_aggregations" in tool_names assert "get_schema" in tool_names assert "query_data" in tool_names + assert "aggregate_data" in tool_names def test_get_tools_uses_city_name_directly(self, arcgis_config): plugin = ArcGISPlugin(arcgis_config)