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

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.contrib.postgres.fields import JSONField
 from django.core.validators import MinLengthValidator
 from django.core.validators import MinLengthValidator
 from django.db import models
 from django.db import models
+from django.db.models.signals import post_save
+from django.dispatch import receiver
 from django.utils import timezone
 from django.utils import timezone
 
 
 
 
@@ -31,23 +33,24 @@ class UserConfig(models.Model):
         ordering = ['user']
         ordering = ['user']
         verbose_name = verbose_name_plural = 'User Preferences'
         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:
         Retrieve a configuration parameter specified by its dotted path. Example:
 
 
             userconfig.get('foo.bar.baz')
             userconfig.get('foo.bar.baz')
 
 
         :param path: Dotted path to the configuration key. For example, 'foo.bar' returns self.data['foo']['bar'].
         :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
         d = self.data
         keys = path.split('.')
         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:
         for key in keys:
             if type(d) is dict:
             if type(d) is dict:
                 d = d.get(key)
                 d = d.get(key)
             else:
             else:
-                return None
+                return default
 
 
         return d
         return d
 
 
@@ -116,6 +119,15 @@ class UserConfig(models.Model):
             self.save()
             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):
 class Token(models.Model):
     """
     """
     An API token used for user authentication. This extends the stock model to allow each user to have multiple tokens.
     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):
     def setUp(self):
 
 
         user = User.objects.create_user(username='testuser')
         user = User.objects.create_user(username='testuser')
-        initial_data = {
+        user.config.data = {
             'a': True,
             'a': True,
             'b': {
             'b': {
                 'foo': 101,
                 '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):
     def test_get(self):
         userconfig = self.userconfig
         userconfig = self.userconfig
@@ -58,12 +59,12 @@ class UserConfigTest(TestCase):
         userconfig.set('b.baz', 'abc')
         userconfig.set('b.baz', 'abc')
         self.assertEqual(userconfig.data['d'], 'abc')
         self.assertEqual(userconfig.data['d'], 'abc')
         self.assertEqual(userconfig.data['b']['baz'], 'abc')
         self.assertEqual(userconfig.data['b']['baz'], 'abc')
-        self.assertIsNone(userconfig.pk)
 
 
         # Set a value and commit to the database
         # Set a value and commit to the database
         userconfig.set('a', 'def', commit=True)
         userconfig.set('a', 'def', commit=True)
+
+        userconfig.refresh_from_db()
         self.assertEqual(userconfig.data['a'], 'def')
         self.assertEqual(userconfig.data['a'], 'def')
-        self.assertIsNotNone(userconfig.pk)
 
 
         # Attempt to change a branch node to a leaf node
         # Attempt to change a branch node to a leaf node
         with self.assertRaises(TypeError):
         with self.assertRaises(TypeError):