Skip to content

Commit 3bb6e3a

Browse files
sodresodre
authored
✨ feat(schema): add GSI3 sparse index for entity config queries (#287)
## Summary - Add GSI3 to efficiently query entities with custom limit configurations - Previously required full table Scan (O(n) cost), now O(k) where k = entities with custom configs - GSI3 uses sparse index pattern: only entity-level config records have GSI3 attributes - Add `list_entities_with_custom_limits()` API method with pagination support - Add CLI command: `entity list --with-custom-limits <resource>` ## Test plan - [x] Unit tests for schema key builders (GSI3PK, GSI3SK) - [x] Unit tests for repository `list_entities_with_custom_limits()` - [x] Unit tests for `RateLimiter` and `SyncRateLimiter` methods - [x] Unit tests for CLI command with pagination - [x] ValidationError coverage for CLI command - [x] Integration test with LocalStack Closes #235 🤖 Generated with [Claude Code](https://claude.ai/code)
2 parents 203395a + d25b342 commit 3bb6e3a

11 files changed

Lines changed: 598 additions & 2 deletions

File tree

‎CLAUDE.md‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,7 @@ Primary mitigation: cascade defaults to `False`.
312312
- All entities, buckets, limits, usage in one table
313313
- GSI1: Parent -> Children lookups
314314
- GSI2: Resource aggregation (capacity tracking)
315+
- GSI3: Entity config queries (sparse - only entity configs indexed)
315316
- Uses TransactWriteItems for atomicity
316317

317318
### Exception Design
@@ -418,6 +419,7 @@ docs/
418419
| Get system config (limits + on_unavailable) | `PK=SYSTEM#, SK=#CONFIG` |
419420
| Get resource config (limits) | `PK=RESOURCE#{resource}, SK=#CONFIG` |
420421
| Get entity config (limits) | `PK=ENTITY#{id}, SK=#CONFIG#{resource}` |
422+
| List entities with custom limits | GSI3: `GSI3PK=ENTITY_CONFIG#{resource}` |
421423

422424
**Optimized read patterns (issue #133):**
423425
- `acquire()` uses `BatchGetItem` to fetch all buckets for entity + parent in a single round trip
@@ -446,7 +448,7 @@ Limit configs use a three-level hierarchy with precedence: **Entity > Resource >
446448
|-------|-----|-----|--------|------|
447449
| System | `set_system_defaults(limits, on_unavailable)` | `get_system_defaults()` | `delete_system_defaults()` | - |
448450
| Resource | `set_resource_defaults(resource, limits)` | `get_resource_defaults(resource)` | `delete_resource_defaults(resource)` | `list_resources_with_defaults()` |
449-
| Entity | `set_limits(entity_id, limits, resource)` | `get_limits(entity_id, resource)` | `delete_limits(entity_id, resource)` | - |
451+
| Entity | `set_limits(entity_id, limits, resource)` | `get_limits(entity_id, resource)` | `delete_limits(entity_id, resource)` | `list_entities_with_custom_limits(resource)` |
450452

451453
**CLI commands for managing stored limits:**
452454

@@ -461,7 +463,7 @@ zae-limiter resource set-defaults gpt-4 -l tpm:50000 -l rpm:500
461463
zae-limiter entity set-limits user-123 --resource gpt-4 -l rpm:1000
462464
```
463465

464-
Each level also has `get-*` and `delete-*` subcommands. Use `zae-limiter resource list` to list resources with defaults.
466+
Each level also has `get-*` and `delete-*` subcommands. Use `zae-limiter resource list` to list resources with defaults. Use `zae-limiter entity list --with-custom-limits <resource>` to list entities with custom limits for a specific resource.
465467

466468
Limit configs use composite items (v0.8.0+, ADR-114 for configs). All limits for a config level are stored in a single item:
467469

‎src/zae_limiter/cli.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3197,6 +3197,104 @@ async def _delete() -> None:
31973197
asyncio.run(_delete())
31983198

31993199

3200+
@entity.command(
3201+
"list",
3202+
epilog="""\b
3203+
Examples:
3204+
\b
3205+
# List entities with custom limits for gpt-4
3206+
zae-limiter entity list --with-custom-limits gpt-4
3207+
\b
3208+
# List with a maximum of 10 results
3209+
zae-limiter entity list --with-custom-limits claude-3 --limit 10
3210+
""",
3211+
)
3212+
@click.option(
3213+
"--with-custom-limits",
3214+
"resource",
3215+
required=True,
3216+
help="List entities with custom limits for this resource",
3217+
)
3218+
@click.option(
3219+
"--limit",
3220+
type=int,
3221+
default=None,
3222+
help="Maximum number of entities to return",
3223+
)
3224+
@click.option(
3225+
"--name",
3226+
"-n",
3227+
default="limiter",
3228+
help="Stack identifier used as the CloudFormation stack name. Default: limiter",
3229+
)
3230+
@click.option(
3231+
"--region",
3232+
help="AWS region (default: use boto3 defaults)",
3233+
)
3234+
@click.option(
3235+
"--endpoint-url",
3236+
help="AWS endpoint URL (e.g., http://localhost:4566 for LocalStack)",
3237+
)
3238+
def entity_list(
3239+
resource: str,
3240+
limit: int | None,
3241+
name: str,
3242+
region: str | None,
3243+
endpoint_url: str | None,
3244+
) -> None:
3245+
"""List entities with custom limit configurations.
3246+
3247+
Uses GSI3 sparse index for efficient queries.
3248+
3249+
\f
3250+
3251+
**Examples:**
3252+
```bash
3253+
# List entities with custom limits for gpt-4
3254+
zae-limiter entity list --with-custom-limits gpt-4
3255+
3256+
# List with a maximum of 10 results
3257+
zae-limiter entity list --with-custom-limits claude-3 --limit 10
3258+
```
3259+
"""
3260+
from .exceptions import ValidationError
3261+
from .repository import Repository
3262+
3263+
async def _list() -> None:
3264+
try:
3265+
repo = Repository(name, region, endpoint_url)
3266+
except ValidationError as e:
3267+
click.echo(f"Error: {e.reason}", err=True)
3268+
sys.exit(1)
3269+
3270+
try:
3271+
cursor: str | None = None
3272+
total_count = 0
3273+
while True:
3274+
entities, cursor = await repo.list_entities_with_custom_limits(
3275+
resource=resource, limit=limit, cursor=cursor
3276+
)
3277+
for entity_id in entities:
3278+
click.echo(entity_id)
3279+
total_count += 1
3280+
3281+
if cursor is None:
3282+
break
3283+
3284+
if total_count == 0:
3285+
click.echo(f"No entities with custom limits for resource '{resource}'")
3286+
except ValidationError as e:
3287+
click.echo(f"Error: {e.reason}", err=True)
3288+
sys.exit(1)
3289+
except Exception as e:
3290+
click.echo(f"Error: Failed to list entities: {e}", err=True)
3291+
sys.exit(1)
3292+
finally:
3293+
await repo.close()
3294+
3295+
asyncio.run(_list())
3296+
3297+
32003298
# ---------------------------------------------------------------------------
32013299
# Local development commands
32023300
# ---------------------------------------------------------------------------

‎src/zae_limiter/infra/cfn_template.yaml‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,10 @@ Resources:
295295
AttributeType: S
296296
- AttributeName: GSI2SK
297297
AttributeType: S
298+
- AttributeName: GSI3PK
299+
AttributeType: S
300+
- AttributeName: GSI3SK
301+
AttributeType: S
298302

299303
KeySchema:
300304
- AttributeName: PK
@@ -321,6 +325,15 @@ Resources:
321325
Projection:
322326
ProjectionType: ALL
323327

328+
- IndexName: GSI3
329+
KeySchema:
330+
- AttributeName: GSI3PK
331+
KeyType: HASH
332+
- AttributeName: GSI3SK
333+
KeyType: RANGE
334+
Projection:
335+
ProjectionType: KEYS_ONLY
336+
324337
TimeToLiveSpecification:
325338
AttributeName: ttl
326339
Enabled: true

‎src/zae_limiter/limiter.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1179,6 +1179,42 @@ async def delete_limits(
11791179
await self._ensure_initialized()
11801180
await self._repository.delete_limits(entity_id, resource, principal=principal)
11811181

1182+
async def list_entities_with_custom_limits(
1183+
self,
1184+
resource: str,
1185+
limit: int | None = None,
1186+
cursor: str | None = None,
1187+
) -> tuple[list[str], str | None]:
1188+
"""
1189+
List all entities that have custom limit configurations.
1190+
1191+
Uses GSI3 sparse index for efficient queries. Only entities with
1192+
custom limits for the specified resource are returned.
1193+
1194+
Args:
1195+
resource: Resource to filter by.
1196+
limit: Maximum number of entities to return. None for all.
1197+
cursor: Pagination cursor from previous call.
1198+
1199+
Returns:
1200+
Tuple of (entity_ids, next_cursor). next_cursor is None if no more results.
1201+
1202+
Example:
1203+
# Get all entities with custom limits for gpt-4
1204+
entities, cursor = await limiter.list_entities_with_custom_limits("gpt-4")
1205+
for entity_id in entities:
1206+
print(entity_id)
1207+
1208+
# Paginate through results
1209+
while cursor:
1210+
more, cursor = await limiter.list_entities_with_custom_limits(
1211+
"gpt-4", cursor=cursor
1212+
)
1213+
entities.extend(more)
1214+
"""
1215+
await self._ensure_initialized()
1216+
return await self._repository.list_entities_with_custom_limits(resource, limit, cursor)
1217+
11821218
# -------------------------------------------------------------------------
11831219
# Resource-level defaults management
11841220
# -------------------------------------------------------------------------
@@ -1951,6 +1987,15 @@ def delete_limits(
19511987
"""Delete stored limit configs for an entity."""
19521988
self._run(self._limiter.delete_limits(entity_id, resource, principal=principal))
19531989

1990+
def list_entities_with_custom_limits(
1991+
self,
1992+
resource: str,
1993+
limit: int | None = None,
1994+
cursor: str | None = None,
1995+
) -> tuple[list[str], str | None]:
1996+
"""List all entities that have custom limit configurations."""
1997+
return self._run(self._limiter.list_entities_with_custom_limits(resource, limit, cursor))
1998+
19541999
# -------------------------------------------------------------------------
19552000
# Resource-level defaults management
19562001
# -------------------------------------------------------------------------

‎src/zae_limiter/repository.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -959,6 +959,9 @@ async def set_limits(
959959
"entity_id": {"S": entity_id},
960960
"resource": {"S": resource},
961961
"config_version": {"N": "1"},
962+
# GSI3 attributes for sparse indexing (entity config queries)
963+
"GSI3PK": {"S": schema.gsi3_pk_entity_config(resource)},
964+
"GSI3SK": {"S": schema.gsi3_sk_entity(entity_id)},
962965
}
963966

964967
# Add l_* attributes for each limit
@@ -1039,6 +1042,62 @@ async def delete_limits(
10391042
resource=resource,
10401043
)
10411044

1045+
async def list_entities_with_custom_limits(
1046+
self,
1047+
resource: str,
1048+
limit: int | None = None,
1049+
cursor: str | None = None,
1050+
) -> tuple[list[str], str | None]:
1051+
"""
1052+
List all entities that have custom limit configurations for a resource.
1053+
1054+
Uses GSI3 sparse index for efficient queries. Only entity-level configs
1055+
have GSI3 attributes, so this query returns only entities with custom
1056+
limits (not system or resource defaults).
1057+
1058+
Args:
1059+
resource: Resource to filter by (required).
1060+
limit: Maximum number of entities to return. None for all.
1061+
cursor: Pagination cursor from previous call. None for first page.
1062+
1063+
Returns:
1064+
Tuple of (entity_ids, next_cursor). next_cursor is None if no more results.
1065+
"""
1066+
import base64
1067+
import json
1068+
1069+
client = await self._get_client()
1070+
1071+
query_params: dict[str, Any] = {
1072+
"TableName": self.table_name,
1073+
"IndexName": schema.GSI3_NAME,
1074+
"KeyConditionExpression": "GSI3PK = :pk",
1075+
"ExpressionAttributeValues": {":pk": {"S": schema.gsi3_pk_entity_config(resource)}},
1076+
}
1077+
1078+
if limit is not None:
1079+
query_params["Limit"] = limit
1080+
if cursor is not None:
1081+
# Decode cursor (base64 encoded LastEvaluatedKey)
1082+
query_params["ExclusiveStartKey"] = json.loads(base64.b64decode(cursor))
1083+
1084+
response = await client.query(**query_params)
1085+
1086+
entity_ids: list[str] = []
1087+
for item in response.get("Items", []):
1088+
entity_id = item.get("GSI3SK", {}).get("S")
1089+
if entity_id:
1090+
entity_ids.append(entity_id)
1091+
1092+
# Encode next cursor if more results
1093+
next_cursor: str | None = None
1094+
if "LastEvaluatedKey" in response:
1095+
next_cursor = base64.b64encode(
1096+
json.dumps(response["LastEvaluatedKey"]).encode()
1097+
).decode()
1098+
1099+
return entity_ids, next_cursor
1100+
10421101
# -------------------------------------------------------------------------
10431102
# Resource-level limit config operations (composite format, ADR-114)
10441103
# -------------------------------------------------------------------------

‎src/zae_limiter/repository_protocol.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,27 @@ async def delete_limits(
374374
"""
375375
...
376376

377+
async def list_entities_with_custom_limits(
378+
self,
379+
resource: str,
380+
limit: int | None = None,
381+
cursor: str | None = None,
382+
) -> tuple[list[str], str | None]:
383+
"""
384+
List all entities that have custom limit configurations for a resource.
385+
386+
Uses GSI3 sparse index for efficient queries.
387+
388+
Args:
389+
resource: Resource to filter by (required).
390+
limit: Maximum number of entities to return. None for all.
391+
cursor: Pagination cursor from previous call. None for first page.
392+
393+
Returns:
394+
Tuple of (entity_ids, next_cursor). next_cursor is None if no more results.
395+
"""
396+
...
397+
377398
# -------------------------------------------------------------------------
378399
# Resource-level defaults
379400
# -------------------------------------------------------------------------

‎src/zae_limiter/schema.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,15 @@
66
DEFAULT_TABLE_NAME = "rate_limits"
77
GSI1_NAME = "GSI1" # For parent -> children lookups
88
GSI2_NAME = "GSI2" # For resource aggregation
9+
GSI3_NAME = "GSI3" # For entity config queries (sparse)
910

1011
# Key prefixes
1112
ENTITY_PREFIX = "ENTITY#"
1213
PARENT_PREFIX = "PARENT#"
1314
CHILD_PREFIX = "CHILD#"
1415
RESOURCE_PREFIX = "RESOURCE#"
1516
SYSTEM_PREFIX = "SYSTEM#"
17+
ENTITY_CONFIG_PREFIX = "ENTITY_CONFIG#" # For GSI3 sparse index
1618

1719
# Sort key prefixes
1820
SK_META = "#META"
@@ -208,6 +210,16 @@ def gsi2_sk_usage(window_key: str, entity_id: str) -> str:
208210
return f"USAGE#{window_key}#{entity_id}"
209211

210212

213+
def gsi3_pk_entity_config(resource: str) -> str:
214+
"""Build GSI3 partition key for entity config lookup by resource."""
215+
return f"{ENTITY_CONFIG_PREFIX}{resource}"
216+
217+
218+
def gsi3_sk_entity(entity_id: str) -> str:
219+
"""Build GSI3 sort key for entity config (just entity_id)."""
220+
return entity_id
221+
222+
211223
def pk_audit(entity_id: str) -> str:
212224
"""Build partition key for audit log records."""
213225
return f"{AUDIT_PREFIX}{entity_id}"
@@ -245,6 +257,8 @@ def get_table_definition(table_name: str) -> dict[str, Any]:
245257
{"AttributeName": "GSI1SK", "AttributeType": "S"},
246258
{"AttributeName": "GSI2PK", "AttributeType": "S"},
247259
{"AttributeName": "GSI2SK", "AttributeType": "S"},
260+
{"AttributeName": "GSI3PK", "AttributeType": "S"},
261+
{"AttributeName": "GSI3SK", "AttributeType": "S"},
248262
],
249263
"KeySchema": [
250264
{"AttributeName": "PK", "KeyType": "HASH"},
@@ -267,6 +281,14 @@ def get_table_definition(table_name: str) -> dict[str, Any]:
267281
],
268282
"Projection": {"ProjectionType": "ALL"},
269283
},
284+
{
285+
"IndexName": GSI3_NAME,
286+
"KeySchema": [
287+
{"AttributeName": "GSI3PK", "KeyType": "HASH"},
288+
{"AttributeName": "GSI3SK", "KeyType": "RANGE"},
289+
],
290+
"Projection": {"ProjectionType": "KEYS_ONLY"},
291+
},
270292
],
271293
"StreamSpecification": {
272294
"StreamEnabled": True,

0 commit comments

Comments
 (0)