[ENG-4546] [OSF Institutions] Institution Default Storage Region(s) -…

… Phase 1 (#10405)

* Allow Institutions to designate allowed storage regions and use to constrain default regions for users created through institutional SSO

admin/institutions/forms.py

@@ -7,7 +7,7 @@ class Meta:
         model = Institution
 
         exclude = [
-            'is_deleted', 'contributors'
+            'is_deleted', 'contributors', 'storage_regions',
         ]
 
 

api/institutions/authentication.py

@@ -9,6 +9,9 @@
 from rest_framework.authentication import BaseAuthentication
 from rest_framework.exceptions import AuthenticationFailed, PermissionDenied
 
+from addons.osfstorage.models import UserSettings as OSFStorageUserSettings
+from addons.osfstorage.models import Region
+
 from api.base.authentication import drf
 from api.base import exceptions, settings
 
@@ -386,4 +389,18 @@ def authenticate(self, request):
                 sso_department=department,
             )
 
+        # Storage region is only updated if the user is created via institutional SSO; the region will be set to the
+        # institution's preferred one if the user's current region is not in the institution's default region list.
+        if is_created:
+            user_settings = OSFStorageUserSettings.objects.get(owner=user)
+            institution_region_list = institution.storage_regions.all()
+            if institution_region_list and user_settings.default_region not in institution_region_list:
+                try:
+                    user_settings.default_region = institution_region_list.get(institutionstorageregion__is_preferred=True)
+                    user_settings.save()
+                except Region.DoesNotExist:
+                    message = f'Institution SSO Warning: Institution {institution._id} does not have a preferred default region'
+                    sentry.log_message(message)
+                    logger.error(message)
+
         return user, None

api_tests/institutions/views/test_institution_auth.py

@@ -14,9 +14,9 @@
 from framework.auth.core import get_user
 from framework.auth.views import send_confirm_email
 
-from osf.models import OSFUser, InstitutionAffiliation
+from osf.models import OSFUser, InstitutionAffiliation, InstitutionStorageRegion
 from osf.models.institution import SsoFilterCriteriaAction
-from osf_tests.factories import InstitutionFactory, ProjectFactory, UserFactory
+from osf_tests.factories import InstitutionFactory, ProjectFactory, UserFactory, RegionFactory
 
 from tests.base import capture_signals
 
@@ -77,6 +77,46 @@ def institution():
     return InstitutionFactory()
 
 
+@pytest.fixture()
+def institution_region():
+    return RegionFactory()
+
+
+@pytest.fixture()
+def institution_region_preferred():
+    return RegionFactory()
+
+
+@pytest.fixture()
+def user_default_region():
+    user = UserFactory()
+    return user.addons_osfstorage_user_settings.default_region
+
+
+@pytest.fixture()
+def institution_without_user_default_region(institution_region, institution_region_preferred):
+    institution = InstitutionFactory()
+    institution.storage_regions.add(institution_region)
+    InstitutionStorageRegion.objects.create(
+        institution=institution,
+        storage_region=institution_region_preferred,
+        is_preferred=True
+    )
+    return institution
+
+
+@pytest.fixture()
+def institution_with_default_user_region(user_default_region, institution_region_preferred):
+    institution = InstitutionFactory()
+    institution.storage_regions.add(user_default_region)
+    InstitutionStorageRegion.objects.create(
+        institution=institution,
+        storage_region=institution_region_preferred,
+        is_preferred=True
+    )
+    return institution
+
+
 @pytest.fixture()
 def institution_primary_type_1():
     institution = InstitutionFactory()
@@ -453,6 +493,50 @@ def test_user_external_unconfirmed(self, app, institution, url_auth_institution)
         assert not user.has_usable_password()
 
 
+@pytest.mark.django_db
+class TestInstitutionStorageRegion:
+
+    def test_region_updated_for_new_user(self, app, institution_region_preferred, institution_without_user_default_region, url_auth_institution):
+        username = '[email protected]'
+        assert OSFUser.objects.filter(username=username).count() == 0
+        res = app.post(url_auth_institution, make_payload(institution_without_user_default_region, username))
+        assert res.status_code == 204
+        user = OSFUser.objects.get(username=username)
+        assert user.addons_osfstorage_user_settings.default_region == institution_region_preferred
+
+    def test_region_not_updated_for_new_user(self, app, user_default_region, institution_region_preferred, institution_with_default_user_region, url_auth_institution):
+        username = '[email protected]'
+        assert OSFUser.objects.filter(username=username).count() == 0
+        res = app.post(url_auth_institution, make_payload(institution_with_default_user_region, username))
+        assert res.status_code == 204
+        user = OSFUser.objects.filter(username=username).first()
+        assert user.addons_osfstorage_user_settings.default_region == user_default_region
+        assert user.addons_osfstorage_user_settings.default_region != institution_region_preferred
+
+    def test_region_not_updated_for_existing_user_affiliated(self, app, institution_region, institution_region_preferred,
+                                                             institution_without_user_default_region, url_auth_institution):
+        username = '[email protected]'
+        user = make_user(username, 'Foo Bar')
+        user.save()
+        res = app.post(url_auth_institution, make_payload(institution_without_user_default_region, username))
+        assert res.status_code == 204
+        user.reload()
+        assert user.addons_osfstorage_user_settings.default_region != institution_region
+        assert user.addons_osfstorage_user_settings.default_region != institution_region_preferred
+
+    def test_region_not_updated_for_existing_user_not_affiliated(self, app, institution_region, institution_region_preferred,
+                                                                 institution_without_user_default_region, url_auth_institution):
+        username = '[email protected]'
+        user = make_user(username, 'Bar Foo')
+        user.add_or_update_affiliated_institution(institution_without_user_default_region)
+        user.save()
+        res = app.post(url_auth_institution, make_payload(institution_without_user_default_region, username))
+        assert res.status_code == 204
+        user.reload()
+        assert user.addons_osfstorage_user_settings.default_region != institution_region
+        assert user.addons_osfstorage_user_settings.default_region != institution_region_preferred
+
+
 @pytest.mark.django_db
 class TestInstitutionAuthnSharedSSOCriteriaType2:
 

osf/management/commands/set_institution_storage_regions.py

@@ -0,0 +1,108 @@
+import logging
+
+from django.core.management.base import BaseCommand
+from django.db import transaction
+
+from osf.models import Institution, InstitutionStorageRegion
+from addons.osfstorage.models import Region
+
+logger = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    """Set storage regions for institutions.
+    """
+
+    def add_arguments(self, parser):
+        super(Command, self).add_arguments(parser)
+        parser.add_argument(
+            '-d',
+            '--dry',
+            action='store_true',
+            dest='dry_run',
+            help='If true, check institution and region only'
+        )
+        parser.add_argument(
+            '-i',
+            '--institution',
+            type=str,
+            required=True,
+            help='Select the institution to add the storage region to'
+        )
+        parser.add_argument(
+            '-r',
+            '--region',
+            type=str,
+            required=True,
+            help='Select the storage region to be added to the institution'
+        )
+        parser.add_argument(
+            '-p',
+            '--preferred',
+            action='store_true',
+            dest='is_preferred',
+            help='Set the storage region as the preferred choice for the institution'
+        )
+
+    def handle(self, *args, **options):
+        dry_run = options.get('dry_run', False)
+        if dry_run:
+            logger.warning('Dry Run: This is a dry-run pass!')
+        institution_id = options['institution']
+        region_id = options['region']
+        is_preferred = options.get('is_preferred', False)
+        with transaction.atomic():
+            set_institution_storage_regions(institution_id, region_id, is_preferred=is_preferred)
+            if dry_run:
+                raise RuntimeError('Dry run -- transaction rolled back')
+
+
+def set_institution_storage_regions(institution_id, region_id, is_preferred=False):
+
+    # Verify institution and region
+    try:
+        institution = Institution.objects.get(_id=institution_id)
+        region = Region.objects.get(_id=region_id)
+    except (Institution.DoesNotExist, Region.DoesNotExist) as e:
+        logger.error(f'Institution and/or Region not found: error={e}')
+        return
+    # Get or set region for institution
+    if region in institution.storage_regions.all():
+        logger.warning(f'Region [{region._id}] already set for Institution [{institution._id}]')
+        institution_storage_region = InstitutionStorageRegion.objects.get(
+            institution=institution,
+            storage_region=region
+        )
+        if institution_storage_region.is_preferred:
+            logger.warning(f'Region [{region._id}] already set as preferred for Institution [{institution._id}]')
+            return
+    else:
+        institution_storage_region = InstitutionStorageRegion.objects.create(
+            institution=institution,
+            storage_region=region
+        )
+        logger.info(f'Region [{region._id}] has been added to Institution [{institution._id}]')
+
+    # Make sure there is only one preferred region
+    try:
+        existing_preferred_institution_storage_region = InstitutionStorageRegion.objects.get(
+            institution=institution,
+            is_preferred=True,
+        )
+    # Case 1: always set the region as preferred if there is no preferred region for the institution;
+    #         this executes even if the option `-p` / `--preferred` is not provided
+    except InstitutionStorageRegion.DoesNotExist:
+        institution_storage_region.is_preferred = True
+        institution_storage_region.save()
+        logger.info(f'Region [{region._id}] has been set as preferred choice for Institution [{institution._id}]')
+        return
+    # Case 2: do nothing and return if preferred region exists and if `is_preferred` is not set
+    if not is_preferred:
+        return
+    # Case 3: if `is_preferred` is set, clear the existing preferred region before setting the new one
+    existing_preferred_institution_storage_region.is_preferred = False
+    existing_preferred_institution_storage_region.save()
+    logger.info(f'The old preferred region has been removed from Institution [{institution._id}]')
+    institution_storage_region.is_preferred = True
+    institution_storage_region.save()
+    logger.info(f'Region [{region._id}] has been set as the preferred choice for Institution [{institution._id}]')

osf/migrations/0014_institution_storage_regions.py

@@ -0,0 +1,42 @@
+# Generated by Django 3.2.17 on 2023-06-14 15:30
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django_extensions.db.fields
+import osf.models.base
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('addons_osfstorage', '0002_auto_20220817_1915'),
+        ('osf', '0013_institution_support_email'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='InstitutionStorageRegion',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('created', django_extensions.db.fields.CreationDateTimeField(auto_now_add=True, verbose_name='created')),
+                ('modified', django_extensions.db.fields.ModificationDateTimeField(auto_now=True, verbose_name='modified')),
+                ('is_preferred', models.BooleanField(default=False)),
+                ('institution', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='osf.institution')),
+                ('storage_region', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='addons_osfstorage.region')),
+            ],
+            bases=(models.Model, osf.models.base.QuerySetExplainMixin),
+        ),
+        migrations.AddField(
+            model_name='institution',
+            name='storage_regions',
+            field=models.ManyToManyField(related_name='institutions', through='osf.InstitutionStorageRegion', to='addons_osfstorage.Region'),
+        ),
+        migrations.AddConstraint(
+            model_name='institutionstorageregion',
+            constraint=models.UniqueConstraint(fields=('institution', 'storage_region'), name='unique_institution_and_region'),
+        ),
+        migrations.AddConstraint(
+            model_name='institutionstorageregion',
+            constraint=models.UniqueConstraint(condition=models.Q(('is_preferred', True)), fields=('institution', 'is_preferred'), name='unique_institution_preferred_region'),
+        ),
+    ]

osf/models/__init__.py

@@ -58,4 +58,5 @@
 from osf.models.outcomes import Outcome  # noqa
 from osf.models.outcome_artifacts import OutcomeArtifact  # noqa
 from osf.models.institution_affiliation import InstitutionAffiliation  # noqa
+from osf.models.institution_storage_region import InstitutionStorageRegion  # noqa
 from osf.models.metadata import GuidMetadataRecord  # noqa

osf/models/institution.py

@@ -16,6 +16,7 @@
 from osf.models import base
 from osf.models.contributor import InstitutionalContributor
 from osf.models.institution_affiliation import InstitutionAffiliation
+from osf.models.institution_storage_region import InstitutionStorageRegion
 from osf.models.mixins import Loggable, GuardianMixin
 from osf.models.storage import InstitutionAssetFile
 from osf.models.validators import validate_email
@@ -78,6 +79,13 @@ class Institution(DirtyFieldsMixin, Loggable, base.ObjectIDMixin, base.BaseModel
         default=''
     )
 
+    # Default Storage Region
+    storage_regions = models.ManyToManyField(
+        'addons_osfstorage.Region',
+        through=InstitutionStorageRegion,
+        related_name='institutions'
+    )
+
     # Verified employment/education affiliation source for `via-orcid` institutions
     orcid_record_verified_source = models.CharField(max_length=255, blank=True, default='')
 

osf/models/institution_storage_region.py

@@ -0,0 +1,27 @@
+from django.db import models
+
+from osf.models.base import BaseModel
+
+
+class InstitutionStorageRegion(BaseModel):
+
+    institution = models.ForeignKey('Institution', on_delete=models.CASCADE)
+    storage_region = models.ForeignKey('addons_osfstorage.Region', on_delete=models.CASCADE)
+    is_preferred = models.BooleanField(default=False)
+
+    class Meta:
+        constraints = [
+            models.UniqueConstraint(fields=['institution', 'storage_region'], name='unique_institution_and_region'),
+            models.UniqueConstraint(
+                fields=['institution', 'is_preferred'],
+                condition=models.Q(is_preferred=True),
+                name='unique_institution_preferred_region'
+            ),
+        ]
+
+    def __repr__(self):
+        return f'<{self.__class__.__name__}(institution={self.institution._id}, ' \
+               f'storage_region={self.storage_region._id}, is_preferred={self.is_preferred}>'
+
+    def __str__(self):
+        return f'{self.institution._id}::{self.storage_region._id}::{self.is_preferred}'