diff --git a/netbox/users/migrations/0005_create_userconfigs.py b/netbox/users/migrations/0005_create_userconfigs.py new file mode 100644 index 000000000..39ce174f6 --- /dev/null +++ b/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 + ), + ] diff --git a/netbox/users/models.py b/netbox/users/models.py index 228b5aace..d401ad68e 100644 --- a/netbox/users/models.py +++ b/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. diff --git a/netbox/users/tests/test_models.py b/netbox/users/tests/test_models.py index 55dba997b..a6d0916ac 100644 --- a/netbox/users/tests/test_models.py +++ b/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):