Add pricing tables API functionality

- Introduced a new `pricing_tables` module with endpoints for creating, retrieving, and deleting pricing tables, accessible based on user roles.
- Updated the main application to include the new pricing tables router.
- Added corresponding schemas for pricing table creation and response.
- Implemented tests to validate the functionality and access control for pricing table operations.

Author: PhiRho

app/api/pricing_tables.py

@@ -0,0 +1,90 @@
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlalchemy.orm import Session
+from datetime import datetime, UTC
+
+from app.db.database import get_db
+from app.db.models import DBSystemSecret
+from app.core.security import check_system_admin, get_role_min_team_admin
+from app.schemas.models import PricingTableCreate, PricingTableResponse
+
+router = APIRouter(
+    tags=["pricing-tables"]
+)
+
+@router.post("", response_model=PricingTableResponse, status_code=status.HTTP_201_CREATED, dependencies=[Depends(check_system_admin)])
+@router.post("/", response_model=PricingTableResponse, status_code=status.HTTP_201_CREATED, dependencies=[Depends(check_system_admin)])
+async def create_pricing_table(
+    pricing_table: PricingTableCreate,
+    db: Session = Depends(get_db)
+):
+    """
+    Create or update the current pricing table. Only accessible by system admin users.
+    There can only be one active pricing table at a time.
+    """
+    # Check if a pricing table already exists
+    existing_table = db.query(DBSystemSecret).filter(DBSystemSecret.key == "CurrentPricingTable").first()
+
+    if existing_table:
+        # Update existing table
+        existing_table.value = pricing_table.pricing_table_id
+        existing_table.updated_at = datetime.now(UTC)
+        db.commit()
+        db.refresh(existing_table)
+        return PricingTableResponse(
+            pricing_table_id=existing_table.value,
+            updated_at=existing_table.updated_at
+        )
+    else:
+        # Create new table
+        db_table = DBSystemSecret(
+            key="CurrentPricingTable",
+            value=pricing_table.pricing_table_id,
+            description="Current Stripe pricing table ID",
+            created_at=datetime.now(UTC)
+        )
+        db.add(db_table)
+        db.commit()
+        db.refresh(db_table)
+        return PricingTableResponse(
+            pricing_table_id=db_table.value,
+            updated_at=db_table.created_at
+        )
+
+@router.get("", response_model=PricingTableResponse, dependencies=[Depends(get_role_min_team_admin)])
+@router.get("/", response_model=PricingTableResponse, dependencies=[Depends(get_role_min_team_admin)])
+async def get_pricing_table(
+    db: Session = Depends(get_db)
+):
+    """
+    Get the current pricing table ID. Only accessible by team admin users or higher privileges.
+    """
+    pricing_table = db.query(DBSystemSecret).filter(DBSystemSecret.key == "CurrentPricingTable").first()
+    if not pricing_table:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="No pricing table found"
+        )
+    return PricingTableResponse(
+        pricing_table_id=pricing_table.value,
+        updated_at=pricing_table.updated_at or pricing_table.created_at
+    )
+
+@router.delete("", dependencies=[Depends(check_system_admin)])
+@router.delete("/", dependencies=[Depends(check_system_admin)])
+async def delete_pricing_table(
+    db: Session = Depends(get_db)
+):
+    """
+    Delete the current pricing table. Only accessible by system admin users.
+    """
+    pricing_table = db.query(DBSystemSecret).filter(DBSystemSecret.key == "CurrentPricingTable").first()
+    if not pricing_table:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="No pricing table found"
+        )
+
+    db.delete(pricing_table)
+    db.commit()
+
+    return {"message": "Pricing table deleted successfully"}

app/main.py

@@ -5,7 +5,7 @@
 from fastapi.openapi.docs import get_swagger_ui_html
 from fastapi.openapi.utils import get_openapi
 from prometheus_fastapi_instrumentator import Instrumentator, metrics
-from app.api import auth, private_ai_keys, users, regions, audit, teams, billing, products
+from app.api import auth, private_ai_keys, users, regions, audit, teams, billing, products, pricing_tables
 from app.core.config import settings
 from app.db.database import get_db
 from app.middleware.audit import AuditLogMiddleware
@@ -134,6 +134,7 @@ async def health_check():
 app.include_router(teams.router, prefix="/teams", tags=["teams"])
 app.include_router(billing.router, prefix="/billing", tags=["billing"])
 app.include_router(products.router, prefix="/products", tags=["products"])
+app.include_router(pricing_tables.router, prefix="/pricing-tables", tags=["pricing-tables"])
 
 @app.get("/", include_in_schema=False)
 async def custom_swagger_ui_html():

app/schemas/models.py

@@ -311,4 +311,13 @@ class Product(ProductBase):
 
 class PricingTableSession(BaseModel):
     client_secret: str
+    model_config = ConfigDict(from_attributes=True)
+
+class PricingTableCreate(BaseModel):
+    pricing_table_id: str
+    model_config = ConfigDict(from_attributes=True)
+
+class PricingTableResponse(BaseModel):
+    pricing_table_id: str
+    updated_at: datetime
     model_config = ConfigDict(from_attributes=True)

tests/test_pricing_tables.py

@@ -0,0 +1,238 @@
+import pytest
+from fastapi import status
+from datetime import datetime, UTC
+from app.db.models import DBSystemSecret, DBUser
+from app.core.security import get_password_hash
+
+def test_create_pricing_table_system_admin(client, db, test_admin):
+    """Test that a system admin can create a pricing table"""
+    # Login as system admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_admin.email, "password": "adminpassword"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Create pricing table
+    response = client.post(
+        "/pricing-tables",
+        json={"pricing_table_id": "test_pricing_table_123"},
+        headers=headers
+    )
+    assert response.status_code == status.HTTP_201_CREATED
+    data = response.json()
+    assert data["pricing_table_id"] == "test_pricing_table_123"
+    assert "updated_at" in data
+
+    # Verify in database
+    pricing_table = db.query(DBSystemSecret).filter(
+        DBSystemSecret.key == "CurrentPricingTable"
+    ).first()
+    assert pricing_table is not None
+    assert pricing_table.value == "test_pricing_table_123"
+
+def test_create_pricing_table_team_admin(client, db, test_team_admin):
+    """Test that a team admin cannot create a pricing table"""
+    # Login as team admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_team_admin.email, "password": "password123"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Try to create pricing table
+    response = client.post(
+        "/pricing-tables",
+        json={"pricing_table_id": "test_pricing_table_123"},
+        headers=headers
+    )
+    assert response.status_code == status.HTTP_403_FORBIDDEN
+
+def test_get_pricing_table_team_admin(client, db, test_team_admin):
+    """Test that a team admin can get the pricing table"""
+    # Create pricing table as system admin
+    system_admin = DBUser(
+        email="system_admin@test.com",
+        hashed_password=get_password_hash("testpassword"),
+        is_admin=True
+    )
+    db.add(system_admin)
+    db.commit()
+
+    # Create pricing table
+    pricing_table = DBSystemSecret(
+        key="CurrentPricingTable",
+        value="test_pricing_table_123",
+        description="Test pricing table",
+        created_at=datetime.now(UTC)
+    )
+    db.add(pricing_table)
+    db.commit()
+
+    # Login as team admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_team_admin.email, "password": "password123"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Get pricing table
+    response = client.get("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_200_OK
+    data = response.json()
+    assert data["pricing_table_id"] == "test_pricing_table_123"
+    assert "updated_at" in data
+
+def test_get_pricing_table_not_found(client, db, test_team_admin):
+    """Test getting pricing table when none exists"""
+    # Login as team admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_team_admin.email, "password": "password123"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Try to get pricing table
+    response = client.get("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_404_NOT_FOUND
+
+def test_delete_pricing_table_system_admin(client, db, test_admin):
+    """Test that a system admin can delete the pricing table"""
+    # Create pricing table
+    pricing_table = DBSystemSecret(
+        key="CurrentPricingTable",
+        value="test_pricing_table_123",
+        description="Test pricing table",
+        created_at=datetime.now(UTC)
+    )
+    db.add(pricing_table)
+    db.commit()
+
+    # Login as system admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_admin.email, "password": "adminpassword"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Delete pricing table
+    response = client.delete("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_200_OK
+    assert response.json()["message"] == "Pricing table deleted successfully"
+
+    # Verify deleted from database
+    pricing_table = db.query(DBSystemSecret).filter(
+        DBSystemSecret.key == "CurrentPricingTable"
+    ).first()
+    assert pricing_table is None
+
+def test_delete_pricing_table_team_admin(client, db, test_team_admin):
+    """Test that a team admin cannot delete the pricing table"""
+    # Create pricing table
+    pricing_table = DBSystemSecret(
+        key="CurrentPricingTable",
+        value="test_pricing_table_123",
+        description="Test pricing table",
+        created_at=datetime.now(UTC)
+    )
+    db.add(pricing_table)
+    db.commit()
+
+    # Login as team admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_team_admin.email, "password": "password123"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Try to delete pricing table
+    response = client.delete("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_403_FORBIDDEN
+
+def test_update_existing_pricing_table(client, db, test_admin):
+    """Test updating an existing pricing table"""
+    # Create initial pricing table
+    pricing_table = DBSystemSecret(
+        key="CurrentPricingTable",
+        value="initial_pricing_table",
+        description="Test pricing table",
+        created_at=datetime.now(UTC)
+    )
+    db.add(pricing_table)
+    db.commit()
+
+    # Login as system admin
+    response = client.post(
+        "/auth/login",
+        data={"username": test_admin.email, "password": "adminpassword"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Update pricing table
+    response = client.post(
+        "/pricing-tables",
+        json={"pricing_table_id": "updated_pricing_table"},
+        headers=headers
+    )
+    assert response.status_code == status.HTTP_201_CREATED
+    data = response.json()
+    assert data["pricing_table_id"] == "updated_pricing_table"
+    assert "updated_at" in data
+
+    # Verify only one pricing table exists with updated value
+    pricing_tables = db.query(DBSystemSecret).filter(
+        DBSystemSecret.key == "CurrentPricingTable"
+    ).all()
+    assert len(pricing_tables) == 1
+    assert pricing_tables[0].value == "updated_pricing_table"
+
+def test_key_creator_cannot_access_pricing_table(client, db, test_team_key_creator):
+    """Test that a key_creator cannot access the pricing table"""
+    # Create pricing table as system admin
+    pricing_table = DBSystemSecret(
+        key="CurrentPricingTable",
+        value="test_pricing_table_123",
+        description="Test pricing table",
+        created_at=datetime.now(UTC)
+    )
+    db.add(pricing_table)
+    db.commit()
+
+    # Login as key_creator
+    response = client.post(
+        "/auth/login",
+        data={"username": test_team_key_creator.email, "password": "password123"}
+    )
+    assert response.status_code == status.HTTP_200_OK
+    token = response.json()["access_token"]
+    headers = {"Authorization": f"Bearer {token}"}
+
+    # Try to get pricing table
+    response = client.get("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_403_FORBIDDEN
+
+    # Try to create pricing table
+    response = client.post(
+        "/pricing-tables",
+        json={"pricing_table_id": "new_pricing_table"},
+        headers=headers
+    )
+    assert response.status_code == status.HTTP_403_FORBIDDEN
+
+    # Try to delete pricing table
+    response = client.delete("/pricing-tables", headers=headers)
+    assert response.status_code == status.HTTP_403_FORBIDDEN