diff --git a/.github/workflows/lambda.yml b/.github/workflows/lambda.yml index 7a97f6d8..f06708b9 100644 --- a/.github/workflows/lambda.yml +++ b/.github/workflows/lambda.yml @@ -1,7 +1,7 @@ name: Post-Deploy Lambda on: - deployment: + deployment_status: jobs: deploy-lambdas: @@ -11,14 +11,20 @@ jobs: - name: Show deployment info run: | - echo "Deployment environment: $DEPLOYMENT_ENVIRONMENT" + echo "Deployment environment: ${{ github.event.deployment.environment }}" - - name: Run Lambda deploy + - name: Run Lambda production deploy + if: > + github.event.deployment.environment == 'documentcloud-prod' && + github.event.deployment_status.state == 'success' run: | - if [[ "$DEPLOYMENT_ENVIRONMENT" == "documentcloud-staging" ]]; then - echo "Deploying staging lambda updates" - bash config/aws/lambda/codeship_deploy_lambdas.sh staging-lambda --staging - else - echo "Deploying production lambda updates" - bash config/aws/lambda/codeship_deploy_lambdas.sh prod-lambda - fi + echo "Deploying production lambda updates" + bash config/aws/lambda/codeship_deploy_lambdas.sh prod-lambda + + - name: Run Lambda staging deploy + if: > + github.event.deployment.environment == 'documentcloud-staging' && + github.event.deployment_status.state == 'success' + run: | + echo "Deploying staging lambda updates" + bash config/aws/lambda/codeship_deploy_lambdas.sh staging-lambda --staging diff --git a/documentcloud/organizations/migrations/0018_alter_organization_merged.py b/documentcloud/organizations/migrations/0018_alter_organization_merged.py new file mode 100644 index 00000000..62733267 --- /dev/null +++ b/documentcloud/organizations/migrations/0018_alter_organization_merged.py @@ -0,0 +1,20 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:06 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('organizations', '0017_organization_merged'), + ] + + operations = [ + migrations.AlterField( + model_name='organization', + name='merged', + field=models.ForeignKey(blank=True, help_text='The organization this organization was merged in to', null=True, on_delete=django.db.models.deletion.PROTECT, related_name='+', to=settings.SQUARELET_ORGANIZATION_MODEL), + ), + ] diff --git a/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py b/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py new file mode 100644 index 00000000..2cd7cb47 --- /dev/null +++ b/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py @@ -0,0 +1,30 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:32 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('organizations', '0018_alter_organization_merged'), + ] + + operations = [ + migrations.AddField( + model_name='organization', + name='members', + field=models.ManyToManyField(blank=True, help_text='Organizations which are members of this organization (useful for trade associations or other member groups)', related_name='groups', to=settings.SQUARELET_ORGANIZATION_MODEL), + ), + migrations.AddField( + model_name='organization', + name='parent', + field=models.ForeignKey(blank=True, help_text='The parent organization', null=True, on_delete=django.db.models.deletion.PROTECT, related_name='children', to=settings.SQUARELET_ORGANIZATION_MODEL, verbose_name='parent'), + ), + migrations.AddField( + model_name='organization', + name='share_resources', + field=models.BooleanField(default=True, help_text='Share resources (subscriptions, credits) with all children and member organizations. Global toggle that applies to all relationships.', verbose_name='share resources'), + ), + ] diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 483de406..35ad78f0 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -130,6 +130,19 @@ def merge(self, uuid): self.addons.update(organization=other) self.visual_addons.update(organization=other) + # transfer children to the other organization + self.children.update(parent=other) + + # transfer group memberships + groups = self.groups.all() + other.groups.add(*groups) + self.groups.clear() + + # transfer members + members = self.members.all() + other.members.add(*members) + self.members.clear() + self.merged = other def calc_ai_credits_per_month(self, users): @@ -143,24 +156,80 @@ def calc_ai_credits_per_month(self, users): @transaction.atomic def use_ai_credits(self, amount, user_id, note): - """Try to deduct AI credits from the organization's balance""" + """Try to deduct AI credits from the organization's balance + + Consumes AI credits in priority order: + 1. Own monthly AI credits + 2. Own regular (purchased) AI credits + 3. Parent monthly AI credits (if parent.share_resources=True) + 4. Parent regular AI credits (if parent.share_resources=True) + 5. Group monthly AI credits (for each group where group.share_resources=True) + 6. Group regular AI credits (for each group where group.share_resources=True) + + Args: + amount: Number of AI credits to consume + user_id: ID of the user consuming the credits + note: Description of what the credits are being used for + + Returns: + dict: {"monthly": count, "regular": count} - breakdown of consumed credits + + Raises: + InsufficientAICreditsError: If not enough AI credits available across + all sources + """ initial_amount = amount ai_credit_count = {"monthly": 0, "regular": 0} - organization = Organization.objects.select_for_update().get(pk=self.pk) - - ai_credit_count["monthly"] = min(amount, organization.monthly_ai_credits) - amount -= ai_credit_count["monthly"] - ai_credit_count["regular"] = min(amount, organization.number_ai_credits) - amount -= ai_credit_count["regular"] + # Lock this organization and related organizations for update to prevent + # race conditions + organization = Organization.objects.select_for_update().get(pk=self.pk) + if organization.parent and organization.parent.share_resources: + parent = Organization.objects.select_for_update().get( + pk=organization.parent_id + ) + else: + parent = None + groups = organization.groups.filter(share_resources=True).select_for_update() + + def deduct_credits(amount, organization, field): + """Helper to deduct AI credits from a specific field on an organization""" + # Calculate how much to deduct: take up to the amount requested, + # but no more than what's available in this field + deduct_amount = min(amount, getattr(organization, field)) + amount -= deduct_amount + setattr(organization, field, getattr(organization, field) - deduct_amount) + # Return remaining amount needed and how much we deducted + return amount, deduct_amount + + # Build list of organizations to consume from in priority order + organizations = [organization] + if parent: + organizations.append(parent) + organizations.extend(groups) + + # Consume AI credits from each organization in priority order + for current_organization in organizations: + # For each organization, consume monthly credits first, then regular + for field, count in [ + ("monthly_ai_credits", "monthly"), + ("number_ai_credits", "regular"), + ]: + amount, deduct_amount = deduct_credits( + amount, current_organization, field + ) + ai_credit_count[count] += deduct_amount + if amount == 0: + break + current_organization.save() + if amount == 0: + break if amount > 0: + # Raising an error here will cancel the current atomic transaction + # No changes to the organizations will be committed to the database raise InsufficientAICreditsError(amount) - organization.monthly_ai_credits -= ai_credit_count["monthly"] - organization.number_ai_credits -= ai_credit_count["regular"] - organization.save() - organization.ai_credit_logs.create( user_id=user_id, organization=organization, @@ -170,6 +239,24 @@ def use_ai_credits(self, amount, user_id, note): return ai_credit_count + def get_total_number_ai_credits(self): + """Get total number AI credits including parent and groups""" + number_ai_credits = self.number_ai_credits + if self.parent and self.parent.share_resources: + number_ai_credits += self.parent.number_ai_credits + for group in self.groups.filter(share_resources=True): + number_ai_credits += group.number_ai_credits + return number_ai_credits + + def get_total_monthly_ai_credits(self): + """Get total monthly AI credits including parent and groups""" + monthly_ai_credits = self.monthly_ai_credits + if self.parent and self.parent.share_resources: + monthly_ai_credits += self.parent.monthly_ai_credits + for group in self.groups.filter(share_resources=True): + monthly_ai_credits += group.monthly_ai_credits + return monthly_ai_credits + class AICreditLog(models.Model): """Log usage of AI Credits""" diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index 30407323..d61a4090 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/documentcloud/organizations/tests/test_models.py @@ -2,6 +2,7 @@ import pytest # DocumentCloud +from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.models import Organization from documentcloud.organizations.tests.factories import OrganizationFactory from documentcloud.users.models import User @@ -50,7 +51,7 @@ def test_merge_fks(self): if f.is_relation and f.auto_created ] ) - == 6 + == 8 ) # Many to many relations defined on the Organization model assert ( @@ -61,5 +62,197 @@ def test_merge_fks(self): if f.many_to_many and not f.auto_created ] ) - == 1 + == 2 ) + + +class TestOrganizationCollective: + """Tests for Organization collective resource sharing""" + + @pytest.mark.django_db() + def test_use_ai_credits_with_parent(self): + """Test using AI credits with parent's resources when own resources exhausted""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, parent=parent_org + ) + + # Use 15 credits - 10 from child monthly, 5 from child regular + result = child_org.use_ai_credits(15, user.pk, "Test") + + child_org.refresh_from_db() + parent_org.refresh_from_db() + + assert result == {"monthly": 10, "regular": 5} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert parent_org.monthly_ai_credits == 50 + assert parent_org.number_ai_credits == 25 + + @pytest.mark.django_db() + def test_use_ai_credits_parent_no_sharing(self): + """Test that resources are not shared when parent.share_resources=False""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=False + ) + child_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, parent=parent_org + ) + + # Try to use 15 credits - should fail after child's 15 credits + with pytest.raises(InsufficientAICreditsError): + child_org.use_ai_credits(20, user.pk, "Test") + + @pytest.mark.django_db() + def test_use_ai_credits_with_groups(self): + """Test using AI credits with group's resources""" + user = UserFactory() + group_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=True + ) + child_org = OrganizationFactory(monthly_ai_credits=10, number_ai_credits=5) + child_org.groups.add(group_org) + + # Use 20 credits - should use 10 from child monthly, 5 from child regular, + # 5 from group monthly + result = child_org.use_ai_credits(20, user.pk, "Test") + + child_org.refresh_from_db() + group_org.refresh_from_db() + + assert result == {"monthly": 15, "regular": 5} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert group_org.monthly_ai_credits == 45 + + @pytest.mark.django_db() + def test_use_ai_credits_with_multiple_groups(self): + """Test using AI credits from multiple groups""" + user = UserFactory() + group1 = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + group2 = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + child_org = OrganizationFactory(monthly_ai_credits=10, number_ai_credits=0) + child_org.groups.add(group1, group2) + + # Use 40 credits - should use 5 from child, then from groups + result = child_org.use_ai_credits(40, user.pk, "Test") + + child_org.refresh_from_db() + group1.refresh_from_db() + group2.refresh_from_db() + + assert result == {"monthly": 30, "regular": 10} + assert child_org.monthly_ai_credits == 0 + # Groups are consumed in arbitrary order + assert group1.monthly_ai_credits + group2.monthly_ai_credits == 20 + assert group1.number_ai_credits + group2.number_ai_credits == 10 + + @pytest.mark.django_db() + def test_use_ai_credits_parent_and_groups(self): + """Test using AI credits with both parent and groups""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + group_org = OrganizationFactory( + monthly_ai_credits=30, number_ai_credits=15, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=5, number_ai_credits=0, parent=parent_org + ) + child_org.groups.add(group_org) + + # Use 60 credits: 5 child monthly, 20 parent monthly, 10 parent regular, + # 25 group monthly + result = child_org.use_ai_credits(60, user.pk, "Test") + + child_org.refresh_from_db() + parent_org.refresh_from_db() + group_org.refresh_from_db() + + assert result == {"monthly": 50, "regular": 10} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert parent_org.monthly_ai_credits == 0 + assert parent_org.number_ai_credits == 0 + assert group_org.monthly_ai_credits == 5 + assert group_org.number_ai_credits == 15 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_own_only(self): + """Test get_total_number_ai_credits with no parent or groups""" + org = OrganizationFactory(number_ai_credits=100) + assert org.get_total_number_ai_credits() == 100 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_with_parent(self): + """Test get_total_number_ai_credits including parent""" + parent_org = OrganizationFactory(number_ai_credits=50, share_resources=True) + child_org = OrganizationFactory(number_ai_credits=25, parent=parent_org) + + assert child_org.get_total_number_ai_credits() == 75 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_parent_no_sharing(self): + """Test get_total_number_ai_credits when parent doesn't share""" + parent_org = OrganizationFactory(number_ai_credits=50, share_resources=False) + child_org = OrganizationFactory(number_ai_credits=25, parent=parent_org) + + assert child_org.get_total_number_ai_credits() == 25 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_with_groups(self): + """Test get_total_number_ai_credits including groups""" + group1 = OrganizationFactory(number_ai_credits=30, share_resources=True) + group2 = OrganizationFactory(number_ai_credits=20, share_resources=True) + org = OrganizationFactory(number_ai_credits=10) + org.groups.add(group1, group2) + + assert org.get_total_number_ai_credits() == 60 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_own_only(self): + """Test get_total_monthly_ai_credits with no parent or groups""" + org = OrganizationFactory(monthly_ai_credits=50) + assert org.get_total_monthly_ai_credits() == 50 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_with_parent(self): + """Test get_total_monthly_ai_credits including parent""" + parent_org = OrganizationFactory(monthly_ai_credits=100, share_resources=True) + child_org = OrganizationFactory(monthly_ai_credits=25, parent=parent_org) + + assert child_org.get_total_monthly_ai_credits() == 125 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_with_groups(self): + """Test get_total_monthly_ai_credits including groups""" + group1 = OrganizationFactory(monthly_ai_credits=40, share_resources=True) + group2 = OrganizationFactory(monthly_ai_credits=30, share_resources=True) + org = OrganizationFactory(monthly_ai_credits=15) + org.groups.add(group1, group2) + + assert org.get_total_monthly_ai_credits() == 85 + + @pytest.mark.django_db() + def test_insufficient_ai_credits_with_parent(self): + """Test InsufficientAICreditsError even with parent resources""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=5, number_ai_credits=2, parent=parent_org + ) + + # Try to use more credits than available (total is 22, trying to use 25) + with pytest.raises(InsufficientAICreditsError): + child_org.use_ai_credits(25, user.pk, "Test") diff --git a/documentcloud/projects/migrations/0013_alter_project_options.py b/documentcloud/projects/migrations/0013_alter_project_options.py new file mode 100644 index 00000000..283ddcd7 --- /dev/null +++ b/documentcloud/projects/migrations/0013_alter_project_options.py @@ -0,0 +1,17 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:06 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('projects', '0012_auto_20210407_1801'), + ] + + operations = [ + migrations.AlterModelOptions( + name='project', + options={'ordering': ('slug',), 'permissions': (('add_remove_project', 'Can add & remove documents from a project'), ('change_project_all', 'Can edit all fields on a project (not just pinned)'))}, + ), + ] diff --git a/requirements/base.txt b/requirements/base.txt index e6cf0655..8ae79beb 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -429,7 +429,7 @@ sqlparse==0.4.4 # via # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.14 # via -r requirements/base.in stack-data==0.3.0 # via ipython diff --git a/requirements/local.txt b/requirements/local.txt index c0b7be23..660893b7 100644 --- a/requirements/local.txt +++ b/requirements/local.txt @@ -736,7 +736,7 @@ sqlparse==0.4.4 # -r requirements/./base.txt # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.14 # via -r requirements/./base.txt stack-data==0.3.0 # via diff --git a/requirements/production.txt b/requirements/production.txt index 513ec384..ded6e691 100644 --- a/requirements/production.txt +++ b/requirements/production.txt @@ -608,7 +608,7 @@ sqlparse==0.4.4 # -r requirements/./base.txt # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.14 # via -r requirements/./base.txt stack-data==0.3.0 # via