zae-limiter 0.1.0__py3-none-any.whl

zae_limiter/exceptions.py

@@ -0,0 +1,214 @@
+"""Exceptions for zae-limiter."""
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .models import LimitStatus
+
+
+class RateLimitError(Exception):
+    """Base exception for rate limiting errors."""
+
+    pass
+
+
+class RateLimitExceeded(RateLimitError):  # noqa: N818
+    """
+    Raised when one or more rate limits would be exceeded.
+
+    Provides full visibility into ALL limits that were checked,
+    both passed and failed, to help callers understand the full picture.
+
+    Attributes:
+        statuses: Status of ALL limits checked (both passed and failed)
+        violations: Only the limits that were exceeded
+        passed: Only the limits that passed
+        retry_after_seconds: Time until ALL requested capacity is available
+        primary_violation: The violation with longest retry time (bottleneck)
+    """
+
+    def __init__(self, statuses: list["LimitStatus"]) -> None:
+        self.statuses = statuses
+        self.violations = [s for s in statuses if s.exceeded]
+        self.passed = [s for s in statuses if not s.exceeded]
+
+        if not self.violations:
+            raise ValueError("RateLimitExceeded requires at least one violation")
+
+        self.primary_violation = max(self.violations, key=lambda v: v.retry_after_seconds)
+        self.retry_after_seconds = self.primary_violation.retry_after_seconds
+
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        v = self.primary_violation
+        exceeded_names = ", ".join(s.limit_name for s in self.violations)
+        return (
+            f"Rate limit exceeded for {v.entity_id}/{v.resource}: "
+            f"[{exceeded_names}]. "
+            f"Retry after {self.retry_after_seconds:.1f}s"
+        )
+
+    def as_dict(self) -> dict[str, Any]:
+        """
+        Serialize for JSON API responses.
+
+        Returns a dictionary suitable for returning in a 429 response body.
+        """
+        return {
+            "error": "rate_limit_exceeded",
+            "message": str(self),
+            "retry_after_seconds": self.retry_after_seconds,
+            "retry_after_ms": int(self.retry_after_seconds * 1000),
+            "limits": [
+                {
+                    "entity_id": s.entity_id,
+                    "resource": s.resource,
+                    "limit_name": s.limit_name,
+                    "capacity": s.limit.capacity,
+                    "burst": s.limit.burst,
+                    "available": s.available,
+                    "requested": s.requested,
+                    "exceeded": s.exceeded,
+                    "retry_after_seconds": s.retry_after_seconds,
+                }
+                for s in self.statuses
+            ],
+        }
+
+    @property
+    def retry_after_header(self) -> str:
+        """Value for HTTP Retry-After header (integer seconds)."""
+        return str(int(self.retry_after_seconds) + 1)  # round up
+
+
+class RateLimiterUnavailable(RateLimitError):  # noqa: N818
+    """
+    Raised when DynamoDB is unavailable and failure_mode=FAIL_CLOSED.
+
+    This indicates a transient infrastructure issue, not a rate limit.
+    """
+
+    def __init__(self, message: str, cause: Exception | None = None) -> None:
+        self.cause = cause
+        super().__init__(message)
+
+
+class EntityNotFoundError(RateLimitError):
+    """Raised when an entity is not found."""
+
+    def __init__(self, entity_id: str) -> None:
+        self.entity_id = entity_id
+        super().__init__(f"Entity not found: {entity_id}")
+
+
+class EntityExistsError(RateLimitError):
+    """Raised when trying to create an entity that already exists."""
+
+    def __init__(self, entity_id: str) -> None:
+        self.entity_id = entity_id
+        super().__init__(f"Entity already exists: {entity_id}")
+
+
+class StackCreationError(Exception):
+    """Raised when CloudFormation stack creation fails."""
+
+    def __init__(
+        self, stack_name: str, reason: str, events: list[dict[str, Any]] | None = None
+    ) -> None:
+        self.stack_name = stack_name
+        self.reason = reason
+        self.events = events or []
+        super().__init__(f"Stack {stack_name} creation failed: {reason}")
+
+
+class StackAlreadyExistsError(StackCreationError):
+    """Raised when stack already exists (informational)."""
+
+    pass
+
+
+# ---------------------------------------------------------------------------
+# Version-related exceptions
+# ---------------------------------------------------------------------------
+
+
+class VersionError(RateLimitError):
+    """Base class for version-related errors."""
+
+    pass
+
+
+class VersionMismatchError(VersionError):
+    """
+    Raised when client and infrastructure versions are incompatible.
+
+    This error indicates that the client library version doesn't match
+    the deployed infrastructure and auto-update is disabled or failed.
+    """
+
+    def __init__(
+        self,
+        client_version: str,
+        schema_version: str,
+        lambda_version: str | None,
+        message: str,
+        can_auto_update: bool = False,
+    ) -> None:
+        self.client_version = client_version
+        self.schema_version = schema_version
+        self.lambda_version = lambda_version
+        self.can_auto_update = can_auto_update
+        super().__init__(self._format_message(message))
+
+    def _format_message(self, message: str) -> str:
+        return (
+            f"Version mismatch: client={self.client_version}, "
+            f"schema={self.schema_version}, "
+            f"lambda={self.lambda_version or 'unknown'}. {message}"
+        )
+
+
+class IncompatibleSchemaError(VersionError):
+    """
+    Raised when schema version requires manual migration.
+
+    This indicates a major version difference that cannot be
+    automatically reconciled.
+    """
+
+    def __init__(
+        self,
+        client_version: str,
+        schema_version: str,
+        message: str,
+        migration_guide_url: str | None = None,
+    ) -> None:
+        self.client_version = client_version
+        self.schema_version = schema_version
+        self.migration_guide_url = migration_guide_url
+        msg = (
+            f"Incompatible schema: client {client_version} is not compatible "
+            f"with schema {schema_version}. {message}"
+        )
+        if migration_guide_url:
+            msg += f" See: {migration_guide_url}"
+        super().__init__(msg)
+
+
+class InfrastructureNotFoundError(VersionError):
+    """
+    Raised when expected infrastructure doesn't exist.
+
+    This typically means the CloudFormation stack or DynamoDB table
+    hasn't been deployed yet.
+    """
+
+    def __init__(self, table_name: str, stack_name: str | None = None) -> None:
+        self.table_name = table_name
+        self.stack_name = stack_name
+        msg = f"Infrastructure not found for table '{table_name}'"
+        if stack_name:
+            msg += f" (stack: {stack_name})"
+        msg += ". Run 'zae-limiter deploy' or use create_stack=True."
+        super().__init__(msg)

zae_limiter/infra/__init__.py

@@ -0,0 +1,10 @@
+"""Infrastructure as code for zae-limiter."""
+
+from pathlib import Path
+
+CFN_TEMPLATE_PATH = Path(__file__).parent / "cfn_template.yaml"
+
+
+def get_cfn_template() -> str:
+    """Get the CloudFormation template as a string."""
+    return CFN_TEMPLATE_PATH.read_text()

zae_limiter/infra/cfn_template.yaml

@@ -0,0 +1,255 @@
+AWSTemplateFormatVersion: '2010-09-09'
+Description: >
+  zae-limiter infrastructure - DynamoDB table with streams and
+  Lambda aggregator for usage snapshots.
+
+Parameters:
+  TableName:
+    Type: String
+    Default: rate_limits
+    Description: Name of the DynamoDB table
+
+  SnapshotWindows:
+    Type: String
+    Default: hourly,daily
+    Description: Comma-separated list of snapshot windows (hourly, daily, monthly)
+
+  SnapshotRetentionDays:
+    Type: Number
+    Default: 90
+    Description: Number of days to retain usage snapshots
+    MinValue: 1
+    MaxValue: 3650
+
+  LambdaMemorySize:
+    Type: Number
+    Default: 256
+    Description: Memory size for the aggregator Lambda
+    MinValue: 128
+    MaxValue: 3008
+
+  LambdaTimeout:
+    Type: Number
+    Default: 60
+    Description: Timeout for the aggregator Lambda in seconds
+    MinValue: 1
+    MaxValue: 900
+
+  EnableAggregator:
+    Type: String
+    Default: 'true'
+    AllowedValues:
+      - 'true'
+      - 'false'
+    Description: Whether to deploy the aggregator Lambda
+
+  SchemaVersion:
+    Type: String
+    Default: '1.0.0'
+    Description: Schema version for the rate limiter infrastructure
+
+Conditions:
+  DeployAggregator: !Equals [!Ref EnableAggregator, 'true']
+
+Resources:
+  # -------------------------------------------------------------------------
+  # DynamoDB Table
+  # -------------------------------------------------------------------------
+  RateLimitsTable:
+    Type: AWS::DynamoDB::Table
+    Properties:
+      TableName: !Ref TableName
+      BillingMode: PAY_PER_REQUEST
+
+      AttributeDefinitions:
+        - AttributeName: PK
+          AttributeType: S
+        - AttributeName: SK
+          AttributeType: S
+        - AttributeName: GSI1PK
+          AttributeType: S
+        - AttributeName: GSI1SK
+          AttributeType: S
+        - AttributeName: GSI2PK
+          AttributeType: S
+        - AttributeName: GSI2SK
+          AttributeType: S
+
+      KeySchema:
+        - AttributeName: PK
+          KeyType: HASH
+        - AttributeName: SK
+          KeyType: RANGE
+
+      GlobalSecondaryIndexes:
+        - IndexName: GSI1
+          KeySchema:
+            - AttributeName: GSI1PK
+              KeyType: HASH
+            - AttributeName: GSI1SK
+              KeyType: RANGE
+          Projection:
+            ProjectionType: ALL
+
+        - IndexName: GSI2
+          KeySchema:
+            - AttributeName: GSI2PK
+              KeyType: HASH
+            - AttributeName: GSI2SK
+              KeyType: RANGE
+          Projection:
+            ProjectionType: ALL
+
+      TimeToLiveSpecification:
+        AttributeName: ttl
+        Enabled: true
+
+      StreamSpecification:
+        StreamEnabled: true
+        StreamViewType: NEW_AND_OLD_IMAGES
+
+      Tags:
+        - Key: Application
+          Value: zae-limiter
+
+  # -------------------------------------------------------------------------
+  # Lambda Aggregator
+  # -------------------------------------------------------------------------
+  AggregatorLogGroup:
+    Type: AWS::Logs::LogGroup
+    Condition: DeployAggregator
+    Properties:
+      LogGroupName: !Sub /aws/lambda/${TableName}-aggregator
+      RetentionInDays: 30
+
+  AggregatorRole:
+    Type: AWS::IAM::Role
+    Condition: DeployAggregator
+    Properties:
+      RoleName: !Sub ${TableName}-aggregator-role
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: lambda.amazonaws.com
+            Action: sts:AssumeRole
+
+      ManagedPolicyArns:
+        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+
+      Policies:
+        - PolicyName: DynamoDBAccess
+          PolicyDocument:
+            Version: '2012-10-17'
+            Statement:
+              # Read/write access to table
+              - Effect: Allow
+                Action:
+                  - dynamodb:GetItem
+                  - dynamodb:PutItem
+                  - dynamodb:UpdateItem
+                  - dynamodb:Query
+                Resource:
+                  - !GetAtt RateLimitsTable.Arn
+                  - !Sub ${RateLimitsTable.Arn}/index/*
+
+              # Stream access
+              - Effect: Allow
+                Action:
+                  - dynamodb:GetRecords
+                  - dynamodb:GetShardIterator
+                  - dynamodb:DescribeStream
+                  - dynamodb:ListStreams
+                Resource:
+                  - !GetAtt RateLimitsTable.StreamArn
+
+  AggregatorFunction:
+    Type: AWS::Lambda::Function
+    Condition: DeployAggregator
+    DependsOn: AggregatorLogGroup
+    Properties:
+      FunctionName: !Sub ${TableName}-aggregator
+      Description: Aggregates rate limiter usage into hourly/daily snapshots
+      Runtime: python3.12
+      Handler: zae_limiter.aggregator.handler.handler
+      MemorySize: !Ref LambdaMemorySize
+      Timeout: !Ref LambdaTimeout
+      Role: !GetAtt AggregatorRole.Arn
+
+      Environment:
+        Variables:
+          TABLE_NAME: !Ref TableName
+          SNAPSHOT_WINDOWS: !Ref SnapshotWindows
+          SNAPSHOT_TTL_DAYS: !Ref SnapshotRetentionDays
+          ZAE_LIMITER_SCHEMA_VERSION: !Ref SchemaVersion
+
+      # Note: Code must be deployed separately via SAM, CDK, or manual upload
+      # This is a placeholder that will need to be replaced
+      Code:
+        ZipFile: |
+          def handler(event, context):
+              # Placeholder - replace with actual deployment
+              return {"statusCode": 200, "body": "Placeholder"}
+
+      Tags:
+        - Key: Application
+          Value: zae-limiter
+
+  AggregatorEventSourceMapping:
+    Type: AWS::Lambda::EventSourceMapping
+    Condition: DeployAggregator
+    Properties:
+      EventSourceArn: !GetAtt RateLimitsTable.StreamArn
+      FunctionName: !Ref AggregatorFunction
+      StartingPosition: LATEST
+      BatchSize: 100
+      MaximumBatchingWindowInSeconds: 5
+
+      # Filter to only process MODIFY events on BUCKET records
+      FilterCriteria:
+        Filters:
+          - Pattern: >-
+              {
+                "eventName": ["MODIFY"],
+                "dynamodb": {
+                  "NewImage": {
+                    "SK": {
+                      "S": [{"prefix": "#BUCKET#"}]
+                    }
+                  }
+                }
+              }
+
+Outputs:
+  TableName:
+    Description: DynamoDB table name
+    Value: !Ref RateLimitsTable
+    Export:
+      Name: !Sub ${AWS::StackName}-TableName
+
+  TableArn:
+    Description: DynamoDB table ARN
+    Value: !GetAtt RateLimitsTable.Arn
+    Export:
+      Name: !Sub ${AWS::StackName}-TableArn
+
+  StreamArn:
+    Description: DynamoDB stream ARN
+    Value: !GetAtt RateLimitsTable.StreamArn
+    Export:
+      Name: !Sub ${AWS::StackName}-StreamArn
+
+  AggregatorFunctionArn:
+    Condition: DeployAggregator
+    Description: Aggregator Lambda function ARN
+    Value: !GetAtt AggregatorFunction.Arn
+    Export:
+      Name: !Sub ${AWS::StackName}-AggregatorArn
+
+  AggregatorFunctionName:
+    Condition: DeployAggregator
+    Description: Aggregator Lambda function name
+    Value: !Ref AggregatorFunction
+    Export:
+      Name: !Sub ${AWS::StackName}-AggregatorName

zae_limiter/infra/lambda_builder.py

@@ -0,0 +1,85 @@
+"""Build Lambda deployment packages for zae-limiter aggregator."""
+
+import io
+import zipfile
+from pathlib import Path
+
+
+def build_lambda_package() -> bytes:
+    """
+    Build Lambda deployment package from installed zae_limiter package.
+
+    Creates a zip file containing the entire zae_limiter package, which
+    includes all necessary code for the aggregator Lambda function.
+
+    The package only depends on boto3, which is provided by the Lambda
+    runtime, so no external dependencies need to be bundled.
+
+    Returns:
+        Zip file contents as bytes
+    """
+    import zae_limiter
+
+    # Find installed package location
+    package_path = Path(zae_limiter.__file__).parent
+
+    zip_buffer = io.BytesIO()
+
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
+        # Add all Python files from the package
+        for py_file in package_path.rglob("*.py"):
+            # Create archive name: zae_limiter/...
+            arcname = py_file.relative_to(package_path.parent)
+            zf.write(py_file, arcname)
+
+        # Also include the CloudFormation template for reference
+        # (not used by Lambda, but useful for debugging)
+        cfn_template = package_path / "infra" / "cfn_template.yaml"
+        if cfn_template.exists():
+            arcname = cfn_template.relative_to(package_path.parent)
+            zf.write(cfn_template, arcname)
+
+    zip_buffer.seek(0)
+    return zip_buffer.getvalue()
+
+
+def write_lambda_package(output_path: str | Path) -> int:
+    """
+    Build and write Lambda package to a file.
+
+    Args:
+        output_path: Path where to write the zip file
+
+    Returns:
+        Size of the written file in bytes
+    """
+    zip_bytes = build_lambda_package()
+    output_path = Path(output_path)
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_bytes(zip_bytes)
+
+    return len(zip_bytes)
+
+
+def get_package_info() -> dict[str, str | int]:
+    """
+    Get information about the Lambda package without building it.
+
+    Returns:
+        Dict with package metadata
+    """
+    import zae_limiter
+
+    package_path = Path(zae_limiter.__file__).parent
+
+    # Count files
+    py_files = list(package_path.rglob("*.py"))
+    total_size = sum(f.stat().st_size for f in py_files)
+
+    return {
+        "package_path": str(package_path),
+        "python_files": len(py_files),
+        "uncompressed_size": total_size,
+        "handler": "zae_limiter.aggregator.handler.handler",
+    }