Merge branch 'main' into nm/post-collect-step

fixbackend/app.py

@@ -76,8 +76,8 @@ async def custom_ui(hash: str) -> Response:
         headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains; preload"
         headers["Content-Security-Policy"] = (
             "default-src 'self' https://cdn.fix.security;"
-            f" connect-src 'self' data: https://capture.trackjs.com https://ph.fix.security;"
-            f" script-src 'self' 'nonce-{nonce}' https://cdn.fix.security https://www.googletagmanager.com;"
+            f" connect-src 'self' data: https://cdn.fix.security https://capture.trackjs.com https://ph.fix.security;"
+            f" script-src 'self' 'nonce-{nonce}' https://ph.fix.security https://cdn.fix.security https://www.googletagmanager.com;"
             f" style-src 'self' 'nonce-{nonce}' https://cdn.fix.security;"
             " font-src 'self' data: https://cdn.fix.security;"
             " img-src 'self' data: https://cdn.fix.security https://usage.trackjs.com https://i.ytimg.com https://www.googletagmanager.com/;"
@@ -330,8 +330,8 @@ async def root(_: Request) -> Response:
             headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains; preload"
             headers["Content-Security-Policy"] = (
                 "default-src 'self' https://cdn.fix.security;"
-                f" connect-src 'self' data: https://capture.trackjs.com https://ph.fix.security;"
-                f" script-src 'self' 'nonce-{nonce}' https://cdn.fix.security https://www.googletagmanager.com;"
+                f" connect-src 'self' data: https://cdn.fix.security https://capture.trackjs.com https://ph.fix.security;"
+                f" script-src 'self' 'nonce-{nonce}' https://ph.fix.security https://cdn.fix.security https://www.googletagmanager.com;"
                 f" style-src 'self' 'nonce-{nonce}' https://cdn.fix.security;"
                 " font-src 'self' data: https://cdn.fix.security;"
                 " img-src 'self' data: https://cdn.fix.security https://usage.trackjs.com https://i.ytimg.com https://www.googletagmanager.com/;"

fixbackend/app_dependencies.py

@@ -84,6 +84,7 @@
 from fixbackend.workspaces.invitation_repository import InvitationRepositoryImpl
 from fixbackend.workspaces.repository import WorkspaceRepositoryImpl
 from fixbackend.workspaces.trial_end_service import TrialEndService
+from fixbackend.workspaces.free_tier_deletion_service import FreeTierCleanupService
 
 log = logging.getLogger(__name__)
 
@@ -441,8 +442,10 @@ async def dispatcher_dependencies(cfg: Config) -> FixDependencies:
         ),
     )
 
-    # uncomment once aws marketplace suscriptions are available on prd
     deps.add(SN.trial_end_service, TrialEndService(workspace_repo, session_maker, cloud_account_service))
+    deps.add(
+        SN.free_tier_cleanup_service, FreeTierCleanupService(workspace_repo, session_maker, cloud_account_service, cfg)
+    )
 
     gcp_account_repo = deps.add(
         SN.gcp_service_account_repo,

fixbackend/config.py

@@ -84,6 +84,7 @@ class Config(BaseSettings):
     stripe_api_key: Optional[str]
     stripe_webhook_key: Optional[str]
     customer_support_users: List[str]
+    free_tier_cleanup_timeout_days: int
 
     def frontend_cdn_origin(self) -> str:
         return f"{self.cdn_endpoint}/{self.cdn_bucket}/{self.fixui_sha}"
@@ -191,6 +192,11 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> Namespace:
     parser.add_argument(
         "--customer-support-users", nargs="+", default=os.environ.get("CUSTOMER_SUPPORT_USERS", "").split(",")
     )
+    parser.add_argument(
+        "--free-tier-cleanup-timeout-days",
+        type=int,
+        default=int(os.environ.get("FREE_TIER_CLEANUP_TIMEOUT_DAYS", "7")),
+    )
     return parser.parse_known_args(argv if argv is not None else sys.argv[1:])[0]
 
 

fixbackend/dependencies.py

@@ -81,6 +81,7 @@ class ServiceNames:
     azure_subscription_repo = "azure_subscription_repo"
     azure_subscription_service = "azure_subscription_service"
     trial_end_service = "trial_end_service"
+    free_tier_cleanup_service = "free_tier_cleanup_service"
 
 
 class FixDependencies(Dependencies):

fixbackend/inventory/inventory_router.py

@@ -21,7 +21,7 @@
 
 from fixbackend.dependencies import FixDependencies, FixDependency, ServiceNames
 from fixbackend.graph_db.models import GraphDatabaseAccess
-from fixbackend.ids import NodeId
+from fixbackend.ids import NodeId, ProductTier
 from fixbackend.inventory.inventory_service import InventoryService
 from fixbackend.inventory.inventory_schemas import (
     CompletePathRequest,
@@ -350,9 +350,14 @@ async def stream() -> AsyncIterator[str]:
         return StreamOnSuccessResponse(stream(), media_type=media_type)
 
     @router.get("/workspace-info", tags=["report"])
-    async def workspace_info(graph_db: CurrentGraphDbDependency) -> InventorySummaryRead:
+    async def workspace_info(
+        graph_db: CurrentGraphDbDependency, workspace: UserWorkspaceDependency
+    ) -> InventorySummaryRead:
         now = utc()
-        info = await inventory().inventory_summary(graph_db, now, timedelta(days=7))
+        duration = timedelta(days=7)
+        if workspace.current_product_tier() == ProductTier.Free:
+            duration = timedelta(days=31)
+        info = await inventory().inventory_summary(graph_db, now, duration)
         return InventorySummaryRead(
             resources_per_account_timeline=info.resources_per_account_timeline,
             score_progress=info.score_progress,

fixbackend/workspaces/free_tier_deletion_service.py

@@ -0,0 +1,89 @@
+#  Copyright (c) 2024. Some Engineering
+#  This program is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Affero General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Affero General Public License for more details.
+#
+#  You should have received a copy of the GNU Affero General Public License
+#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+
+import logging
+from datetime import timedelta
+from typing import Any, Optional
+
+from fixcloudutils.asyncio.periodic import Periodic
+from fixcloudutils.service import Service
+
+from fixbackend.cloud_accounts.service import CloudAccountService
+from fixbackend.config import ProductTierSettings
+from fixbackend.ids import ProductTier
+from fixbackend.types import AsyncSessionMaker
+from fixbackend.workspaces.repository import WorkspaceRepository
+from fixbackend.config import Config
+
+log = logging.getLogger(__name__)
+
+
+class FreeTierCleanupService(Service):
+
+    def __init__(
+        self,
+        workspace_repository: WorkspaceRepository,
+        session_maker: AsyncSessionMaker,
+        cloud_account_service: CloudAccountService,
+        config: Config,
+    ):
+        self.workspace_repository = workspace_repository
+        self.cloud_account_service = cloud_account_service
+        self.session_maker = session_maker
+        self.config = config
+        self.periodic: Optional[Periodic] = Periodic(
+            "clean_up_free_tiers",
+            self.cleanup_free_tiers,
+            frequency=timedelta(minutes=60),
+            first_run=timedelta(seconds=30),
+        )
+        self.free_tier_cleanup_timeout = timedelta(days=self.config.free_tier_cleanup_timeout_days)
+
+    async def start(self) -> Any:
+        if self.periodic:
+            await self.periodic.start()
+
+    async def stop(self) -> None:
+        if self.periodic:
+            await self.periodic.stop()
+
+    async def cleanup_free_tiers(self) -> None:
+        workspaces = await self.workspace_repository.list_overdue_free_tier_cleanup(
+            been_in_free_tier_for=self.free_tier_cleanup_timeout
+        )
+        for workspace in workspaces:
+            account_limit = ProductTierSettings[ProductTier.Free].account_limit
+
+            if account_limit is None:
+                continue
+
+            accounts = await self.cloud_account_service.list_accounts(workspace.id)
+
+            if len(accounts) <= account_limit:
+                continue
+
+            log.info(
+                f"Cleaning up workspace {workspace.id}"
+                " because it has been in free tier for"
+                f"{self.free_tier_cleanup_timeout}."
+            )
+            for i, account in enumerate(accounts):
+                if i < account_limit:
+                    continue
+
+                log.info(
+                    f"Deleting cloud account {account.id} on workspace {workspace.id} " "because it is over the limit."
+                )
+                await self.cloud_account_service.delete_cloud_account(workspace.owner_id, account.id, workspace.id)

fixbackend/workspaces/models/__init__.py

@@ -41,6 +41,7 @@ class Workspace:
     )
     current_cycle_ends_at: Optional[datetime] = None  # when the active product tier ends, typically end of the month
     move_to_free_acknowledged_at: Optional[datetime] = None  # only set for list_workspaces when the user_id is provided
+    tier_updated_at: Optional[datetime] = None
 
     # this is the product tier that is active for the workspace at the moment
     # it is based on the highest tier we saw during the billing cycle

fixbackend/workspaces/models/orm.py

@@ -43,6 +43,7 @@ class Organization(Base, CreatedUpdatedMixin):
     created_at: Mapped[datetime] = mapped_column(UTCDateTime, server_default=func.now(), index=True)
     highest_current_cycle_tier: Mapped[Optional[str]] = mapped_column(String(length=64), nullable=True)
     current_cycle_ends_at: Mapped[Optional[datetime]] = mapped_column(UTCDateTime, nullable=True)
+    tier_updated_at: Mapped[Optional[datetime]] = mapped_column(UTCDateTime, nullable=True, index=True)
 
     def to_model(self, move_to_free_acknowledged_at: Optional[datetime] = None) -> models.Workspace:
         return models.Workspace(
@@ -62,6 +63,7 @@ def to_model(self, move_to_free_acknowledged_at: Optional[datetime] = None) -> m
             created_at=self.created_at,
             updated_at=self.updated_at,
             move_to_free_acknowledged_at=move_to_free_acknowledged_at,
+            tier_updated_at=self.tier_updated_at,
         )
 
 

fixbackend/workspaces/repository.py

@@ -128,6 +128,11 @@ async def ack_move_to_free(self, workspace_id: WorkspaceId, user_id: UserId) ->
         """Acknowledge that the workspace moved to free tier."""
         raise NotImplementedError
 
+    @abstractmethod
+    async def list_overdue_free_tier_cleanup(self, been_in_free_tier_for: timedelta) -> Sequence[Workspace]:
+        """List all workspaces where the free tier cleanup is overdue."""
+        raise NotImplementedError
+
 
 class WorkspaceRepositoryImpl(WorkspaceRepository):
     def __init__(
@@ -157,7 +162,12 @@ async def create_workspace(self, name: str, slug: str, owner: User) -> Workspace
                 slug = f"{slug}-{workspace_id}"
 
             organization = orm.Organization(
-                id=workspace_id, name=name, slug=slug, tier=ProductTier.Trial.value, owner_id=owner.id
+                id=workspace_id,
+                name=name,
+                slug=slug,
+                tier=ProductTier.Trial.value,
+                owner_id=owner.id,
+                tier_updated_at=utc(),
             )
             member_relationship = orm.OrganizationMembers(user_id=owner.id)
             organization.members.append(member_relationship)
@@ -331,6 +341,7 @@ async def do_tx(session: AsyncSession) -> Workspace:
                 )
                 workspace.current_cycle_ends_at = last_billing_cycle_instant
             workspace.highest_current_cycle_tier = max(active_tier, new_tier)
+            workspace.tier_updated_at = utc()
 
             await session.commit()
             await session.refresh(workspace)
@@ -420,6 +431,21 @@ async def ack_move_to_free(self, workspace_id: WorkspaceId, user_id: UserId) ->
 
             return evolve(workspace, move_to_free_acknowledged_at=now)
 
+    async def list_overdue_free_tier_cleanup(
+        self, been_in_free_tier_for: timedelta, window: timedelta = timedelta(days=1)
+    ) -> Sequence[Workspace]:
+        """List all workspaces where the free tier cleanup is overdue."""
+        async with self.session_maker() as session:
+            statement = (
+                select(orm.Organization)
+                .where(orm.Organization.tier == ProductTier.Free.value)
+                .where(orm.Organization.tier_updated_at < utc() - been_in_free_tier_for)
+                .where(orm.Organization.tier_updated_at > utc() - been_in_free_tier_for - window)
+            )
+            results = await session.execute(statement)
+            workspaces = results.unique().scalars().all()
+            return [ws.to_model() for ws in workspaces]
+
 
 async def get_workspace_repository(fix: FixDependency) -> WorkspaceRepository:
     return fix.service(ServiceNames.workspace_repo, WorkspaceRepositoryImpl)

migrations/versions/2024-07-30T17:09:40Z_free_tier_account_removal.py

@@ -0,0 +1,29 @@
+"""free tier account  removal
+
+Revision ID: dbe8f626f045
+Revises: 2e39a90ed4ac
+Create Date: 2024-07-30 17:09:40.293980+00:00
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from fixbackend.sqlalechemy_extensions import UTCDateTime
+
+
+# revision identifiers, used by Alembic.
+revision: str = "dbe8f626f045"
+down_revision: Union[str, None] = "2e39a90ed4ac"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column("organization", sa.Column("tier_updated_at", UTCDateTime(timezone=True), nullable=True))
+    op.create_index(op.f("ix_organization_tier_updated_at"), "organization", ["tier_updated_at"], unique=False)
+
+
+def downgrade() -> None:
+    pass

tests/fixbackend/conftest.py

@@ -236,6 +236,7 @@ def default_config() -> Config:
         stripe_api_key=None,
         stripe_webhook_key=None,
         customer_support_users=[],
+        free_tier_cleanup_timeout_days=7,
     )
 
 

tests/fixbackend/workspaces/repository_test.py

@@ -283,6 +283,34 @@ async def test_expired_trials(
     ]
 
 
+@pytest.mark.asyncio
+async def test_overdue_free_tiers(
+    workspace_repository: WorkspaceRepository,
+    workspace: Workspace,
+    async_session_maker: AsyncSessionMaker,
+) -> None:
+
+    async with async_session_maker() as session:
+        statement = select(orm.Organization).where(orm.Organization.id == workspace.id)
+        results = await session.execute(statement)
+        org = results.unique().scalar_one_or_none()
+        assert org
+        org.created_at = utc() - datetime.timedelta(days=1)
+        org.tier = ProductTier.Free
+        await session.commit()
+        await session.refresh(org)
+        workspace = org.to_model()
+
+    assert (
+        await workspace_repository.list_overdue_free_tier_cleanup(been_in_free_tier_for=datetime.timedelta(days=14))
+        == []
+    )
+
+    assert await workspace_repository.list_overdue_free_tier_cleanup(
+        been_in_free_tier_for=datetime.timedelta(seconds=0)
+    ) == [workspace]
+
+
 @pytest.mark.asyncio
 async def test_ack_move_to_free(
     workspace_repository: WorkspaceRepository, user: User, user_repository: UserRepository