Просмотр исходного кода

Automatically create UserConfig for users

Jeremy Stretch 5 лет назад
Родитель
Сommit
f3012ed839

+ 27 - 0
netbox/users/migrations/0005_create_userconfigs.py

@@ -0,0 +1,27 @@
+from django.contrib.auth import get_user_model
+from django.db import migrations
+
+
+def create_userconfigs(apps, schema_editor):
+    """
+    Create an empty UserConfig instance for each existing User.
+    """
+    User = get_user_model()
+    UserConfig = apps.get_model('users', 'UserConfig')
+    UserConfig.objects.bulk_create(
+        [UserConfig(user_id=user.pk) for user in User.objects.all()]
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('users', '0004_userconfig'),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            code=create_userconfigs,
+            reverse_code=migrations.RunPython.noop
+        ),
+    ]

+ 15 - 3
netbox/users/models.py

@@ -5,6 +5,8 @@ from django.contrib.auth.models import User
 from django.contrib.postgres.fields import JSONField
 from django.core.validators import MinLengthValidator
 from django.db import models
+from django.db.models.signals import post_save
+from django.dispatch import receiver
 from django.utils import timezone
 
 
@@ -31,23 +33,24 @@ class UserConfig(models.Model):
         ordering = ['user']
         verbose_name = verbose_name_plural = 'User Preferences'
 
-    def get(self, path):
+    def get(self, path, default=None):
         """
         Retrieve a configuration parameter specified by its dotted path. Example:
 
             userconfig.get('foo.bar.baz')
 
         :param path: Dotted path to the configuration key. For example, 'foo.bar' returns self.data['foo']['bar'].
+        :param default: Default value to return for a nonexistent key (default: None).
         """
         d = self.data
         keys = path.split('.')
 
-        # Iterate down the hierarchy, returning None for any invalid keys
+        # Iterate down the hierarchy, returning the default value if any invalid key is encountered
         for key in keys:
             if type(d) is dict:
                 d = d.get(key)
             else:
-                return None
+                return default
 
         return d
 
@@ -116,6 +119,15 @@ class UserConfig(models.Model):
             self.save()
 
 
+@receiver(post_save, sender=User)
+def create_userconfig(instance, created, **kwargs):
+    """
+    Automatically create a new UserConfig when a new User is created.
+    """
+    if created:
+        UserConfig(user=instance).save()
+
+
 class Token(models.Model):
     """
     An API token used for user authentication. This extends the stock model to allow each user to have multiple tokens.

+ 5 - 4
netbox/users/tests/test_models.py

@@ -9,7 +9,7 @@ class UserConfigTest(TestCase):
     def setUp(self):
 
         user = User.objects.create_user(username='testuser')
-        initial_data = {
+        user.config.data = {
             'a': True,
             'b': {
                 'foo': 101,
@@ -27,8 +27,9 @@ class UserConfigTest(TestCase):
                 }
             }
         }
+        user.config.save()
 
-        self.userconfig = UserConfig(user=user, data=initial_data)
+        self.userconfig = user.config
 
     def test_get(self):
         userconfig = self.userconfig
@@ -58,12 +59,12 @@ class UserConfigTest(TestCase):
         userconfig.set('b.baz', 'abc')
         self.assertEqual(userconfig.data['d'], 'abc')
         self.assertEqual(userconfig.data['b']['baz'], 'abc')
-        self.assertIsNone(userconfig.pk)
 
         # Set a value and commit to the database
         userconfig.set('a', 'def', commit=True)
+
+        userconfig.refresh_from_db()
         self.assertEqual(userconfig.data['a'], 'def')
-        self.assertIsNotNone(userconfig.pk)
 
         # Attempt to change a branch node to a leaf node
         with self.assertRaises(TypeError):