Manually set pk to next highest when pointer may not be reset

Author: AlanCoding

ansible_base/rbac/api/queries.py

@@ -1,9 +1,12 @@
+from typing import Union
+
 from django.db.models import Model
 
 from ..models import DABContentType, get_evaluation_model
+from ..remote import RemoteObject
 
 
-def assignment_qs_user_to_obj(actor: Model, obj: Model):
+def assignment_qs_user_to_obj(actor: Model, obj: Union[Model, RemoteObject]):
     """Querset of assignments (team or user) that grants the actor any form of permission to obj"""
     evaluation_cls = get_evaluation_model(obj)
     ct = DABContentType.objects.get_for_model(obj)
@@ -18,7 +21,7 @@ def assignment_qs_user_to_obj(actor: Model, obj: Model):
     return (global_assignment_qs | obj_assignment_qs).distinct()
 
 
-def assignment_qs_user_to_obj_perm(actor: Model, obj: Model, permission: Model):
+def assignment_qs_user_to_obj_perm(actor: Model, obj: Union[Model, RemoteObject], permission: Model):
     """Queryset of assignments that grants this specific permission to this specific object"""
     evaluation_cls = get_evaluation_model(obj)
     ct = DABContentType.objects.get_for_model(obj)

ansible_base/rbac/management/__init__.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from collections import defaultdict
 
 from django.apps import apps as global_apps
@@ -42,6 +43,30 @@ def create_dab_permissions(app_config, verbosity=2, interactive=True, using=DEFA
     sync_DAB_permissions(verbosity=verbosity, using=using, apps=global_apps)
 
 
+def is_safe_identifier(name: str) -> bool:
+    """Returns True or False, name is a valid identifier in postgres, just for safety"""
+    return re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', name) is not None
+
+
+def reset_ct_sequence(ct_cls):
+    table = ct_cls._meta.db_table
+    pk_column = ct_cls._meta.pk.column
+    # This case should never be hit, but for database safety we have this
+    if not is_safe_identifier(table) or not is_safe_identifier(pk_column):
+        raise ValueError(f"Unsafe identifier: {table}.{pk_column}")
+
+    logger.info('Updating the serial sequence of DABContentType model')
+    with connection.cursor() as cursor:
+        cursor.execute(
+            f"""
+            SELECT setval(
+                pg_get_serial_sequence(%s, %s),
+                (SELECT MAX({pk_column}) FROM {table})
+            )""",
+            [table, pk_column],
+        )
+
+
 def sync_DAB_permissions(verbosity=2, using=DEFAULT_DB_ALIAS, apps=global_apps):
     """Idepotent method to set database types and permissions for DAB RBAC
 
@@ -51,7 +76,7 @@ def sync_DAB_permissions(verbosity=2, using=DEFAULT_DB_ALIAS, apps=global_apps):
     DABContentType = apps.get_model("dab_rbac", "DABContentType")
     Permission = apps.get_model("dab_rbac", "DABPermission")
 
-    create_DAB_contenttypes(verbosity=verbosity, using=using, apps=apps)
+    new_cts = create_DAB_contenttypes(verbosity=verbosity, using=using, apps=apps)
 
     # This will hold the permissions we're looking for as (content_type, (codename, name))
     searched_perms = []
@@ -106,15 +131,5 @@ def sync_DAB_permissions(verbosity=2, using=DEFAULT_DB_ALIAS, apps=global_apps):
 
     # Reset the sequence to avoid PK collision later
     if connection.vendor == 'postgresql':
-        table = DABContentType._meta.db_table
-        pk_column = DABContentType._meta.pk.column
-        with connection.cursor() as cursor:
-            cursor.execute(
-                f"""
-                SELECT setval(
-                    pg_get_serial_sequence(%s, %s),
-                    (SELECT MAX({pk_column}) FROM {table})
-                )
-            """,
-                [table, pk_column],
-            )
+        if new_cts:
+            reset_ct_sequence(DABContentType)

ansible_base/rbac/management/create_types.py

@@ -21,13 +21,15 @@ def create_DAB_contenttypes(
     verbosity=2,
     using=DEFAULT_DB_ALIAS,
     apps=global_apps,
-):
+) -> list:
     """Create DABContentType for models in the given app.
 
     This is significantly different from the ContentType post-migrate method,
     because that creates types for all apps, and so this is only called one app at a time.
     DAB RBAC runs its post_migration logic just once, because the model list
     comes from the permission registry.
+
+    Returns a list of the _new_ content types created
     """
     DABContentType = apps.get_model("dab_rbac", "DABContentType")
     ContentType = apps.get_model("contenttypes", "ContentType")
@@ -52,9 +54,12 @@ def create_DAB_contenttypes(
             real_ct = ContentType.objects.get_for_model(model)
             if not DABContentType.objects.filter(id=real_ct.id).exists():
                 ct_item_data['id'] = real_ct.id
+            else:
+                current_max_id = DABContentType.objects.order_by('-id').values_list('id', flat=True).first() or 0
+                ct_item_data['id'] = current_max_id + 1
             ct_data.append(ct_item_data)
     if not ct_data:
-        return
+        return []
 
     # To make usage earier in a transitional period, we will set the content type
     # of any new entries created here to the id of its corresponding ContentType
@@ -67,6 +72,7 @@ def create_DAB_contenttypes(
         for ct in cts:
             logger.debug("Adding DAB content type " f"'{ct.service}:{ct.app_label} | {ct.model}'")
 
+    updated_ct = 0
     for ct in DABContentType.objects.all():
         if not permission_registry.is_registered(ct.model_class()):
             logger.warning(f'{ct.model} is a stale content type in DAB RBAC')
@@ -76,3 +82,9 @@ def create_DAB_contenttypes(
             if ct.parent_content_type != parent_content_type:
                 ct.parent_content_type = parent_content_type
                 ct.save(update_fields=['parent_content_type'])
+                updated_ct += 1
+    if updated_ct:
+        # If this happens outside of the migration when the entries were created, that would be notable
+        logger.info(f'Updated the parent reference of {updated_ct} content types')
+
+    return cts

test_app/tests/rbac/api/test_access_lists.py

@@ -1,8 +1,8 @@
 import pytest
 
-from ansible_base.rbac.models import RoleDefinition
 from ansible_base.lib.utils.response import get_relative_url
 from ansible_base.rbac import permission_registry
+from ansible_base.rbac.models import RoleDefinition
 from test_app.models import Team, User